"""HPO sweep for jepa-fx-risk HEPA backbone.

Runs train.py with different JEPA_* env overrides, logs results to
results/hpo/hpo_results.jsonl. Each config writes its metrics.json then
the result is appended to the JSONL.

Usage:
    python scripts/hpo_sweep.py
    python scripts/hpo_sweep.py --dry-run   # print configs, don't train
"""
import argparse
import json
import os
import subprocess
import sys
from datetime import datetime
from itertools import product
from pathlib import Path

# ── Search space ──────────────────────────────────────────────────────────────
SEARCH_SPACE = {
    "JEPA_D_MODEL": [64, 128, 256],
    "JEPA_DEPTH":   [2, 4],
    "JEPA_WINDOW":  [120, 240, 480],
}
# Fixed: PATCH_LEN=24 (1-day patches), N_HEADS=4, EPOCHS=300, PHASE1_EPOCHS=200

PYTHON = str(Path(sys.executable))
OUT_DIR = Path("results/hpo")


def configs():
    """Yield all configs as dicts of JEPA_* env overrides."""
    keys = list(SEARCH_SPACE.keys())
    for vals in product(*SEARCH_SPACE.values()):
        yield dict(zip(keys, vals))


def run_config(cfg: dict, metrics_path: str = "metrics.json") -> dict:
    env = {**os.environ, **{k: str(v) for k, v in cfg.items()}}
    result = subprocess.run(
        [PYTHON, "train.py"],
        env=env,
        capture_output=True,
        text=True,
    )
    if result.returncode != 0:
        return {"config": cfg, "error": result.stderr[-500:]}
    stdout_last = result.stdout.strip().split("\n")[-1]
    with open(metrics_path) as f:
        m = json.load(f)
    return {
        "config": cfg,
        "val_vol_r2":  m.get("val_vol_r2"),
        "phase1_r2":   m.get("phase1_r2"),
        "stdout_last": stdout_last,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    out_file = OUT_DIR / "hpo_results.jsonl"

    all_cfgs = list(configs())
    print(f"HPO sweep: {len(all_cfgs)} configs")
    for i, cfg in enumerate(all_cfgs):
        label = " ".join(f"{k.replace('JEPA_','')}={v}" for k, v in cfg.items())
        print(f"\n[{i+1}/{len(all_cfgs)}] {label}")
        if args.dry_run:
            continue
        ts = datetime.utcnow().isoformat()
        row = run_config(cfg)
        row["ts"] = ts
        with open(out_file, "a") as f:
            f.write(json.dumps(row) + "\n")
        if "error" in row:
            print(f"  ERROR: {row['error'][:200]}")
        else:
            print(f"  val_vol_r2={row['val_vol_r2']:.4f}  phase1_r2={row['phase1_r2']:.4f}")

    if not args.dry_run:
        # Print leaderboard
        rows = [json.loads(l) for l in open(out_file) if l.strip()]
        rows = [r for r in rows if "error" not in r]
        rows.sort(key=lambda r: r.get("phase1_r2", -999), reverse=True)
        print("\n── Leaderboard (by phase1_r2) ─────────────────────────")
        for r in rows[:5]:
            cfg_str = " ".join(f"{k.replace('JEPA_','')}={v}" for k,v in r["config"].items())
            print(f"  {r['phase1_r2']:.4f}  {cfg_str}")


if __name__ == "__main__":
    main()