"""Read eval.json across budgets and produce the success-vs-demos figure + table."""
import json, os, argparse
import numpy as np
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

ap = argparse.ArgumentParser()
ap.add_argument("--budgets", nargs="+", type=int, default=[10, 50, 100, 1000])
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--tag", default="state")
ap.add_argument("--outdir", default="results")
a = ap.parse_args()
os.makedirs(a.outdir, exist_ok=True)

xs, ys, los, his, rows = [], [], [], [], []
for b in a.budgets:
    p = f"runs/gcbc_{a.tag}_b{b:04d}_seed{a.seed}/eval.json"
    if not os.path.exists(p):
        print("missing", p); continue
    d = json.load(open(p)); o = d["overall"]
    xs.append(b); ys.append(o["rate"]); los.append(o["ci"][0]); his.append(o["ci"][1])
    rows.append((b, o["k"], o["n"], o["rate"], o["ci"][0], o["ci"][1]))

ys = np.array(ys); los = np.array(los); his = np.array(his)
plt.figure(figsize=(7, 4.5))
plt.errorbar(xs, ys, yerr=[ys - los, his - ys], fmt="o-", capsize=4, lw=2, color="#2563eb")
plt.xscale("log")
plt.xlabel("number of demonstrations (log scale)")
plt.ylabel("online success rate")
plt.title("Cube-single GCBC (state) — success vs. demonstrations\n95% Wilson CI, 50 rollouts x 5 tasks")
plt.grid(True, which="both", alpha=0.3); plt.ylim(0, max(0.5, his.max() * 1.15))
plt.tight_layout()
fig = os.path.join(a.outdir, f"success_vs_demos_{a.tag}.png")
plt.savefig(fig, dpi=130); print("figure ->", fig)

csv = os.path.join(a.outdir, f"scaling_{a.tag}.csv")
with open(csv, "w") as f:
    f.write("demos,successes,rollouts,success_rate,ci_low,ci_high\n")
    for r in rows: f.write("%d,%d,%d,%.4f,%.4f,%.4f\n" % r)
print("table ->", csv)
print("\nDemos  Success   95% Wilson CI")
for b, k, n, p, lo, hi in rows:
    print(f"{b:5d}  {p:6.1%}   [{lo:.1%}, {hi:.1%}]  ({k}/{n})")