"""T_grok vs clamped norm N for each width variant (p=97), eps 1e-8 vs 1e-16. Reads width_dose_response.csv
(run width_summarize.py first) and the free-run logs for W_c / init norm. Writes fig_width_knee.png."""
import csv, os, collections
import matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
os.chdir(os.path.dirname(os.path.abspath(__file__)))
BLUE,ORANGE,AQUA,YELLOW='#2a78d6','#eb6834','#1baf7a','#eda100'
VAR=[('ref',256,64),('A',128,64),('B',512,64),('C',256,32),('D',256,128)]
rows=list(csv.DictReader(open('width_dose_response.csv')))
def free(v,h,e):
    s=("" if h==256 else f"_h{h}")+("" if e==64 else f"_e{e}")
    fn='../temp-clamp/logs/none0_lr0.001_T1_seed0_wd1.csv' if v=='ref' else f'logs/p97_none0_lr0.001_T1_b20.98_adamw_wd1{s}_seed0.csv'
    R=list(csv.DictReader(open(fn))); g=lambda th: next(float(r['norm_total']) for r in R if float(r['test_acc'])>=th)
    return float(R[0]['norm_total']),g(0.5),g(0.99)
fig,axs=plt.subplots(1,5,figsize=(19,4.2),sharey=True)
for ax,(v,h,e) in zip(axs,VAR):
    init,w5,w99=free(v,h,e)
    ax.axvspan(w99,w5,color=AQUA,alpha=0.12,lw=0)
    ax.axvline(w99,color=AQUA,lw=1.5,ls='-',label='free-run W_c (test 0.99 .. 0.5)'); ax.axvline(w5,color=AQUA,lw=1.5,ls='-')
    ax.axvline(init,color=YELLOW,lw=2,ls='--',label='default init norm')
    for eps,c,lab in (('1e-8',BLUE,'eps 1e-8'),('1e-16',ORANGE,'eps 1e-16')):
        R=[r for r in rows if r['variant']==v and r['eps']==eps and float(r['N'])<=125]  # ref N=150: both eps censored, omitted
        byN=collections.defaultdict(list)
        for r in R: byN[float(r['N'])].append((int(r['step_or_last']),int(r['censored'])))
        Ns=sorted(byN); means=[sum(t for t,_ in byN[N])/len(byN[N]) for N in Ns]
        ax.plot(Ns,means,color=c,lw=2,label=lab,zorder=3)
        for N in Ns:
            for t,cz in byN[N]:
                ax.plot([N],[t],marker='^' if cz else 'o',ms=8,mfc='white' if cz else c,mec=c,mew=2 if cz else 1.5,ls='none',zorder=4)
    ax.set_yscale('log'); ax.set_xlabel('clamped total norm N')
    ax.set_title(f"{'reference' if v=='ref' else 'variant '+v}: d_hid {h}, d_emb {e}",fontsize=10)
    ax.grid(alpha=0.25,lw=0.6); [ax.spines[s].set_visible(False) for s in ('top','right')]
axs[0].set_ylabel('steps to test acc >= 0.99 (log)')
h,l=axs[1].get_legend_handles_labels()
h.append(plt.Line2D([],[],marker='^',mfc='white',mec='gray',ls='none',ms=8)); l.append('not grokked by cap (lower bound)')
fig.legend(h,l,loc='lower center',ncol=5,frameon=False,fontsize=9)
fig.suptitle('(a+b) mod 97, total norm clamped to N: knee in steps-to-grok with eps 1e-8 (absent with 1e-16), by width',fontsize=11)
fig.tight_layout(rect=(0,0.08,1,0.95)); fig.savefig('fig_width_knee.png',dpi=150); print('saved fig_width_knee.png')