"""Primer deck figures: schedule, 72M blowup (real logs), loop-gain, beta window.
PNGs to ../assets/ for the pptx build."""
import re
import matplotlib
matplotlib.use('Agg')
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch

INK, GRAY = '#3a3a3a', '#8a8a8a'
BOXF, BOXE = '#f4f4f4', '#9a9a9a'
BLUE, BLUEF = '#2c6fbb', '#eaf1fa'
ORAN, ORANF = '#d95f02', '#fdeee2'
RED = '#b03a2e'
CREAM, CREME = '#faf6ee', '#c9b895'
GREEN = '#2e7d32'
A = '/home/yurenh2/ept/assets/'

def newfig(w, h, xmax=100, ymax=62):
    fig = plt.figure(figsize=(w, h))
    ax = fig.add_axes([0, 0, 1, 1]); ax.set_xlim(0, xmax); ax.set_ylim(0, ymax); ax.axis('off')
    return fig, ax

def box(ax, xc, yc, w, s, fc=BOXF, ec=BOXE, size=10, lw=1.2, h=4.6, tc=INK, wt='normal'):
    ax.add_patch(FancyBboxPatch((xc - w / 2, yc - h / 2), w, h,
                 boxstyle='round,pad=0.3,rounding_size=0.7', fc=fc, ec=ec, lw=lw))
    ax.text(xc, yc, s, fontsize=size, color=tc, ha='center', va='center', fontweight=wt)

def txt(ax, x, y, s, size=9.5, color=INK, ha='center', va='center', w='normal', style='normal'):
    ax.text(x, y, s, fontsize=size, color=color, ha=ha, va=va, fontweight=w, fontstyle=style)

def arrow(ax, x0, y0, x1, y1, color=INK, lw=1.4, ms=12, cs=None):
    ax.annotate('', xy=(x1, y1), xytext=(x0, y0),
                arrowprops=dict(arrowstyle='-|>', color=color, lw=lw, mutation_scale=ms,
                                shrinkA=0.5, shrinkB=0.5, connectionstyle=cs))

# ================= fig S3: the update schedule =================
fig, ax = newfig(11.8, 6.0)
dx, zx = 15, 41
txt(ax, dx, 58.5, 'mismatches derive TOP → BOTTOM', 10, ORAN, w='bold')
txt(ax, zx, 58.5, 'states rebuild BOTTOM → TOP', 10, BLUE, w='bold')
dl = [r'$d_3 = -\beta\,\nabla\mathrm{CE}$', r'$d_2 = J_3^{\top} d_3$', r'$d_1 = J_2^{\top} d_2$']
dn = ['needs: loss only', 'needs: block 3 only', 'needs: block 2 only']
for i, (s, n) in enumerate(zip(dl, dn)):
    y = 51 - i * 9
    box(ax, dx, y, 20, s, fc=ORANF, ec=ORAN)
    txt(ax, dx, y - 3.4, n, 7.6, GRAY)
    if i < 2: arrow(ax, dx, y - 2.4 - 2.2, dx, y - 9 + 2.5, color=ORAN)
zl = [r'$z_1 = f_1(\mathrm{emb}) + d_1$', r'$z_2 = f_2(z_1) + d_2$', r'$z_3 = f_3(z_2) + d_3$']
for i, s in enumerate(zl):
    y = 33 + i * 9
    box(ax, zx, y, 20, s, fc=BLUEF, ec=BLUE)
    if i < 2: arrow(ax, zx, y + 2.4, zx, y + 9 - 2.5, color=BLUE)
txt(ax, 28, 22.5, 'block $l$ touches NEIGHBORS only — $d$ from the block after it,\n'
                  '$z$ from the block before it.  $\\Delta\\theta_l$ fires as soon as $d_l$ arrives.\n'
                  'Repeat $K$ times ($K{=}3$ measured sufficient; $K{=}8$ identical).', 9)

