"""Schematic: the J-lens view of the residual stream + depth regimes + looping."""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch, Rectangle
from matplotlib.patches import ConnectionPatch

fig = plt.figure(figsize=(13, 8.5))
gs = fig.add_gridspec(1, 2, width_ratios=[1.15, 1], wspace=0.05)
axL = fig.add_subplot(gs[0]); axR = fig.add_subplot(gs[1])
for a in (axL, axR):
    a.set_xlim(0, 10); a.set_ylim(0, 10); a.axis("off")

# ---------- LEFT: residual stream with regime bands (gemma-4-E2B, 35 layers) ----------
axL.text(5, 9.6, "The residual stream under the J-lens", ha="center",
         fontsize=15, fontweight="bold")
axL.text(5, 9.18, "gemma-4-E2B  ·  35 layers, d=1536", ha="center",
         fontsize=9.5, color="#555")

# central stream bar
sx, sw = 4.15, 1.3
axL.add_patch(Rectangle((sx, 0.7), sw, 7.7, facecolor="#f2f2f2",
                        edgecolor="#888", lw=1.2, zorder=1))
axL.annotate("", xy=(sx+sw/2, 8.75), xytext=(sx+sw/2, 0.55),
             arrowprops=dict(arrowstyle="-|>", color="#444", lw=2), zorder=0)
axL.text(sx+sw/2, 0.32, "token embedding (input)", ha="center", fontsize=8, color="#333")
axL.text(sx+sw/2, 8.95, "logits (output)", ha="center", fontsize=8, color="#333")

def y(layer): return 0.8 + (layer/34)*7.35
bands = [
    ("transduction", 0, 5,  "#d9d9d9", "detokenization /\nlexical assembly",
     ["J-lens: noise / surface form"]),
    ("sensor", 6, 13, "#bcd4f0", "input written\ninto the stream",
     ["J-lens: current input token", "(echo ↑ 26%)"]),
    ("workspace", 14, 30, "#bfe3c0", "abstract concepts\nheld & broadcast",
     ["J-lens: unspoken content — 'spider', 'big'",
      "≈ 1.08B params (58% of decoder)",
      "swap-gates fire here · ignition threshold"]),
    ("motor", 31, 34, "#f2c4c4", "output staged\nfor unembedding",
     ["J-lens: next output token (→68%)"]),
]
for name, l0, l1, c, fn, lens_lines in bands:
    yb, yt = y(l0)-0.10, y(l1)+0.10
    ymid = (yb+yt)/2
    axL.add_patch(Rectangle((sx, yb), sw, yt-yb, facecolor=c, edgecolor="#666",
                            lw=1, alpha=0.95, zorder=2))
    axL.text(sx+sw/2, ymid, f"L{l0}–{l1}", ha="center", va="center",
             fontsize=8.5, fontweight="bold", zorder=3)
    # band name above-left inside its band
    axL.text(sx-0.2, ymid, name, ha="right", va="center",
             fontsize=10.5, fontweight="bold", color="#222")
    # function (right, top) + lens lines stacked below
    axL.text(sx+sw+0.3, ymid+0.28, fn, ha="left", va="center", fontsize=8, color="#222")
    for j, ln in enumerate(lens_lines):
        axL.text(sx+sw+0.3, ymid-0.16-0.28*j, ln, ha="left", va="center",
                 fontsize=6.9, color="#556", style="italic")

# three verbs: write in (sensor), hold (workspace), read out (motor) — offset from
# the band-name row so arrows don't clip the labels
axL.annotate("", xy=(sx-0.02, y(12.3)), xytext=(sx-0.7, y(12.3)),
             arrowprops=dict(arrowstyle="-|>", color="#2b6cb0", lw=1.8))
axL.text(sx-0.72, y(12.3)+0.2, "write in", ha="right", fontsize=7.5, color="#2b6cb0")
axL.add_patch(FancyArrowPatch((sx-0.08, y(28.5)), (sx-0.08, y(25.5)),
              connectionstyle="arc3,rad=-1.0", arrowstyle="-|>",
              color="#2f855a", lw=1.7, mutation_scale=11))
