"""Figure: clamped 1-layer transformer (p=113), total norm held at 50 or 60, AdamW eps 1e-8 vs 1e-16. Data: /work/grokking/baseline/var_clamp/N*/log.csv"""
import csv, matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
BLUE,ORANGE='#2a78d6','#eb6834'; INK='#0b0b0b'; INK2='#52514e'; GRID='#e6e5e2'
def load(d):
    r=list(csv.DictReader(open(f'/work/grokking/baseline/var_clamp/{d}/log.csv'))); return [int(x['step']) for x in r],[float(x['test_acc']) for x in r]
plt.rcParams.update({'font.size':9,'axes.edgecolor':INK2,'axes.labelcolor':INK,'xtick.color':INK2,'ytick.color':INK2,'text.color':INK})
fig,axs=plt.subplots(1,2,figsize=(10,3.6),sharey=True,gridspec_kw=dict(wspace=0.08))
for ax,N in zip(axs,[50,60]):
    for eps,col in [('1e-8',BLUE),('1e-16',ORANGE)]:
        s,a=load(f'N{N}_eps{eps}'); ax.plot(s,a,color=col,lw=2,label=f'eps = {eps}')
        t99=next((x for x,y in zip(s,a) if y>=0.99),None)
        if t99: ax.scatter([t99],[0.99],s=40,color=col,edgecolor='#fcfcfb',linewidth=1.5,zorder=4); ax.annotate(f'{t99:,}',(t99,0.99),xytext=(0,-14 if eps=='1e-8' else 8),textcoords='offset points',ha='center',fontsize=8,color=INK2)
    ax.set_title(f'norm held at N = {N}',loc='left',fontsize=10,fontweight='bold'); ax.set_xlabel('training step'); ax.set_ylim(0,1.05)
    ax.grid(True,color=GRID,lw=1); ax.set_axisbelow(True); ax.spines[['top','right']].set_visible(False)
axs[0].set_ylabel('test accuracy'); axs[0].legend(frameon=False,loc='lower right',fontsize=8)
fig.suptitle('Clamped one-layer transformer, (a+b) mod 113: step at which test accuracy first reaches 0.99 is marked (one seed per run)',x=0.01,ha='left',fontsize=10)
fig.savefig('fig_transformer_clamp.png',dpi=150,bbox_inches='tight',facecolor='#fcfcfb'); print('ok')