"""Mechanism figure: test acc, train-set softmax saturation, raw gradient norm for N=50 and N=88 (eps 1e-8 vs 1e-16)."""
import csv, numpy as np, matplotlib
matplotlib.use('Agg'); import matplotlib.pyplot as plt
L='/work/grokking/temp-clamp/logs/'
runs=[('N=50, eps=1e-8', L+'total50_lr0.001_T1_seed0.csv', '#4C78A8', '-'),
      ('N=88, eps=1e-8', L+'total88_lr0.001_T1_seed0.csv', '#E45756', '-'),
      ('N=88, eps=1e-16', L+'total88_lr0.001_T1_seed0_eps1e-16.csv', '#E45756', '--')]
def load(p):
    r=list(csv.DictReader(open(p))); 
    return {k:np.array([float(x[k]) for x in r]) for k in r[0]}
NPARAM=97*64+(128*256+256)+(256*97+97)
fig,ax=plt.subplots(3,1,figsize=(6.4,7.2),sharex=True)
for lab,p,c,ls in runs:
    d=load(p); s=d['step']
    ax[0].plot(s,d['test_acc'],c,ls=ls,label=lab)
    ax[1].plot(s,d['frac_sat'],c,ls=ls)
    ax[2].semilogy(s,d['grad_norm']/np.sqrt(NPARAM),c,ls=ls)
ax[2].axhline(1e-8,color='k',lw=0.8,ls=':'); ax[2].text(4200,1.4e-8,'AdamW default eps = 1e-8',fontsize=8)
ax[0].set_ylabel('test accuracy'); ax[1].set_ylabel('fraction of training samples\nwith p(correct) > 1 - 1e-4'); ax[2].set_ylabel('RMS per-coordinate\ngradient (before Adam)')
ax[2].set_xlabel('training step (full batch)'); ax[0].legend(loc='lower right',fontsize=8,frameon=False)
ax[0].set_xlim(0,6000)
for a in ax: a.spines['top'].set_visible(False); a.spines['right'].set_visible(False)
fig.suptitle('Clamped-norm MLP, (a+b) mod 97: saturation, gradient underflow below eps, delayed grokking',fontsize=9)
fig.tight_layout(); fig.savefig('/work/publish/grokking-adam-eps/fig_mechanism.png',dpi=150)
# print key numbers
for lab,p,c,ls in runs:
    d=load(p); s=d['step']; g=d['grad_norm']/np.sqrt(NPARAM)
    i99=np.where(d['test_acc']>=0.99)[0]; isat=np.where(d['frac_sat']>=0.99)[0]
    print(lab,'test99 step',s[i99[0]] if len(i99) else None,'sat99 step',s[isat[0]] if len(isat) else None,'min rms grad %.2e'%g.min(),'median rms grad after sat %.2e'%(np.median(g[isat[0]:]) if len(isat) else float('nan')))