sx = 76
txt(ax, sx, 58.5, 'the same settle, three schedules', 10.5, INK, w='bold')
rows = [
    ('sequential  (our trainer today)', r'$O(K\!\cdot\!d)$', '3 × 12 = 36 block-ops', BOXF, BOXE),
    ('pipelined wavefront  (systolic; multi-GPU lever)', r'$O(K\!+\!d\!-\!1)$', '3 + 12 − 1 = 14 slots', BLUEF, BLUE),
    ('continuous-time analog:  all blocks relax AT ONCE', r'$\sim O(K)$', 'depth-serial factor gone;\n'
     r'$K$ dissolves into settling time $\tau$', ORANF, ORAN),
]
for i, (name, comp, det, fc, ec) in enumerate(rows):
    y = 50 - i * 12.5
    box(ax, sx, y, 44, '', fc=fc, ec=ec, h=9.5)
    txt(ax, sx - 20.5, y + 2.2, name, 9.3, ha='left', w='bold')
    txt(ax, sx - 20.5, y - 1.8, det, 8.6, GRAY, ha='left')
    txt(ax, sx + 15.5, y + 2.2, comp, 12)
txt(ax, sx, 9.5, 'open question we track: does $\\tau$ itself grow with depth?\n'
                 '(one transit down the chain is unavoidable)', 8.4, GRAY)
fig.savefig(A + 'fig_p3_schedule.png', dpi=170); plt.close(fig)

# ================= fig S4: what happened at 72M (real logs) =================
STEP_RE = re.compile(r'^step\s+(\d+)/\d+ \| train [\d.]+ val ([\d.]+) \(best ([\d.]+)\).*?beta=([\d.e+-]+)')
def parse(fn):
    out = {}
    for line in open('/home/yurenh2/ept/ep_run/runs/' + fn, errors='replace'):
        m = STEP_RE.match(line)
        if m: out[int(m.group(1))] = (float(m.group(2)), float(m.group(4)))
    return out

base = parse('fw72m.log')
resc = parse('fw72m_c.log'); resc.update(parse('fw72m_c2.log'))
bs = np.array(sorted(base)); bval = np.array([base[s][0] for s in bs])
rs = np.array(sorted(k for k in resc if k >= 195000)); rval = np.array([resc[s][0] for s in rs])
print('base max val post-195k:', bval[bs > 195000].max(), '| resc pts:', len(rs))
beta_steps = np.array(sorted([k for k in base if k < 195000] + list(rs)))
beta_v = np.array([base[k][1] if k < 195000 else resc[k][1] for k in beta_steps])

fig = plt.figure(figsize=(9.2, 6.0))
axv = fig.add_axes([0.085, 0.12, 0.815, 0.78])
axv.plot(bs / 1000, np.clip(bval, 0, 6.42), color=BLUE, lw=1.0, zorder=3)
axv.plot(rs / 1000, np.clip(rval, 0, 6.42), color='#999999', lw=0.9, zorder=2)
axv.text(233, 4.32, 'rescue resumes\n(new recipe from 195k / 215k)', fontsize=8, color='#777777', ha='right')
axv.axhline(3.2884, color=GRAY, lw=1.0, ls='--')
axv.text(2, 3.19, 'BP twin final 3.2884', fontsize=8.5, color=GRAY, va='top')
axv.plot([185], [3.7117], marker='*', ms=14, color=GREEN, zorder=5)
axv.text(180, 3.50, 'best 3.7117 @185k  (never beaten again)', fontsize=8.5, color=GREEN, ha='right')
axv.annotate('195k: relaxation diverges — val 3.9 → 6.0,\nguard storm (15k step-skips); the same\n$\\beta$ was safe for the previous 135k steps',
             xy=(204, 5.95), xytext=(112, 5.62), fontsize=9.5, color=RED,
             arrowprops=dict(arrowstyle='-|>', color=RED))
axv.axvline(60, color=GRAY, lw=0.8, ls=':')
axv.axvline(215, color=GRAY, lw=0.8, ls=':')
axv.text(61, 6.30, 'floor raised\n3e-4 → 1e-3: fine', fontsize=8, color=GRAY, va='top')
axv.text(146, 4.78, 'governor caps $\\beta$: survives,\nbut starved to 2e-5 — best frozen', fontsize=8.5, color=INK)
axv.set_xlabel('step (×1000)', fontsize=10); axv.set_ylabel('val CE', fontsize=10, color=BLUE)
axv.set_ylim(3.0, 6.5); axv.set_xlim(0, 238)
axv.tick_params(labelsize=9)
axb = axv.twinx()
axb.plot(beta_steps / 1000, beta_v, color=ORAN, lw=1.3, alpha=0.9, drawstyle='steps-post')
axb.set_yscale('log'); axb.set_ylim(8e-6, 6e-3)
axb.set_ylabel(r'$\beta$ (log)', fontsize=10, color=ORAN)
axb.tick_params(labelsize=8, colors=ORAN)
for sp in ['top']: axv.spines[sp].set_visible(False); axb.spines[sp].set_visible(False)
axv.set_title('fw72m — 72.11M params × 1.44B FineWeb tokens, fully BP-free (what actually happened)',
              fontsize=11, color=INK, pad=10)
