"""Small-multiples: pass@1 vs loop depth k, easy (pastel) + hard (dark),
one panel per regime arm. All on the same 250-item MBPP eval, e400.
Per-depth cells fall back to log-parsed values until its JSON lands.
"""
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"

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

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\n(training-free, ρ=0.3)",
     cells("eval_code_untrained.json"), "#2b6cb0"),
    ("trained merge\n(fixed A,B · curriculum)",
     cells("eval_code_trained.json"), "#2b6cb0"),
    ("unconstrained rec\n(learned A,B · noise h₀ · rand-k) ρ→4.5",
     cells("eval_code_rec16_partial.json"), "#c53030"),
    ("Parcae rec\n(ρ<1 enforced · learned B · noise h₀) ρ→0.29",
     cells("eval_code_parcae16_e400.json"), "#2f855a"),
    ("per-depth merges\n(time-varying · anchored · curriculum)",
     cells("eval_code_pd4_e400.json")
     or cells_from_log("/home/nils/gpuq/spark-gpu0/perdepth_arm.log"),
     "#805ad5"),
]

fig, axes = plt.subplots(1, 5, figsize=(15.5, 3.6), sharey=True)
for ax, (title, kc, dark) in zip(axes, PANELS):
    ax.grid(True, color="#ececec", lw=0.7)
    ax.set_axisbelow(True)
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    if not kc:
        ax.set_title(title + "\n(pending)", fontsize=9)
        continue
    ks = sorted(kc)
    xs = range(len(ks))
    easy = [kc[k][0] for k in ks]
    hard = [kc[k][1] for k in ks]
    # pastel = dark color at low alpha via manual blend to white
    ax.plot(xs, easy, "-o", color=dark, alpha=0.32, lw=2.4, ms=5,
            label="easy")
    ax.plot(xs, hard, "-s", color=dark, lw=2.4, ms=5, label="hard")
    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.set_xticks(list(xs))
    ax.set_xticklabels([str(k) for k in ks], fontsize=8)
    ax.set_title(title, fontsize=9)
    ax.set_xlabel("k", fontsize=9)
    ax.set_ylim(0, 1.02)
axes[0].set_ylabel("pass@1 (easy pastel / hard dark)")
axes[0].legend(fontsize=8, frameon=False, loc="center right")
fig.suptitle("Depth curves by regime — same frozen band, data, and 250-item "
             "eval; dotted lines = k=0 base levels", fontsize=11, y=1.04)
fig.tight_layout()
fig.savefig(OUT / "fig_kcurves.png", dpi=140, facecolor="white",
            bbox_inches="tight")
print("wrote", OUT / "fig_kcurves.png")