"""Figure: T_grok(N) for optimizer variants on the p=97 clamp. Left: beta2 = 0.98 with eps 1e-8, eps 1e-16, eps 0, Adam-atan2 (a=b=1).
Right: beta2 = 0.98 / 0.99 / 0.999, eps 1e-8 (solid) vs eps 1e-16 (dashed), plus Adam-atan2 at beta2 = 0.999.
Data: /work/grokking/generality/logs/p97_* and /work/grokking/temp-clamp/dose_response.csv (beta2 0.98 reference). Output: fig_optimizers.png, optimizer_dose_response.csv"""
import csv, glob, re, os, collections, matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
BLUE,ORANGE,AQUA,YELLOW='#2a78d6','#eb6834','#1baf7a','#eda100'; INK='#0b0b0b'; INK2='#52514e'; GRID='#e6e5e2'
rows=[]  # variant,N,seed,T,censored
for f in glob.glob('/work/grokking/generality/logs/p97_total*_seed*.csv'):
    m=re.match(r'p97_total(\d+)_lr0\.001_T1_b2([\d.]+)_(adamw|atan2)(?:_eps([\d.e-]+))?(?:_a([\d.]+)_b([\d.]+))?_seed(\d)\.csv',os.path.basename(f))
    if not m: continue
    N,b2,opt,eps,a,seed=int(m[1]),m[2],m[3],m[4] or '1e-8',m[5],int(m[7])
    if opt=='atan2' and a!='1': continue
    var=f'b2={b2} '+(f'eps={eps}' if opt=='adamw' else 'atan2')
    r=list(csv.DictReader(open(f))); T=next((int(x['step']) for x in r if float(x['test_acc'])>=0.99),None)
    rows.append((var,N,seed,T if T else int(r[-1]['step']),T is None))
for x in csv.DictReader(open('/work/grokking/temp-clamp/dose_response.csv')):
    T=x['T_grok']; cens=not T.isdigit(); rows.append((f"b2=0.98 eps={x['eps']}",int(x['N']),int(x['seed']),int(T) if not cens else int(x['last_step']),cens))
with open('optimizer_dose_response.csv','w',newline='') as fh:
    w=csv.writer(fh); w.writerow(['variant','N','seed','step_or_last','censored']); [w.writerow(r) for r in sorted(rows)]
def series(ax,var,col,ls='-',label=None,dx=0):
    byN=collections.defaultdict(list)
    for r in rows:
        if r[0]==var: byN[r[1]].append(r)
    xs,ys=[],[]
    for N in sorted(byN):
        ok=[r[3] for r in byN[N] if not r[4]]; ce=[r[3] for r in byN[N] if r[4]]
        ax.scatter([N+dx]*len(ok),ok,s=18,color=col,edgecolor='#fcfcfb',linewidth=1,zorder=3)
        if ce: ax.scatter([N+dx]*len(ce),ce,s=28,facecolor='#fcfcfb',edgecolor=col,linewidth=1.5,marker='^',zorder=3)
        if ok and not ce: xs.append(N); ys.append(sum(ok)/len(ok))
    ax.plot(xs,ys,color=col,lw=2,ls=ls,label=label or var,zorder=2)
plt.rcParams.update({'font.size':9,'axes.edgecolor':INK2,'axes.labelcolor':INK,'xtick.color':INK2,'ytick.color':INK2,'text.color':INK})
fig,(a1,a2)=plt.subplots(1,2,figsize=(10.5,4),sharey=True,gridspec_kw=dict(wspace=0.08))
series(a1,'b2=0.98 eps=1e-8',BLUE,label='AdamW, eps = 1e-8 (default)',dx=-0.6); series(a1,'b2=0.98 eps=1e-16',ORANGE,label='AdamW, eps = 1e-16',dx=-0.2)
series(a1,'b2=0.98 eps=0',AQUA,label='AdamW, eps = 0',dx=0.2); series(a1,'b2=0.98 atan2',YELLOW,label='Adam-atan2 (no eps)',dx=0.6)
a1.set_title('beta2 = 0.98: remove eps and the knee goes',loc='left',fontsize=10,fontweight='bold')
series(a2,'b2=0.98 eps=1e-8',BLUE,label='beta2 = 0.98, eps 1e-8',dx=-0.4); series(a2,'b2=0.98 eps=1e-16',BLUE,ls=(0,(4,2)),label='beta2 = 0.98, eps 1e-16',dx=-0.4)
series(a2,'b2=0.99 eps=1e-8',ORANGE,label='beta2 = 0.99, eps 1e-8'); series(a2,'b2=0.99 eps=1e-16',ORANGE,ls=(0,(4,2)),label='beta2 = 0.99, eps 1e-16')
series(a2,'b2=0.999 eps=1e-8',AQUA,label='beta2 = 0.999, eps 1e-8',dx=0.4); series(a2,'b2=0.999 eps=1e-16',AQUA,ls=(0,(4,2)),label='beta2 = 0.999, eps 1e-16',dx=0.4)
series(a2,'b2=0.999 atan2',YELLOW,label='beta2 = 0.999, Adam-atan2',dx=0.8)
a2.set_title('larger beta2: a slowdown that eps does not touch',loc='left',fontsize=10,fontweight='bold')
for ax in (a1,a2):
    ax.set_yscale('log'); ax.set_ylim(600,30000); ax.set_xlim(44,126); ax.set_xlabel('clamped total weight norm N'); ax.grid(True,color=GRID,lw=1); ax.set_axisbelow(True); ax.spines[['top','right']].set_visible(False)
    ax.legend(frameon=False,fontsize=7.5,loc='upper left')
a1.set_ylabel('steps to test accuracy 0.99'); a2.text(0.98,0.03,'open triangles: not reached within budget',transform=a2.transAxes,ha='right',fontsize=7.5,color=INK2)
fig.suptitle('p = 97 MLP under a total-norm clamp, lr 1e-3, wd 0: which optimizer variants show the knee (dots: seeds; lines: means)',x=0.01,ha='left',fontsize=10)
fig.savefig('fig_optimizers.png',dpi=150,bbox_inches='tight',facecolor='#fcfcfb'); print(len(rows),'rows')