"""Experimental designs: a design = list of cells (config overrides + treatment labels) × seeds.
Runs are executed with multiprocessing (classical) or sequentially with a shared LLM cache (API cells),
each writing its own run directory + manifest; outcomes are appended to a parquet table."""
from __future__ import annotations

import json
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

import pandas as pd
import yaml

from ..metrics.outcomes import run_outcomes
from .config import RunConfig
from .runner import run_config


def load_design(path: str | Path) -> dict:
    with open(path) as f:
        return yaml.safe_load(f)


def expand_cells(design: dict) -> list[dict]:
    """Design file: {base: <config path or dict>, seeds: [..] | {start, n}, cells: [{name, overrides, treatment}], factors: {...}}.
    `factors` (optional) is a dict factor -> list of {label, overrides}; the full factorial product is expanded."""
    cells = list(design.get("cells", []))
    factors = design.get("factors")
    if factors:
        import itertools
        names = list(factors)
        for combo in itertools.product(*[factors[n] for n in names]):
            ov, tr, label = {}, {}, []
            for n, lev in zip(names, combo):
                ov.update(lev.get("overrides", {}))
                tr[n] = lev["label"]
                label.append(f"{n}={lev['label']}")
            cells.append({"name": "__".join(label), "overrides": ov, "treatment": tr})
    return cells


def seeds_of(design: dict) -> list[int]:
    s = design.get("seeds", {"start": 1, "n": 30})
    if isinstance(s, list):
        return [int(x) for x in s]
    return list(range(int(s.get("start", 1)), int(s.get("start", 1)) + int(s["n"])))


def cell_config(design: dict, cell: dict, seed: int, base_dir: Path) -> RunConfig:
    base = design["base"]
    cfg = RunConfig.from_yaml(base_dir / base) if isinstance(base, str) else RunConfig.model_validate(base)
    ov = dict(design.get("overrides", {}))
    ov.update(cell.get("overrides", {}))
    cfg = cfg.with_overrides(**ov)
    return cfg.with_overrides(name=f"{design['name']}/{cell['name']}", seed=seed,
                              treatment={**design.get("treatment", {}), **cell.get("treatment", {})})


def _run_one(args):
    cfg_json, out_dir, cache_path = args
    cfg = RunConfig.model_validate(cfg_json)
    out = Path(out_dir)
    if (out / "outcomes.json").exists() and (out / "manifest.json").exists():
        with open(out / "manifest.json") as f:
            man = json.load(f)
        if man.get("config_hash") == cfg.config_hash():       # resume only an identical configuration
            with open(out / "outcomes.json") as f:
                return json.load(f)
        import shutil
        shutil.rmtree(out)
    res = run_config(cfg, out_dir=out, cache_path=cache_path)
    oc = run_outcomes(res)
    oc["run_dir"] = str(out)
    with open(out / "outcomes.json", "w") as f:
        json.dump(oc, f, indent=1, default=float)
    return oc


def run_design(design_path: str | Path, results_root: str | Path = "results/runs", workers: int = 4,
               cache_path: str | None = None, only_cells: list[str] | None = None, max_seeds: int | None = None,
               progress: bool = True) -> pd.DataFrame:
    design_path = Path(design_path)
    design = load_design(design_path)
    cells = expand_cells(design)
    if only_cells:
        cells = [c for c in cells if c["name"] in only_cells]
    seeds = seeds_of(design)
    if max_seeds:
        seeds = seeds[:max_seeds]
    root = Path(results_root) / design["name"]
    root.mkdir(parents=True, exist_ok=True)
    jobs = []
    for c in cells:
        for s in seeds:
            cfg = cell_config(design, c, s, design_path.parent)
            out = root / c["name"] / f"seed{s:04d}"
            jobs.append((cfg.model_dump(mode="json"), str(out), cache_path))
    t0 = time.perf_counter()
    rows = []
    if workers <= 1:
        for i, j in enumerate(jobs):
            rows.append(_run_one(j))
            if progress:
                print(f"  [{i + 1}/{len(jobs)}] {j[1]}  ({time.perf_counter() - t0:.0f}s)", flush=True)
    else:
        with ProcessPoolExecutor(max_workers=workers) as ex:
            futs = {ex.submit(_run_one, j): j for j in jobs}
            for i, f in enumerate(as_completed(futs)):
                rows.append(f.result())
                if progress and (i + 1) % 10 == 0:
                    print(f"  [{i + 1}/{len(jobs)}] done ({time.perf_counter() - t0:.0f}s)", flush=True)
    df = pd.DataFrame(rows)
    df.to_parquet(root / "outcomes.parquet", index=False)
    with open(root / "design.json", "w") as f:
        json.dump({"design": design, "cells": cells, "seeds": seeds, "n_runs": len(jobs), "wall_s": time.perf_counter() - t0}, f, indent=1, default=str)
    return df