"""Collect <runs_dir>/*/eval.json into one markdown table (plus baselines).

usage: python scripts/collect.py [runs_dir=runs] [baseline_json=results/ablation_baseline/eval_all.json]"""
from __future__ import annotations
import json, sys
from pathlib import Path

root = Path(sys.argv[1] if len(sys.argv) > 1 else "runs")
rows = []
for run in sorted(root.glob("*/eval.json")):
    r = json.load(open(run))
    ol = r["openloop_action_mse"]; sp = r["openloop_action_mse_split"]
    cfg = r["config"]; tr = json.load(open(run.parent / "train_args.json")) if (run.parent / "train_args.json").exists() else {}
    rows.append({
        "run": run.parent.name,
        "params_M": r["params_total"] / 1e6,
        "d/L": f"{cfg['dim']}/{cfg['num_layers']}",
        "future_state": cfg["predict_future_state"],
        "train_frac": tr.get("train_fraction"),
        "heldout": r.get("heldout_fraction"),
        "cascade": ol.get("cascade", [float("nan")])[0],
        "coarse_only(τ=.4)": ol.get("coarse_only", [float("nan")])[0],
        "coarse_to_zero": ol.get("coarse_to_zero", [float("nan")])[0],
        "cascade_arm": sp.get("cascade_arm", float("nan")),
        "cascade_hand": sp.get("cascade_hand", float("nan")),
        "fstate_mse": r.get("openloop_future_state_mse", {}).get("cascade", float("nan")),
        "act_loss@τ.2": r["forward_tau_losses"]["0.2"]["action_loss"][0],
    })
base = json.load(open(sys.argv[2] if len(sys.argv) > 2 else Path(__file__).resolve().parents[1] / "results/ablation_baseline/eval_all.json"))
for name, r in base.items():
    ol = r["openloop_action_mse"]; sp = r["openloop_action_mse_split"]
    rows.append({
        "run": f"baseline:{name}", "params_M": 5686.4 if name != "noforce" else 5638.5, "d/L": "3072/30 (Wan2.2 LoRA r4)",
        "future_state": "video", "train_frac": None, "heldout": None,
        "cascade": ol.get("cascade", [float("nan")])[0], "coarse_only(τ=.4)": ol.get("coarse_only", [float("nan")])[0],
        "coarse_to_zero": float("nan"), "cascade_arm": sp.get("cascade_arm", float("nan")), "cascade_hand": sp.get("cascade_hand", float("nan")),
        "fstate_mse": float("nan"), "act_loss@τ.2": r["forward_tau_losses"]["0.2"]["action_loss"][0],
    })
keys = list(rows[0].keys())
print("| " + " | ".join(keys) + " |")
print("|" + "---|" * len(keys))
for row in rows:
    cells = []
    for k in keys:
        v = row[k]
        cells.append(f"{v:.4f}" if isinstance(v, float) and k != "params_M" else (f"{v:.1f}" if isinstance(v, float) else str(v)))
    print("| " + " | ".join(cells) + " |")
json.dump(rows, open(root / "summary.json", "w"), indent=1)