"""Ladder: what the loop adds over the base model and the same implant
applied once (feedforward). Same 250-item MBPP eval throughout.
Per-arm colors follow fig_phase/fig_kcurves; easy pastel, hard dark,
overall neutral.
"""
from pathlib import Path

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

OUT = Path(__file__).resolve().parent.parent / "results-loop"

#        arm                                       overall easy   hard   color
ROWS = [
    ("base model (k=0)",                            0.488, 0.984, 0.036, "#8a8f98"),
    ("+ FF implant\n(same adapter, applied once)",  0.496, 0.934, 0.179, "#8a8f98"),
    ("+ untrained loop (best k)",                   0.476, 0.892, 0.179, "#8a8f98"),
    ("+ trained loop, random-depth\n(randk, k=8)",  0.512, 0.902, 0.393, "#2b6cb0"),
    ("+ trained loop, curriculum\n(merge, k=4)",    0.512, 0.885, 0.464, "#2b6cb0"),
    ("+ plan-distilled implant\n(no recurrence; 8-seed mean)*", 0.555, 0.919, 0.457, "#b7791f"),
]

fig, ax = plt.subplots(figsize=(8.6, 4.6))
H = 0.22
for i, (name, ov, easy, hard, c) in enumerate(ROWS):
    y = len(ROWS) - 1 - i
    ax.barh(y + H, hard, height=H, color=c, zorder=3)
    ax.barh(y, easy, height=H, color=c, alpha=0.32, zorder=3)
    ax.barh(y - H, ov, height=H, color="#c9ccd1", zorder=3)
    for v, dy in ((hard, H), (easy, 0), (ov, -H)):
        ax.text(v + 0.012, y + dy, f"{v:.1%}", va="center", fontsize=8,
                color="#333")
ax.set_yticks([len(ROWS) - 1 - i for i in range(len(ROWS))])
ax.set_yticklabels([r[0] for r in ROWS], fontsize=9)
ax.set_xlim(0, 1.09)
ax.set_xlabel("pass@1", fontsize=9)
ax.xaxis.set_major_formatter(lambda x, _: f"{x:.0%}")

import matplotlib.patches as mpatches
ax.legend(handles=[
    mpatches.Patch(color="#555a61", label="hard (plan-dependent, n=28)"),
    mpatches.Patch(color="#555a61", alpha=0.32, label="easy (n=122)"),
    mpatches.Patch(color="#c9ccd1", label="overall (n=250)"),
], fontsize=8, frameon=False, loc="lower right")

ax.grid(True, axis="x", color="#ececec", lw=0.7)
ax.set_axisbelow(True)
for s in ("top", "right"):
    ax.spines[s].set_visible(False)
ax.set_title("What recurrence adds — and what distillation matches without it\n"
             "(MBPP; FF = trained adapter without recurrence)", fontsize=10.5)
fig.text(0.13, -0.015, "*distill rows evaluated on the 500-item set (hard "
         "n=55, base hard 3.6-5.5%); all other rows on the 250-item set. "
         "Distill trains the same adapter on written plans, deploys with "
         "zero loops.", fontsize=7.5, color="#555")
fig.tight_layout()
fig.savefig(OUT / "fig_loop_vs_ff.png", dpi=140, facecolor="white",
            bbox_inches="tight")
print("wrote", OUT / "fig_loop_vs_ff.png")