"""Cross-model depth-regimes figure: sensor / motor / persistence curves by
layer, one panel per lens-mapped model, workspace band shaded (from
results/REGIMES.json). Replaces the ad-hoc single-model regimes.png.

Panel sources: per-model regimes*.pt diagnostics (exp4_regimes.py output:
dict with per-layer 'sensor', 'motor', 'persist' fractions). Pending models
render as placeholders and fill in on re-render.
"""
import json
from pathlib import Path

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

ROOT = Path(__file__).resolve().parent.parent
REG = json.load(open(ROOT / "results/REGIMES.json"))

PANELS = [
    ("gemma-4-E2B (2B eff.)", "results/regimes.pt", "google/gemma-4-E2B-it"),
    ("gemma-4-E4B (4B eff., E2B ⊂ E4B)", "results/regimes_e4b.pt",
     "google/gemma-4-E4B-it"),
    ("gemma-4-12B", "results-12b/regimes.pt", "google/gemma-4-12B-it"),
    ("gemma-4-26B-A4B (MoE)", "results/regimes_26b_a4b.pt",
     "google/gemma-4-26B-A4B-it"),
    ("gemma-4-31B", "results/regimes_31b.pt", "google/gemma-4-31B-it"),
]

fig, axes = plt.subplots(1, len(PANELS), figsize=(3.4 * len(PANELS), 3.4),
                         sharey=True)
for ax, (title, ptfile, key) in zip(axes, PANELS):
    ax.set_title(title, fontsize=9)
    ax.set_xlabel("layer", fontsize=8)
    ax.grid(True, color="#ececec", lw=0.6)
    ax.set_axisbelow(True)
    for s in ("top", "right"):
        ax.spines[s].set_visible(False)
    p = ROOT / ptfile
    if not p.exists():
        ax.text(0.5, 0.5, "scan pending", transform=ax.transAxes,
                ha="center", color="#999", fontsize=10)
        continue
    d = torch.load(p, map_location="cpu")
    sensor = torch.as_tensor(d["sensor"]).float()
    motor = torch.as_tensor(d["motor"]).float()
    persist = torch.as_tensor(d["persist"]).float()
    L = len(sensor)
    xs = range(L)
    band = REG.get(key, {}).get("workspace_band")
    if band:
        ax.axvspan(band[0], band[1], color="#2f855a", alpha=0.10)
        ax.text(sum(band) / 2, 0.51, "workspace", color="#2f855a",
                fontsize=7, ha="center")
        # sensor regime: contiguous pre-band region where the sensor
        # readout dominates (rule reproduces the hand-shaded E2B bands)
        sens = [i for i, (s, m) in enumerate(zip(sensor, motor))
                if s > m and s >= 0.05 and i < band[0]]
        if sens:
            ax.axvspan(min(sens), max(sens) + 1, color="#2b6cb0", alpha=0.08)
            ax.text((min(sens) + max(sens) + 1) / 2, 0.51, "sensor",
                    color="#2b6cb0", fontsize=7, ha="center")
        if band[1] + 1 < L - 1:
            ax.axvspan(band[1] + 1, L - 1, color="#c53030", alpha=0.08)
            ax.text((band[1] + L) / 2, 0.51, "motor", color="#c53030",
                    fontsize=7, ha="center")
    kv = REG.get(key, {}).get("kv_share_start")
    if isinstance(kv, int):
        ax.axvline(kv, color="#c53030", lw=0.9, ls=":")
    ax.plot(xs, sensor, color="#2b6cb0", lw=1.6, label="sensor")
    ax.plot(xs, motor, color="#c53030", lw=1.6, label="motor")
    ax.plot(xs, persist, color="#2f855a", lw=1.4, ls="--", label="persistence")
axes[0].set_ylabel("fraction", fontsize=9)
axes[0].legend(fontsize=7.5, frameon=False, loc="upper left")
fig.suptitle("Depth regimes under the J-lens across gemma-4 scales — green "
             "shade: workspace band (REGIMES.json); red dotted: KV-share "
             "boundary (entrances above are structurally null)",
             fontsize=10.5, y=1.04)
fig.tight_layout()
out = ROOT / "results/regimes_all_models.png"
fig.savefig(out, dpi=140, facecolor="white", bbox_inches="tight")
print("wrote", out)