axL.text(sx-0.62, y(27), "hold", ha="right", va="center", fontsize=7.5, color="#2f855a")
axL.annotate("", xy=(sx-0.02, y(33.7)), xytext=(sx-0.7, y(33.7)),
             arrowprops=dict(arrowstyle="-|>", color="#c53030", lw=1.8))
axL.text(sx-0.72, y(33.7)+0.2, "read out", ha="right", fontsize=7.5, color="#c53030")

# ---------- RIGHT: the looping result ----------
axR.text(5, 9.7, "Can we loop the workspace?", ha="center",
         fontsize=15, fontweight="bold")

def box(ax, x, y0, w, h, fc, ec, txt, fs=8, fw="normal", tc="#111"):
    ax.add_patch(FancyBboxPatch((x, y0), w, h, boxstyle="round,pad=0.06,rounding_size=0.12",
                 facecolor=fc, edgecolor=ec, lw=1.4))
    ax.text(x+w/2, y0+h/2, txt, ha="center", va="center", fontsize=fs, fontweight=fw, color=tc)

# Case A: band loop (not a self-map) -- fails
axR.text(0.4, 8.9, "A.  Loop the whole band  (L14→30)", fontsize=10.5, fontweight="bold")
box(axR, 0.6, 7.3, 2.2, 1.0, "#bfe3c0", "#2f855a", "band\nL14 → L30")
axR.add_patch(FancyArrowPatch((2.9, 8.15), (2.9, 7.45),
              connectionstyle="arc3,rad=1.3", arrowstyle="-|>", color="#c53030",
              lw=2, mutation_scale=14))
axR.text(4.9, 7.8, "out lives 17 layers\ndownstream of in", fontsize=8, color="#333")
axR.text(0.6, 6.75, "✗  NOT a self-map (out-space ≠ in-space)",
         fontsize=9, color="#c53030", fontweight="bold")
axR.text(0.6, 6.4, "     → collapses in 1 iteration:  '8' → ')'",
         fontsize=8.3, color="#555")

axR.plot([0.4, 9.6], [5.95, 5.95], color="#ccc", lw=1)

# Case B: single tied layer (self-map)
axR.text(0.4, 5.55, "B.  Loop one tied layer   h ← h + β·Δ(h)", fontsize=10.5, fontweight="bold")
box(axR, 0.6, 4.0, 2.2, 1.0, "#bcd4f0", "#2b6cb0", "layer L\nin = out")
axR.add_patch(FancyArrowPatch((2.9, 4.85), (2.9, 4.15),
              connectionstyle="arc3,rad=1.3", arrowstyle="-|>", color="#2b6cb0",
              lw=2, mutation_scale=14))
axR.text(4.9, 4.5, "same residual point\nin and out", fontsize=8, color="#333")
axR.text(0.6, 3.45, "✓  self-map — stable; damping (low-pass) extends it",
         fontsize=9, color="#2f855a", fontweight="bold")
axR.text(0.6, 3.05, "     β=0.5 @ L26 keeps '8' through ~6 loops; |ΔH| ↓ (converges)",
         fontsize=8.3, color="#555")
axR.text(0.6, 2.6, "✗  but frozen → degenerate fixed point; concept never sharpens",
         fontsize=9, color="#c53030", fontweight="bold")
axR.text(0.6, 2.2, "     extra loop compute ≠ extra reasoning",
         fontsize=8.3, color="#555")

# verdict box
box(axR, 0.6, 0.5, 9.0, 1.35, "#fff7e6", "#d69e2e",
    "Verdict:  mechanically loopable at the layer level, and low-pass\n"
    "damping is necessary for stability — but not sufficient. Turning loop-depth\n"
    "into reasoning needs the loop trained in (LoRA-per-loop / Huginn / MoR).",
    fs=8.6, fw="normal", tc="#7a5a00")

fig.savefig("results/jlens_diagram.png", dpi=140, bbox_inches="tight",
            facecolor="white")
print("wrote results/jlens_diagram.png")
