"""fig_beta2_model.py: T_grok at N=88 vs beta2 (data from beta2_sweep.csv) with the ODE model of analysis/staleness_fit.py
overlaid as a dotted line. The model gives relative times; it is anchored to the eps-1e-16 seed mean at beta2 0.98 (one factor
for both eps curves). 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,al=0.0110,0.679   # joint fit from analysis/staleness_fit.py
def simulate(b2,eps,g0=1e-5,T=60000):
    g,v=g0,g0*g0; out=[]
    for t in range(T+1):
        s=g/(np.sqrt(v)+eps); out.append(g); g=g*np.exp(-c*s**al); v=v+(1-b2)*(g*g-v)
    return np.array(out)
def T11(b2,eps):
    g=simulate(b2,eps); lng=np.log(1e-5/g); return int(np.argmax(lng>=11)) if (lng>=11).any() else np.nan
b2s=[0.9,0.95,0.98,0.985,0.99,0.995,0.999]
model={eps:np.array([T11(b,eps) for b in b2s]) for eps in (1e-8,1e-16)}
C={'1e-08':'#2a78d6','1e-16':'#eb6834'}; LAB={'1e-08':'eps 1e-8','1e-16':'eps 1e-16'}
fig,ax=plt.subplots(figsize=(6.2,4.4),dpi=150)
anchor=None
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]+', measured (seed mean)',zorder=2)
    if eps=='1e-16': anchor=mean[int(np.argmin([abs(x-0.02) for x in xs]))]/model[1e-16][b2s.index(0.98)]
for eps in (1e-16,1e-8):
    ax.plot([1-b for b in b2s],model[eps]*anchor,ls=':',lw=2,color=C['1e-08' if eps==1e-8 else '1e-16'],label=f'{LAB["1e-08" if eps==1e-8 else "1e-16"]}, ODE model',zorder=1)
ax.plot([],[],marker='^',mfc='none',mec='#555',ls='none',label='censored at 20000')
ax.set_xscale('log'); ax.set_yscale('log'); ax.invert_xaxis(); ax.set_xlim(0.12,0.0008)
ax.set_xticks([0.1,0.05,0.02,0.01,0.005,0.001]); ax.xaxis.set_minor_locator(matplotlib.ticker.NullLocator()); ax.set_xticklabels(['0.9','0.95','0.98','0.99','0.995','0.999'])
ax.set_xlabel('Adam beta2  (log scale of 1 - beta2)'); ax.set_ylabel('steps to test acc >= 0.99')
ax.set_title('N = 88, p = 97: measured vs ODE model (anchored at beta2 0.98, eps 1e-16)',fontsize=9)
ax.legend(fontsize=8,frameon=False,loc='upper left'); ax.grid(alpha=0.25,lw=0.6); [ax.spines[s].set_visible(False) for s in ('top','right')]
fig.tight_layout(); fig.savefig(f'{D}/fig_beta2_model.png')
print('anchor',anchor); print('b2, model eps1e-16, model eps1e-8 (scaled):'); [print(b,round(model[1e-16][i]*anchor),round(model[1e-8][i]*anchor)) for i,b in enumerate(b2s)]