"""fig_beta2_sweep.py: (left) T_grok at N=88, p=97 vs 1-beta2 for eps 1e-8 / 1e-16; (right) measured gradient collapse rate r vs (1-beta2)/2.
Reads beta2_sweep.csv (written by b2sweep_analyze.py). Run with /work/venv-plot/bin/python."""
import csv, numpy as np, matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
D='/work/grokking/generality'; R=[r for r in csv.DictReader(open(f'{D}/beta2_sweep.csv')) if r['N']=='88']
C={'1e-08':'#2a78d6','1e-16':'#eb6834'}; LAB={'1e-08':'eps 1e-8','1e-16':'eps 1e-16'}
fig,(ax,bx)=plt.subplots(1,2,figsize=(11,4.4),dpi=150)
for eps,dx in (('1e-08',0.96),('1e-16',1.04)):
    rs=[r for r in R if r['eps']==eps]; xs=sorted({1-float(r['beta2']) for r in rs},reverse=True)
    for r in rs:
        x=(1-float(r['beta2']))*dx; y=float(r['step_or_last'])
        if r['censored']=='1': ax.plot(x,y,marker='^',ms=8,mfc='none',mec=C[eps],mew=1.5,ls='none',zorder=3)
        else: ax.plot(x,y,'o',ms=6,color=C[eps],mec='white',mew=1,zorder=3)
    mean=[np.mean([float(r['step_or_last']) for r in rs if abs(1-float(r['beta2'])-x)<1e-9]) for x in xs]
    ax.plot(xs,mean,color=C[eps],lw=2,label=LAB[eps]+' (line: seed mean)',zorder=2)
    ax.text(0.0145 if eps=='1e-08' else 0.032,7600 if eps=='1e-08' else 2450,LAB[eps],color='#333',fontsize=9,va='center')
ax.plot([],[],marker='^',mfc='none',mec='#555',ls='none',label='censored at 20000 (lower bound)')
ax.set_xscale('log'); ax.set_yscale('log'); ax.invert_xaxis(); ax.set_xlim(0.062,0.0008)
ax.annotate('eps 1e-8: all 3 seeds censored\n(mean is a lower bound)',(0.001,20000),xytext=(0.0022,3600),textcoords='data',ha='center',fontsize=8,color='#333',arrowprops=dict(arrowstyle='-',color='#999',lw=0.8))
ax.set_xticks([0.05,0.02,0.015,0.01,0.005,0.001]); ax.xaxis.set_minor_locator(matplotlib.ticker.NullLocator()); ax.set_xticklabels(['0.95','0.98','0.985','0.99','0.995','0.999'],rotation=35)
ax.set_xlabel('Adam beta2  (log scale of 1 - beta2)'); ax.set_ylabel('steps to test acc >= 0.99')
ax.set_title('Time to grok, N=88, p=97 (seeds 0,1; 0.999 also seed 2)',fontsize=10)
ax.legend(fontsize=8,frameon=False,loc='upper left')
for r in R:
    if r['r_per1000']=='nan': continue
    h=(1-float(r['beta2']))/2*1000
    bx.plot(h*(0.95 if r['eps']=='1e-08' else 1.05),float(r['r_per1000']),'o' if r['eps']=='1e-16' else 's',ms=6,color=C[r['eps']],mec='white',mew=1,ls='none')
xx=np.array([0.3,30]); bx.plot(xx,xx,color='#999',lw=1,ls='--'); bx.text(0.35,0.42,'r = (1-beta2)/2',color='#666',fontsize=8,rotation=38)
bx.axhspan(6.9,9.3,color='#1baf7a',alpha=0.15,lw=0); bx.text(0.35,10.2,'undamped r0 (beta2 0.95/0.98, eps 1e-16)',color='#333',fontsize=8)
for eps in C: bx.plot([],[],'o' if eps=='1e-16' else 's',color=C[eps],ls='none',label=LAB[eps])
bx.set_xscale('log'); bx.set_yscale('log'); bx.set_xlim(0.3,40); bx.set_ylim(0.3,40)
bx.set_xlabel('(1 - beta2)/2   per 1000 steps'); bx.set_ylabel('measured collapse rate r of grad RMS, per 1000 steps')
bx.set_title('Collapse rate r vs (1-beta2)/2 (N=88)',fontsize=10)
bx.legend(fontsize=8,frameon=False,loc='lower right')
for a in (ax,bx):
    a.grid(alpha=0.25,lw=0.6); [a.spines[s].set_visible(False) for s in ('top','right')]
fig.tight_layout(); fig.savefig(f'{D}/fig_beta2_sweep.png'); print('saved')