fig.savefig(A + 'fig_p4_blowup.png', dpi=170); plt.close(fig)

# ================= fig S5: the loop =================
fig, ax = newfig(8.8, 6.2, xmax=100, ymax=70)
bx = 26
for i, name in enumerate(['block 1', 'block 2', 'block 3']):
    box(ax, bx, 16 + i * 15, 22, name, fc=BOXF, ec=BOXE, h=8, size=11)
arrow(ax, bx + 12, 20, bx + 12, 27, color=BLUE, lw=2.0); arrow(ax, bx + 12, 35, bx + 12, 42, color=BLUE, lw=2.0)
arrow(ax, bx - 12, 42, bx - 12, 35, color=ORAN, lw=2.0); arrow(ax, bx - 12, 27, bx - 12, 20, color=ORAN, lw=2.0)
ax.text(bx - 17.5, 31, 'errors flow DOWN   $d_l = J_{l+1}^{\\top} d_{l+1}$', fontsize=9, color=ORAN,
        ha='center', va='center', rotation=90)
ax.text(bx + 17.5, 31, 'states rebuild UP   $z_l = f_l(z_{l-1}) + d_l$', fontsize=9, color=BLUE,
        ha='center', va='center', rotation=270)
arrow(ax, bx + 9, 51.5, bx - 9, 51.5, color=RED, lw=2.0, cs='arc3,rad=0.45')
txt(ax, bx, 60, 'top link: $-\\beta\\,\\nabla\\mathrm{CE}$, scale $\\sigma(W_{out})^2$', 9.5, RED)
txt(ax, bx, 6, 'one sweep = down + up:\na closed loop through the stack', 9.5, GRAY)

L = 50
box(ax, 72.5, 56, 45, r'per-sweep gain   $\rho \;\approx\; \beta\cdot\sigma(W_{out})^2\cdot\|J\mathrm{-chains}\|^2$',
    fc=CREAM, ec=CREME, h=8, size=10.5)
txt(ax, L, 46, r'$\rho < 1$:  mismatch decays every sweep → settles', 10, GREEN, ha='left')
txt(ax, L, 41, r'$\rho > 1$:  every sweep AMPLIFIES the residual → diverges', 10, RED, ha='left')
txt(ax, L, 34.8, 'and $\\rho$ RISES during training:', 10, INK, ha='left')
txt(ax, L, 31.3, r'$\sigma(W_{out})$ and $\|J\|$ grow as the fit deepens', 10, INK, ha='left')
txt(ax, L, 24, 'measured:', 9.5, INK, ha='left', w='bold')
txt(ax, L, 20.3, '· damping refuted in 3 doses (spectrum monotone-positive)', 9, INK, ha='left')
txt(ax, L, 16.8, '· lowering $\\beta$ restores convergence instantly', 9, INK, ha='left')
txt(ax, L, 13.3, '· continuous-time flow is unconditionally stable →', 9, INK, ha='left')
txt(ax, L + 1.5, 9.6, 'a DISCRETE-SOLVER disease; analog hardware\nhas no such ceiling', 9, ORAN, ha='left')
fig.savefig(A + 'fig_p5_loop.png', dpi=170); plt.close(fig)

