"""Figure: Adam's variables inside the clamped transformer (N=60), eps 1e-8 (and 1e-16 if available): RMS gradient, RMS sqrt(v_hat), actual RMS step/lr, test acc.
Data: /work/grokking/baseline/var_clamp/probe_N60_eps*/{probe,log}.csv (train_probe.py). Output: fig_transformer_probe.png"""
import csv, os, matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
BLUE,ORANGE,AQUA='#2a78d6','#eb6834','#1baf7a'; INK='#0b0b0b'; INK2='#52514e'; GRID='#e6e5e2'
def load(p): r=list(csv.DictReader(open(p))); return {k:[float(x[k]) for x in r] for k in r[0]}
runs=[e for e in ['1e-8','1e-16'] if os.path.exists(f'/work/grokking/baseline/var_clamp/probe_N60_eps{e}/probe.csv')]
XMAX=8000
plt.rcParams.update({'font.size':9,'axes.edgecolor':INK2,'axes.labelcolor':INK,'xtick.color':INK2,'ytick.color':INK2,'text.color':INK})
fig,ax=plt.subplots(3,len(runs),figsize=(5.2*len(runs),7),sharex=True,squeeze=False,gridspec_kw=dict(hspace=0.12,wspace=0.2))
for j,e in enumerate(runs):
    P=load(f'/work/grokking/baseline/var_clamp/probe_N60_eps{e}/probe.csv'); L=load(f'/work/grokking/baseline/var_clamp/probe_N60_eps{e}/log.csv')
    s=P['step']; a=ax[0,j]; a.plot(s,P['rms_grad'],color=BLUE,lw=1.5,label='RMS gradient g'); a.plot(s,P['rms_sqrt_vhat'],color=ORANGE,lw=1.5,label='RMS sqrt(v̂)')
    a.axhline(float(e),color=INK2,lw=1,ls=(0,(3,3))); a.text(XMAX,float(e)*1.5,f'eps = {e}',ha='right',va='bottom',color=INK2,fontsize=8)
    a.set_yscale('log'); a.set_ylim(1e-17 if e=='1e-16' else 1e-10,1); a.set_title(f'eps = {e}',loc='left',fontsize=10,fontweight='bold')
    a=ax[1,j]; a.plot(s,P['rms_step_over_lr'],color=AQUA,lw=1.5); a.set_yscale('log'); a.set_ylim(5e-4,2)
    a=ax[2,j]; a.plot(L['step'],L['test_acc'],color=INK,lw=1.5); a.set_ylim(0,1.05); a.set_xlabel('training step'); a.set_xlim(0,XMAX)
for a in ax.flat: a.grid(True,color=GRID,lw=1); a.set_axisbelow(True); a.spines[['top','right']].set_visible(False)
ax[0,0].set_ylabel('per-coordinate scale'); ax[1,0].set_ylabel('RMS step / lr (log)'); ax[2,0].set_ylabel('test accuracy'); ax[0,0].legend(frameon=False,loc='lower left',fontsize=8)
fig.suptitle('Clamped transformer, N = 60: each loss spike inflates sqrt(v̂); the step stays at 0.1 to 1% of lr while it decays',x=0.01,ha='left',fontsize=10)
fig.savefig('fig_transformer_probe.png',dpi=150,bbox_inches='tight',facecolor='#fcfcfb'); print('ok',runs)