#!/usr/bin/env python3
"""The scale sweep: retrieval vs YaRN scale, and the two-factor grid.
Every number read from the summary.csv files this session produced."""
import csv, os, glob
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

S={"surface":"#fcfcfb","ink":"#0b0b0b","ink2":"#52514e","grid":"#e4e4e1"}
def style(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 read(path,model):
    if not os.path.exists(path): return None
    r=list(csv.reader(open(path)))
    if len(r)<2: return None
    hdr=r[0]
    for row in r[1:]:
        if row and row[0]=="niah_single" and model in hdr:
            try: return float(row[hdr.index(model)])
            except Exception: return None
    return None

# (scale, prompt) -> dir
SRC={(4,1024):"eshort_cache16k_len1024",(8,1024):"eshort_cache32k_len1024",
     (4,4096):"eshort_cache16k_len4096",(8,4096):"eshort_cache32k_len4096",
     (16,4096):"e64k_len4096",(32,4096):"e128k_len4096",
     (8,8192):"e32k_len8192",(8,16360):"e32k_len16360",(8,32744):"e32k_len32744",
     (16,16360):"e64k_len16360",(16,32744):"e64k_len32744",(16,65512):"e64k_len65512",
     (32,32744):"e128k_len32744",(32,65512):"e128k_len65512"}
PRIOR={(4,8192):84.0,(4,16360):40.0}          # measured on the prior box, same haystack
base={}
for (sc,pl),d in SRC.items():
    v=read(f"{d}/summary.csv","base680")
    if v is not None: base[(sc,pl)]=v
base.update(PRIOR)

fig,(axA,axB,axC)=plt.subplots(1,3,figsize=(19.5,5.6),facecolor=S["surface"])
PL=[1024,4096,8192,16360,32744,65512]
cols=plt.cm.viridis([i/(len(PL)-1) for i in range(len(PL))])
for pl,c in zip(PL,cols):
    pts=sorted((sc,v) for (sc,p),v in base.items() if p==pl)
    if len(pts)<2: continue
    axA.plot([p[0] for p in pts],[p[1] for p in pts],"o-",color=c,lw=2,ms=6,
             label=f"{pl:,} tok")
style(axA); axA.set_xscale("log",base=2); axA.set_xticks([4,8,16,32])
axA.set_xticklabels(["x4","x8","x16","x32"]); axA.set_ylim(0,100)
axA.set_xlabel("YARN_SCALE  (ORIG_MAXPOS held at 4096)",color=S["ink2"],fontsize=10)
axA.set_ylabel("niah_single accuracy (%)",color=S["ink2"],fontsize=10)
axA.set_title("Same untrained weights. Only the RoPE scale changes.",
              color=S["ink"],fontsize=12,loc="left",pad=10)
axA.legend(frameon=False,fontsize=8.5,labelcolor=S["ink2"],loc="lower right")

for pl,c in zip(PL,cols):
    pts=sorted((pl/sc,v) for (sc,p),v in base.items() if p==pl)
    if not pts: continue
    axB.plot([p[0] for p in pts],[p[1] for p in pts],"o-",color=c,lw=2,ms=6,label=f"{pl:,} tok")
style(axB); axB.set_xscale("log",base=2); axB.set_ylim(0,100)
axB.set_xlabel("virtual position of last token  (real / scale)",color=S["ink2"],fontsize=10)
axB.set_ylabel("niah_single accuracy (%)",color=S["ink2"],fontsize=10)
axB.axvspan(2000,4200,color="#f3eceb",zorder=0)
axB.text(2900,6,"pretrain p90\n~2,721 tok",ha="center",fontsize=7.6,color=S["ink2"],style="italic")
axB.set_title("Collapses onto virtual position; spread = distractor load",
              color=S["ink"],fontsize=12,loc="left",pad=10)
axB.legend(frameon=False,fontsize=8.5,labelcolor=S["ink2"],loc="lower left")
fig.text(.005,-.02,"680M base680, ZERO long-context training. niah_single N=100, greedy, "
         "identical haystack (199,754,694 B) across every point. Wilson 95% CI is +/-9-13pp at N=100.",
         fontsize=7.8,color=S["ink2"])

# ---- panel C: the 1.06B control -- config helps, but training still wins at 64k
PLC=[16360,32744,65512]
own=[70,50,31]; corr=[78,73,65]; cliff_own=[54,31,18]; cliff_corr=[66,70,65]
m680=[92,79,46]
x=range(len(PLC))
axC.plot(x,own,"o--",color="#b0b0aa",lw=1.6,ms=6,label="1.06B, its own config")
axC.plot(x,corr,"o-",color="#b3243b",lw=2.4,ms=7,label="1.06B, corrected config")
axC.plot(x,cliff_corr,"s--",color="#eb6834",lw=1.6,ms=6,label='1.06B "collapsed", corrected')
axC.plot(x,m680,"^-",color="#12508f",lw=2.4,ms=8,label="680M base, no long-ctx training")
style(axC); axC.set_xticks(list(x)); axC.set_xticklabels(["16k","32k","64k"])
axC.set_ylim(0,100); axC.set_xlabel("prompt length",color=S["ink2"],fontsize=10)
axC.set_ylabel("niah_single accuracy (%)",color=S["ink2"],fontsize=10)
axC.set_title("Control: training wins where it trained (64k)",color=S["ink"],fontsize=12,loc="left",pad=10)
axC.annotate("1.06B trained to 64k\nbeats the 680M here\nby 19pp",xy=(2,65),xytext=(1.05,30),
             fontsize=8.2,color=S["ink2"],arrowprops=dict(arrowstyle="->",color=S["ink2"],lw=1))
axC.legend(frameon=False,fontsize=8,labelcolor=S["ink2"],loc="upper right")
fig.tight_layout(); fig.savefig("plots_p1/scale_sweep.png",dpi=155,bbox_inches="tight")
print("  wrote plots_p1/scale_sweep.png")
print(f"  points plotted: {len(base)}")