"""Figure: Adam's denominator vs the gradient under the clamp (p=97 MLP, N=82, eps 1e-8, seed 0), beta2 = 0.98 vs 0.999.
Data: /work/grokking/analysis/update_size_b2-*.csv (update_size_b2.py). Output: fig_denominator.png"""
import csv, 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]}
D={0.98:load('/work/grokking/analysis/update_size_b2-0.98_N82_eps1e-8.csv'),0.999:load('/work/grokking/analysis/update_size_b2-0.999_N82_eps1e-8.csv')}
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,2,figsize=(9,7.2),sharex='col',gridspec_kw=dict(hspace=0.12,wspace=0.22))
for j,b2 in enumerate([0.98,0.999]):
    d=D[b2]; s=d['step']
    a=ax[0,j]; a.plot(s,d['rms_grad'],color=BLUE,lw=2,label='RMS gradient g'); a.plot(s,d['rms_sqrt_vhat'],color=ORANGE,lw=2,label='RMS sqrt(v̂)  (Adam denominator)')
    a.axhline(1e-8,color=INK2,lw=1,ls=(0,(3,3))); a.text(s[-1],1.3e-8,'eps = 1e-8',ha='right',va='bottom',color=INK2,fontsize=8)
    a.set_yscale('log'); a.set_ylim(1e-10,2e-3); a.set_title(f'beta2 = {b2}',loc='left',fontsize=10,fontweight='bold')
    a=ax[1,j]; a.plot(s,d['rms_step_over_lr'],color=AQUA,lw=2); a.set_ylim(0,1.05)
    a=ax[2,j]; a.plot(s,d['test_acc'],color=INK,lw=2); a.set_ylim(0,1.05); a.set_xlabel('training step')
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'); ax[2,0].set_ylabel('test accuracy')
ax[0,0].legend(frameon=False,loc='lower left',fontsize=8)
ax[1,0].text(2500,0.08,'drops only once g reaches eps',color=INK2,fontsize=8,ha='right')
ax[1,1].text(600,0.3,'damped from step ~500: sqrt(v̂) lags g',color=INK2,fontsize=8)
fig.suptitle('Clamped MLP, N = 82, eps = 1e-8: the denominator tracks g at beta2 = 0.98 but lags it at beta2 = 0.999',x=0.01,ha='left',fontsize=10)
fig.savefig('fig_denominator.png',dpi=150,bbox_inches='tight',facecolor='#fcfcfb')
print('ok')