#!/usr/bin/env python
"""Parallel experiment sweep runner.

Usage:
  .venv/bin/python -m scripts.run_sweep --exp seeds --jobs jobs.csv
  .venv/bin/python -m scripts.run_sweep --exp seeds --jobs jobs.csv --post-only   # eval/plot only

jobs.csv columns: job,model,seed,flags   (flags = extra train.py args, may be empty)

For each job: trains runs/<exp>/<job>/<model>/seed<seed>, then runs eval + plot.
Skips jobs whose results.json already exists (idempotent resume).
Concurrency default 2 (2 cores, one thread per child). Logs to runs/<exp>/logs/.
Appends a compact row per job to runs/<exp>/summary.csv.
"""
import argparse
import csv
import json
import os
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_venv_py = os.path.join(ROOT, ".venv", "bin", "python")
PY = _venv_py if os.path.exists(_venv_py) else sys.executable


def parse_jobs(path: str) -> list[dict]:
    jobs = []
    with open(path) as fh:
        for row in csv.DictReader(fh):
            jobs.append({
                "job": row["job"].strip(),
                "model": row["model"].strip(),
                "seed": row["seed"].strip(),
                "flags": (row.get("flags") or "").strip().split(),
            })
    return jobs


def out_dir(exp: str, job: dict) -> str:
    return os.path.join(ROOT, "runs", exp, job["job"], job["model"], f"seed{job['seed']}")


def run_one(exp: str, job: dict, post_only: bool) -> dict:
    od = out_dir(exp, job)
    logdir = os.path.join(ROOT, "runs", exp, "logs")
    os.makedirs(logdir, exist_ok=True)
    tag = f"{job['job']}__{job['model']}_s{job['seed']}"
    logpath = os.path.join(logdir, tag + ".log")
    results_path = os.path.join(od, "results.json")
    res = {"job": job["job"], "model": job["model"], "seed": job["seed"]}

    if os.path.exists(results_path):
        res["skipped"] = True
        return res

    t0 = time.time()
    if not post_only:
        cmd = [PY, "-m", "src.train", job["model"], job["seed"],
               "--out_dir", f"runs/{exp}/{job['job']}"] + job["flags"]
        # CPU-resident jobs get 2 threads (sequential tied cell benefits); GPU jobs
        # stay at 1 so the feeder process doesn't contend with concurrent CPU jobs.
        n_threads = "2" if "--device cpu" in job["flags"] else "1"
        env = dict(os.environ, OMP_NUM_THREADS=n_threads, MKL_NUM_THREADS=n_threads)
        with open(logpath, "w") as lf:
            p = subprocess.run(cmd, cwd=ROOT, env=env, stdout=lf, stderr=subprocess.STDOUT)
        if p.returncode != 0:
            res["error"] = f"train rc={p.returncode} (see {logpath})"
            return res

    for cmd in ([PY, "-m", "src.eval", job["model"], job["seed"], "--runs-dir", f"runs/{exp}/{job['job']}"],
                [PY, "-m", "scripts.plot", job["model"], job["seed"], "--runs-dir", f"runs/{exp}/{job['job']}"]):
        with open(logpath, "a") as lf:
            subprocess.run(cmd, cwd=ROOT, stdout=lf, stderr=subprocess.STDOUT)

    if not os.path.exists(results_path):
        res["error"] = f"eval produced no results.json (see {logpath})"
        return res

    with open(results_path) as fh:
        r = json.load(fh)
    sig = r.get("signature", {})
    final = r.get("final", {})
    probe = final.get("probe", {})
    hlt = final.get("halting") or {}
    res.update({
        "minutes": round((time.time() - t0) / 60, 1),
        "signature": sig.get("code"),
        "probe": probe.get("code"),
        "halting": hlt.get("code"),
        "val_em_last": final.get("val_exact_match"),
        "val_em_best": r.get("val_selected", {}).get("val_exact_match"),
        "probe_acc": probe.get("acc"),
        "halt_mean": hlt.get("mean_steps"),
    })
    return res


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--exp", required=True)
    ap.add_argument("--jobs", required=True)
    ap.add_argument("--post-only", action="store_true")
    ap.add_argument("--concurrency", type=int, default=2)
    a = ap.parse_args()
    jobs = parse_jobs(os.path.join(ROOT, a.jobs))
    print(f"{len(jobs)} jobs, concurrency {a.concurrency}, post_only={a.post_only}")

    results = []
    with ThreadPoolExecutor(max_workers=a.concurrency) as pool:
        futs = [pool.submit(run_one, a.exp, j, a.post_only) for j in jobs]
        for f in as_completed(futs):
            r = f.result()
            results.append(r)
            print(f"[{r['job']}/{r['model']}/s{r['seed']}] "
                  + (r.get("error") or r.get("skipped") and "skipped (exists)" or
                     f"{r.get('signature')}/{r.get('probe')}/{r.get('halting')} "
                     f"valEM {r.get('val_em_last')} probe {r.get('probe_acc')} "
                     f"({r.get('minutes')}min)"))

    summary_path = os.path.join(ROOT, "runs", a.exp, "summary.csv")
    with open(summary_path, "w", newline="") as fh:
        keys = ["job", "model", "seed", "minutes", "signature", "probe", "halting",
                "val_em_last", "val_em_best", "probe_acc", "halt_mean", "error", "skipped"]
        w = csv.DictWriter(fh, fieldnames=keys, extrasaction="ignore")
        w.writeheader()
        for r in sorted(results, key=lambda x: (x["job"], x["model"])):
            w.writerow(r)
    print(f"summary: {summary_path}")


if __name__ == "__main__":
    main()