# ================= fig S6: the beta window =================
fig = plt.figure(figsize=(9.2, 6.0))
ax2 = fig.add_axes([0.10, 0.13, 0.86, 0.76])
x = np.linspace(0, 234, 500)
cx = [0, 60, 120, 160, 195, 215, 222, 230, 234]
cy = [1.2e-2, 1.0e-2, 6e-3, 2.5e-3, 9.5e-4, 4e-4, 2.6e-4, 4e-5, 2.2e-5]
ceil = np.exp(np.interp(x, cx, np.log(cy)))
fx_ = [0, 60, 120, 180, 234]; fy = [8e-5, 1.1e-4, 1.6e-4, 2.6e-4, 4.5e-4]
floor = np.exp(np.interp(x, fx_, np.log(fy)))
ax2.fill_between(x, ceil, 2e-2, color=RED, alpha=0.10)
ax2.fill_between(x, 8e-6, floor, color='#666666', alpha=0.12)
ax2.fill_between(x, floor, ceil, color=GREEN, alpha=0.07)
ax2.plot(x, ceil, color=RED, lw=1.6)
ax2.plot(x, floor, color='#555555', lw=1.6)
ax2.text(6, 6.5e-3, r'wall-2:  $\rho>1$ — relaxation diverges', fontsize=10, color=RED)
ax2.text(6, 2.6e-5, 'wall-1:  SNR < 1 — updates dissolve into noise', fontsize=10, color='#555555')
ax2.text(96, 1.35e-3, r'$\beta_{max}(t) \propto \mathrm{margin}/(\sigma^2\|J\|^2)$', fontsize=9.5, color=RED, rotation=-14)
ax2.text(120, 1.05e-4, r'$\beta_{min}(t) \propto \mathrm{noise}/|g|$', fontsize=9.5, color='#555555', rotation=6)
ax2.text(150, 5.5e-4, 'operating window\n(narrows as fit deepens)', fontsize=9.5, color=GREEN, ha='center')
# original schedule
ax2.plot([0, 60], [3e-4, 3e-4], color=INK, lw=2.2)
ax2.plot([60, 60], [3e-4, 1e-3], color=INK, lw=1.0, ls=':')
ax2.plot([60, 195], [1e-3, 1e-3], color=INK, lw=2.2)
ax2.plot([195], [1e-3], marker='x', ms=13, mew=3, color=RED, zorder6b=5) if False else ax2.plot([195], [1e-3], marker='x', ms=13, mew=3, color=RED, zorder=5)
ax2.text(196.5, 1.25e-3, 'fixed floor crosses the\nfalling ceiling → blow', fontsize=8.5, color=RED)
ax2.text(90, 7.6e-4, 'fw72m schedule (fixed floors)', fontsize=8.5, color=INK)
# bcap rescue
bs = x[(x >= 215) & (x <= 234)]
ax2.plot(bs, np.exp(np.interp(bs, [215, 222, 230, 234], np.log([3e-4, 2.2e-4, 3.5e-5, 2e-5]))),
         color=GRAY, lw=1.8, ls='-.')
ax2.text(213, 4.8e-5, 'bcap rescue: survives, starved —\nthe window had CLOSED for the\nold recipe\'s margin', fontsize=8, color=GRAY, ha='right')
# cent path
xm = x[x <= 185]
cent = np.where(xm < 20, 3e-4, 3e-3)
cent = np.minimum(cent, np.exp(np.interp(xm, cx, np.log(cy))) * 0.85)
ax2.plot(xm, cent, color=ORAN, lw=2.2, ls='--')
ax2.text(24, 4.25e-3, 'fw72m_cent (in flight): centered + $\\beta$ rides ITS OWN ceiling via bcap 0.9 — where that ceiling sits is being measured now (48k: ahead)', fontsize=8.6, color=ORAN, ha='left')
ax2.text(186.5, 1.62e-3, '→ ?', fontsize=10, color=ORAN, fontweight='bold')
# measured ceiling points
for (mx, my, lab) in [(195, 1e-3, '≤1e-3'), (222, 3e-4, '≤3e-4'), (228, 2.5e-5, '≈2e-5')]:
    ax2.plot([mx], [my], marker='v', ms=7, color=RED, zorder=5)
ax2.text(231, 2.3e-3, 'measured ceiling\ncrossings ▾', fontsize=8, color=RED, ha='right')
ax2.set_yscale('log'); ax2.set_ylim(8e-6, 2e-2); ax2.set_xlim(0, 236)
ax2.set_xlabel('training progress (×1000 steps)', fontsize=10)
ax2.set_ylabel(r'$\beta$  (log scale)', fontsize=10)
ax2.tick_params(labelsize=9)
for sp in ['top', 'right']: ax2.spines[sp].set_visible(False)
ax2.set_title(r'$\beta$ is squeezed from both sides — and the window narrows as the model fits deeper',
              fontsize=11.5, color=INK, pad=10)
fig.savefig(A + 'fig_p6_window.png', dpi=170); plt.close(fig)
print('figs saved: p3_schedule, p4_blowup, p5_loop, p6_window')