#!/usr/bin/env python3
"""Transposed view: tokens on x, one coloured line per YARN_SCALE."""
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
S={"surface":"#fcfcfb","ink":"#0b0b0b","ink2":"#52514e","grid":"#e4e4e1"}
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)

BASE={4:{1024:99,4096:89,8192:84,16360:40},
      8:{1024:99,4096:91,8192:88,16360:67,32744:34},
      16:{4096:92,16360:92,32744:54,65512:32},
      32:{4096:91,32744:79,65512:46},
      64:{4096:87,32744:79,65512:75}}
P1  ={4:{8192:77,16360:42},
      8:{1024:99,4096:88,8192:83,16360:67,32744:36},
      16:{4096:93,16360:90,32744:56,65512:36},
      32:{4096:96,32744:74,65512:52},
      64:{4096:95,32744:90,65512:82}}
SC=[4,8,16,32,64]
COL={4:"#b3243b",8:"#eb6834",16:"#d8b52a",32:"#3f9a54",64:"#12508f"}
PL=[1024,4096,8192,16360,32744,65512]

fig,(a,b)=plt.subplots(1,2,figsize=(16.6,6.4),facecolor=S["surface"],sharey=True)
for ax,D,ttl in ((a,BASE,"base680  (NO long-context training)"),
                 (b,P1,"p1final  (1,200 steps at 16k)")):
    for s in SC:
        pts=sorted(D[s].items())
        if not pts: continue
        ax.plot([p for p,_ in pts],[v for _,v in pts],"o-",color=COL[s],lw=2.4,ms=8,
                label=f"x{s}",zorder=3)
        ax.plot([p for p,_ in pts],[v for _,v in pts],"o",color="white",ms=3.4,zorder=4)
    sty(ax); ax.set_xscale("log",base=2)
    ax.set_xticks(PL); ax.set_xticklabels(["1k","4k","8k","16k","32k","64k"])
    ax.set_ylim(0,100); ax.set_xlim(900,78000)
    ax.set_xlabel("prompt length (tokens)",color=S["ink2"],fontsize=10.5)
    ax.set_title(ttl,color=S["ink"],fontsize=12,loc="left",pad=10)
a.set_ylabel("niah_single accuracy (%)",color=S["ink2"],fontsize=10.5)
a.legend(frameon=False,fontsize=9.5,labelcolor=S["ink2"],loc="lower left",
         title="YARN_SCALE",title_fontsize=9.5,ncol=2)
a.annotate("every line = same weights,\nonly the RoPE scale differs",xy=(16360,40),xytext=(2100,26),
           fontsize=9,color=S["ink2"],
           arrowprops=dict(arrowstyle="->",color=S["ink2"],lw=1.2,connectionstyle="arc3,rad=0.2"))
b.annotate("x64 holds up where\nx4 and x8 fall apart",xy=(65512,82),xytext=(20000,20),
           fontsize=9,color="#12508f",ha="center",
           arrowprops=dict(arrowstyle="->",color="#12508f",lw=1.2,
                           connectionstyle="arc3,rad=-0.25"))
fig.text(.005,.005,"Markers are MEASURED points only; lines connect them. Coverage is ragged -- 1,024 and 8,192 were "
         "only ever run at x4/x8, so their curves stop early and their true ceiling at high scale is untested. "
         "N=100 per point, identical haystack, Wilson 95% CI +/-9-13pp.",fontsize=8,color=S["ink2"])
fig.tight_layout(); fig.savefig("results_2026-09-16/overlay_by_scale.png",dpi=150,bbox_inches="tight")
print("  wrote results_2026-09-16/overlay_by_scale.png")