"""Figure: clamped dose-response T_grok(N) for p = 59, 97, 113 at AdamW eps 1e-8 vs 1e-16 (betas (0.9,0.98), lr 1e-3, wd 0).
Data: /work/grokking/generality/logs/p59_*,p113_* (clamp_gen.py) and /work/grokking/temp-clamp/dose_response.csv (p=97, run 4).
Output: fig_moduli.png and moduli_dose_response.csv"""
import csv, glob, re, os, collections, matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
BLUE,ORANGE='#2a78d6','#eb6834'; INK='#0b0b0b'; INK2='#52514e'; GRID='#e6e5e2'
rows=[]  # p,N,eps,seed,T,censored,last
for f in glob.glob('/work/grokking/generality/logs/p*_total*_b20.98_adamw*_seed*.csv'):
    m=re.match(r'p(\d+)_total(\d+)_lr0\.001_T1_b20\.98_adamw(?:_eps([\d.e-]+))?_seed(\d)\.csv',os.path.basename(f))
    if not m: continue
    p,N,eps,seed=int(m[1]),int(m[2]),m[3] or '1e-8',int(m[4])
    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((p,N,eps,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((97,int(x['N']),x['eps'],int(x['seed']),int(T) if not cens else int(x['last_step']),cens))
with open('moduli_dose_response.csv','w',newline='') as fh:
    w=csv.writer(fh); w.writerow(['p','N','eps','seed','step_or_last','censored']); [w.writerow(r) for r in sorted(rows)]
WC={59:64,97:88,113:95.5}  # free wd=1 run: total norm when test acc crosses 0.5
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,3,figsize=(11,3.9),sharey=True,gridspec_kw=dict(wspace=0.08))
for ax,p in zip(axs,[59,97,113]):
    for eps,col,dx in [('1e-8',BLUE,-0.35),('1e-16',ORANGE,0.35)]:
        pts=[r for r in rows if r[0]==p and r[2]==eps]
        byN=collections.defaultdict(list)
        for r in pts: byN[r[1]].append(r)
        Ns=sorted(byN); means=[]
        for N in Ns:
            ok=[r[4] for r in byN[N] if not r[5]]; ce=[r[4] for r in byN[N] if r[5]]
            ax.scatter([N+dx]*len(ok),ok,s=22,color=col,edgecolor='#fcfcfb',linewidth=1,zorder=3)
            if ce: ax.scatter([N+dx]*len(ce),ce,s=30,facecolor='#fcfcfb',edgecolor=col,linewidth=1.5,marker='^',zorder=3)
            means.append(sum(ok)/len(ok) if ok and not ce else None)
        xs=[N for N,m in zip(Ns,means) if m]; ys=[m for m in means if m]
        ax.plot(xs,ys,color=col,lw=2,label=f'eps = {eps}',zorder=2)
    ax.axvline(WC[p],color=INK2,lw=1,ls=(0,(3,3))); ax.text(WC[p],22000,' free run\n groks here',fontsize=7.5,color=INK2,va='top')
    ax.set_yscale('log'); ax.set_ylim(800,25000); ax.set_title(f'p = {p}',loc='left',fontsize=10,fontweight='bold'); 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)
axs[0].set_ylabel('steps to test accuracy 0.99'); axs[0].legend(frameon=False,loc='upper left',fontsize=8)
axs[2].text(0.98,0.04,'open triangles: not reached\nwithin the budget (plotted at last step)',transform=axs[2].transAxes,ha='right',fontsize=7.5,color=INK2)
fig.suptitle('The default-eps knee appears at every modulus and vanishes with eps = 1e-16 (MLP, total-norm clamp, AdamW betas (0.9, 0.98))',x=0.01,ha='left',fontsize=10)
fig.savefig('fig_moduli.png',dpi=150,bbox_inches='tight',facecolor='#fcfcfb'); print(len(rows),'runs')