"""Aggregate per-seed eval.json for an arbitrary run prefix; optional state-vs-pixel overlay."""
import json, os, argparse
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from common import wilson_ci

def collect(prefix, budgets, seeds):
    rows = []
    for b in budgets:
        per, K, N = [], 0, 0
        for s in seeds:
            p = f"runs/{prefix}_b{b:04d}_seed{s}/eval.json"
            if not os.path.exists(p): print("missing", p); continue
            o = json.load(open(p))["overall"]; per.append(o["rate"]); K += o["k"]; N += o["n"]
        if not per: continue
        per = np.array(per); pr, lo, hi = wilson_ci(K, N)
        rows.append(dict(demos=b, per_seed=per.tolist(), pooled_k=K, pooled_n=N,
                         pooled_rate=pr, ci_low=lo, ci_high=hi,
                         mean=float(per.mean()),
                         std=float(per.std(ddof=1)) if len(per) > 1 else 0.0))
    return rows

ap = argparse.ArgumentParser()
ap.add_argument("--prefix", default="pix")
ap.add_argument("--label", default="pixel GCBC")
ap.add_argument("--compare_prefix", default=None)
ap.add_argument("--compare_label", default="state GCBC")
ap.add_argument("--budgets", nargs="+", type=int, default=[10, 50, 100, 1000])
ap.add_argument("--seeds", nargs="+", type=int, default=[0, 1, 2])
ap.add_argument("--out", default="results/scaling_pixel.png")
a = ap.parse_args()
os.makedirs(os.path.dirname(a.out), exist_ok=True)

series = [(a.prefix, a.label, "#2563eb")]
if a.compare_prefix: series.append((a.compare_prefix, a.compare_label, "#dc2626"))

plt.figure(figsize=(7.5, 4.8))
allrows = {}
for pref, label, color in series:
    rows = collect(pref, a.budgets, a.seeds); allrows[pref] = rows
    if not rows: continue
    xs = [r["demos"] for r in rows]; y = np.array([r["pooled_rate"] for r in rows])
    lo = np.array([r["ci_low"] for r in rows]); hi = np.array([r["ci_high"] for r in rows])
    plt.errorbar(xs, y, yerr=[y - lo, hi - y], fmt="o-", capsize=4, lw=2, color=color,
                 label=f"{label} (95% Wilson CI)", zorder=3)
    for r in rows:
        plt.scatter([r["demos"]] * len(r["per_seed"]), r["per_seed"], color=color, alpha=.35, s=20, zorder=2)
    print(f"\n=== {label} ===")
    print("Demos   pooled     95% CI            per-seed              mean±std")
    for r in rows:
        ps = " ".join(f"{v:.1%}" for v in r["per_seed"])
        print(f"{r['demos']:5d}  {r['pooled_rate']:6.1%}  [{r['ci_low']:.1%},{r['ci_high']:.1%}]   {ps:22s} {r['mean']:.1%}±{r['std']:.1%}")

plt.xscale("log"); plt.xlabel("number of demonstrations (log scale)")
plt.ylabel("online success rate")
plt.title("Cube-single — success vs. demonstrations\n3 seeds x 100 rollouts x 5 tasks")
plt.grid(True, which="both", alpha=0.3); plt.legend(); plt.tight_layout()
plt.savefig(a.out, dpi=130); print("\nfigure ->", a.out)
json.dump(allrows, open(a.out.replace(".png", ".json"), "w"), indent=2)