"""Paper-side analysis from collected results: gap-vs-scale table + curves.
    python scripts/plot_ladder.py --results results/h200node1 --fig results/h200node1/fig_ladder.png"""
import os, json, glob, argparse
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

p = argparse.ArgumentParser()
p.add_argument("--results", required=True)
p.add_argument("--fig", default=None)
a = p.parse_args()
summary = json.load(open(os.path.join(a.results, "summary.json")))
sizes = sorted({k.rsplit("_", 2)[0] if "_zbp_" in k else k.split("_")[0] for k in summary})
print(f"{'run':22s} {'params':>10s} {'steps':>7s} {'val_loss':>9s} {'tok/s':>9s}")
for k, v in sorted(summary.items()):
    print(f"{k:22s} {v['params'] or 0:>10d} {v['steps']:>7d} {v['val_loss'] or float('nan'):>9.3f} {v['tok_per_s'] or 0:>9.0f}")
fig, axes = plt.subplots(1, 2, figsize=(10, 3.8))
for k in sorted(summary):
    rows = [json.loads(l) for l in open(os.path.join(a.results, f"{k}.jsonl"))]
    ev = [(r["step"], r["val_loss"]) for r in rows if r.get("kind") == "eval" and r["step"] > 0]
    if ev:
        axes[0].plot(*zip(*ev), label=k, lw=1.4)
axes[0].set_xlabel("step"); axes[0].set_ylabel("val loss"); axes[0].legend(fontsize=6); axes[0].grid(alpha=0.3)
# gap vs scale: for each size, zbp - bp final
gaps = {}
for k, v in summary.items():
    if v["val_loss"] is None: continue
    size = k.split("_")[0]
    gaps.setdefault(size, {})[k[len(size) + 1:]] = (v["params"], v["val_loss"])
xs, arms = [], {}
for size, d in sorted(gaps.items(), key=lambda kv: (kv[1].get("bp") or (0, 0))[0]):
    if "bp" not in d: continue
    for arm, (params, loss) in d.items():
        if arm == "bp": continue
        arms.setdefault(arm, []).append((params, loss - d["bp"][1]))
for arm, pts in arms.items():
    pts.sort()
    axes[1].plot(*zip(*pts), "o-", label=f"{arm} − bp")
axes[1].set_xscale("log"); axes[1].axhline(0, color="k", lw=0.8)
axes[1].set_xlabel("params"); axes[1].set_ylabel("val-loss gap vs BP"); axes[1].legend(fontsize=7); axes[1].grid(alpha=0.3)
fig.tight_layout(); out = a.fig or os.path.join(a.results, "fig_ladder.png")
fig.savefig(out, dpi=150); print("saved", out)