"""Small-multiples: pass@1 vs loop depth k, easy (pastel) + hard (dark),
one panel per regime arm — the full design-space sweep. All loop arms on
the same 250-item MBPP eval, e400 checkpoints. Final panel: no-recurrence
references (FF control, plan-distill; distill on the 500-item set).
Pending arms render as empty boxes and fill in on re-render.
"""
import json
import re
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = Path(__file__).resolve().parent.parent / "results-loop"
LOGD = Path("/home/nils/gpuq/spark-gpu0")

def cells(fname):
    try:
        d = json.load(open(OUT / fname))
    except FileNotFoundError:
        return None
    return {int(k): (v["by_label"].get("easy"), v["by_label"].get("hard"))
            for k, v in d["ks"].items()}

def cells_from_log(logfile):
    out = {}
    pat = re.compile(r"k=(\d+): pass@1=[\d.]+  by_label=\{'easy': ([\d.]+), "
                     r"'hard': ([\d.]+)")
    try:
        for line in open(logfile):
            m = pat.search(line.replace('"', "'"))
            if m:
                out[int(m.group(1))] = (float(m.group(2)), float(m.group(3)))
    except FileNotFoundError:
        pass
    return out or None

PANELS = [
    ("untrained α-merge\ntraining-free · ρ=0.3",
     cells("eval_code_untrained.json"), "#8a8f98"),
    ("trained merge  ★\nfixed A,B · curriculum",
     cells("eval_code_trained.json"), "#2b6cb0"),
    ("tied-alpha\nlearned a, B=(1−a) tied · curriculum",
     cells("eval_code_ta_e400.json"), "#0987a0"),
    ("random-depth merge\nfixed A,B · NO curriculum",
     cells("eval_code_merge_rk16_e400.json")
     or cells_from_log(LOGD / "zzz_a_randk.log"), "#3182ce"),
    ("unconstrained rec\nlearned A,B · noise h₀ · ρ→4.5",
     cells("eval_code_rec16_partial.json"), "#c53030"),
    ("Parcae rec\nρ<1 enforced · learned B · ρ→0.29",
     cells("eval_code_parcae16_e400.json"), "#2f855a"),
    ("per-depth merges\ntime-varying · anchored",
     cells("eval_code_pd4_e400.json")
     or cells_from_log(LOGD / "perdepth_arm.log"), "#805ad5"),
    ("noise-s₀ merge\nfixed A,B · curriculum · noise s₀",
     cells("eval_code_merge_ns_e400.json")
     or cells_from_log(LOGD / "zzz_b_noises0.log"), "#b83280"),
    ("4× capacity merge\nMLP h=2048 · curriculum",
     cells("eval_code_merge_h2048_e400.json")
     or cells_from_log(LOGD / "zzz_c_h2048.log"), "#5f6b7a"),
    ("Parcae rec, seed 1",
     cells("eval_code_parcae16_s1_e400.json")
     or cells_from_log(LOGD / "zzz_d_parcae_s1.log"), "#276749"),
]

fig, axes = plt.subplots(2, 5, figsize=(16, 6.6), sharey=True)
axf = axes.flatten()
for ax, (title, kc, dark) in zip(axf, PANELS):
    ax.grid(True, color="#ececec", lw=0.7)
    ax.set_axisbelow(True)
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    ax.set_title(title, fontsize=8.6)
    ax.set_ylim(0, 1.02)
    if not kc:
        ax.text(0.5, 0.5, "pending", transform=ax.transAxes, ha="center",
                color="#999", fontsize=10)
        continue
    ks = sorted(kc)
    xs = range(len(ks))
    ax.plot(xs, [kc[k][0] for k in ks], "-o", color=dark, alpha=0.32,
            lw=2.2, ms=4.5)
    ax.plot(xs, [kc[k][1] for k in ks], "-s", color=dark, lw=2.2, ms=4.5)
    # references: base levels (dotted) and distill hard (dashed amber)
    ax.axhline(kc[ks[0]][0], color="#999", lw=0.8, ls=":")
    ax.axhline(kc[ks[0]][1], color="#999", lw=0.8, ls=":")
    ax.axhline(0.457, color="#b7791f", lw=1.0, ls="--", alpha=0.7)
    ax.set_xticks(list(xs))
    ax.set_xticklabels([str(k) for k in ks], fontsize=7.5)
    ax.set_xlabel("k", fontsize=8)
for i in (0, 5):
    axf[i].set_ylabel("pass@1", fontsize=9)

from matplotlib.lines import Line2D
axf[0].legend(handles=[
    Line2D([], [], color="#555", alpha=0.32, marker="o", lw=2.2, label="easy"),
    Line2D([], [], color="#555", marker="s", lw=2.2, label="hard"),
    Line2D([], [], color="#b7791f", ls="--", lw=1.0,
           label="distill hard (45.7%, no loop)"),
    Line2D([], [], color="#999", ls=":", lw=0.8, label="k=0 base levels"),
], fontsize=7, frameon=False, loc="center right")

fig.suptitle("The design space, measured — depth curves for every regime "
             "(same frozen band, data, 250-item eval; ★ = recommended recipe; "
             "amber dashes = what distillation reaches with no loop at all)",
             fontsize=11, y=1.0)
fig.tight_layout()
fig.savefig(OUT / "fig_kcurves2.png", dpi=140, facecolor="white",
            bbox_inches="tight")
print("wrote", OUT / "fig_kcurves2.png")