#!/usr/bin/env python3
"""Phase-1 LR-schedule ablation. Every number read from p1ev_*/summary.csv."""
import csv,os,math
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
S={"surface":"#fcfcfb","ink":"#0b0b0b","ink2":"#52514e","grid":"#e4e4e1"}
C={"abl_const":"#b3243b","abl_cosine":"#eb6834","p1final":"#12508f","base680":"#9dc3ea"}
LAB={"abl_const":"const 8e-5 (Muon+AdamW)","abl_cosine":"cosine 8e-5->8e-6 (Muon+AdamW)",
     "p1final":"p1final  1.8e-4 WSD (AdamW)","base680":"base680  NO training"}
def sty(ax):
    ax.set_facecolor(S["surface"])
    for s in ax.spines.values(): s.set_color(S["grid"])
    ax.grid(True,color=S["grid"],lw=.7,alpha=.9); ax.tick_params(colors=S["ink2"],labelsize=9)
def wilson(k,n=100,z=1.96):
    p=k/n; d=1+z*z/n; c=(p+z*z/(2*n))/d
    h=z*math.sqrt(p*(1-p)/n+z*z/(4*n*n))/d
    return 100*max(0,c-h),100*min(1,c+h)
PL=[4096,8192,16360]; D={}
for pl in PL:
    r=list(csv.reader(open(f"p1ev_{pl}/summary.csv"))); h=r[0]
    row=[x for x in r[1:] if x and x[0]=="niah_single"][0]
    for k,v in zip(h[1:],row[1:]): D.setdefault(k,{})[pl]=float(v)

fig,(a,b,c)=plt.subplots(1,3,figsize=(17.6,5.6),facecolor=S["surface"])
X=range(len(PL))
for m in ("base680","p1final","abl_const","abl_cosine"):
    ys=[D[m][p] for p in PL]
    lo=[y-wilson(round(y))[0] for y in ys]; hi=[wilson(round(y))[1]-y for y in ys]
    a.errorbar(list(X),ys,yerr=[lo,hi],color=C[m],lw=2.4,ms=8,marker="o",capsize=3,label=LAB[m])
sty(a); a.set_xticks(list(X)); a.set_xticklabels(["4k","8k","16k"]); a.set_ylim(60,105)
a.set_xlabel("prompt length",color=S["ink2"],fontsize=10)
a.set_ylabel("niah_single (%)",color=S["ink2"],fontsize=10)
a.set_title("A. All four models, Wilson 95% CI",color=S["ink"],fontsize=11.5,loc="left",pad=9)
a.legend(frameon=False,fontsize=8,labelcolor=S["ink2"],loc="lower left")

d=[D["abl_const"][p]-D["abl_cosine"][p] for p in PL]
b.bar(list(X),d,color=["#b3243b" if x>0 else "#eb6834" for x in d],width=.5)
b.axhline(0,color=S["ink2"],lw=1)
b.axhspan(-11,11,color="#f0f0ee",zorder=0)
b.text(1,8.4,"Wilson 95% CI band at N=100",ha="center",fontsize=8.4,color=S["ink2"],style="italic")
sty(b); b.set_xticks(list(X)); b.set_xticklabels(["4k","8k","16k"]); b.set_ylim(-14,14)
b.set_ylabel("const - cosine (pp)",color=S["ink2"],fontsize=10)
b.set_title("B. THE ABLATION: schedule shape -- null",color=S["ink"],fontsize=11.5,loc="left",pad=9)
for i,v in enumerate(d): b.text(i,v+(0.7 if v>=0 else -1.6),f"{v:+.0f}",ha="center",fontsize=10,color=S["ink2"])

names=["base680","p1final","abl_const","abl_cosine"]
tot=[sum(D[m][p] for p in PL) for m in names]
c.bar(range(4),tot,color=[C[m] for m in names])
for i,v in enumerate(tot): c.text(i,v+3,f"{v:.0f}",ha="center",fontsize=10.5,color=S["ink2"],weight="bold")
sty(c); c.set_xticks(range(4))
c.set_xticklabels(["base680\n(none)","p1final\n1.8e-4 AdamW","const\n8e-5 Muon","cosine\n8e-5 Muon"],fontsize=8.4)
c.set_ylim(0,310); c.set_ylabel("sum over 4k/8k/16k (pp)",color=S["ink2"],fontsize=10)
c.set_title("C. Both new arms sit BELOW the untrained base",color=S["ink"],fontsize=11.5,loc="left",pad=9)
fig.text(.005,.005,"680M, 4k->16k, 1,200 steps, identical data (ablmix16, token-identical to p1mix) and identical haystack "
  "(199,754,694 B). N=100/point. CAVEAT: the two new arms differ from p1final in BOTH optimizer (Muon vs AdamW) and LR "
  "(8e-5 vs 1.8e-4) -- this run cannot separate them.",fontsize=7.8,color=S["ink2"])
fig.tight_layout(); fig.savefig("plots_abl/phase1_ablation.png",dpi=150,bbox_inches="tight")
print("  wrote plots_abl/phase1_ablation.png")