#!/usr/bin/env python3
"""One prompt, three scales. Why the same weights score 40 or 92."""
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, Rectangle
S={"surface":"#fcfcfb","ink":"#0b0b0b","ink2":"#52514e","grid":"#e4e4e1"}
RED,AMB,BLU="#b3243b","#eb6834","#12508f"
fig,(ax,bx)=plt.subplots(2,1,figsize=(12.6,8.6),facecolor=S["surface"],
                         gridspec_kw={"height_ratios":[1.15,1],"hspace":.42})

# ---- top: the squeeze -------------------------------------------------------
ax.set_facecolor(S["surface"]); ax.axis("off")
ax.set_xlim(0,100); ax.set_ylim(0,100)
ax.text(0,96,"ONE prompt of 16,360 tokens. Same model file. Only YARN_SCALE changes.",
        fontsize=12.5,color=S["ink"],weight="bold",va="top")
# the familiar band
for y,(s,pos,sc,col) in zip((66,40,14),((4,4090,40,RED),(8,2045,67,AMB),(16,1022,92,BLU))):
    ax.add_patch(Rectangle((14,y-4),62,8,fc="#f0f0ee",ec=S["grid"]))
    ax.text(13,y,f"x{s}",fontsize=11,color=col,ha="right",va="center",weight="bold")
    # familiar / rare split at 4096 -> the bar spans 0..4096 perceived
    frac=pos/4096.0
    ax.add_patch(Rectangle((14,y-4),62*frac,8,fc=col,alpha=.30,ec="none"))
    ax.plot([14+62*frac],[y],"o",color=col,ms=11)
    if s == 4:
        ax.text(14+62*frac-2,y-7.0,f"last token perceived\nat position {pos:,}",
                fontsize=8.4,color=col,ha="right",va="top")
    else:
        ax.text(14+62*frac,y+7.0,f"last token perceived\nat position {pos:,}",
                fontsize=8.4,color=col,ha="center",va="bottom")
    ax.text(78,y,f"scores {sc}",fontsize=11,color="white",va="center",
            bbox=dict(fc=col,ec="none",pad=3.4))
ax.plot([76,76],[8,72],color=S["ink2"],lw=1.2,ls="--")
ax.text(78.5,80,"position 4,096\nedge of pretraining",fontsize=8.4,color=S["ink2"],ha="left")
ax.text(14,4,"position 0",fontsize=8.4,color=S["ink2"],ha="center")
ax.text(45,-1,"the 4,096-position range the model was pretrained on",
        fontsize=9,color=S["ink2"],ha="center",style="italic")

# ---- bottom: where the pretraining data actually was ------------------------
bx.set_facecolor(S["surface"])
for s in bx.spines.values(): s.set_color(S["grid"])
bx.grid(True,color=S["grid"],lw=.7,alpha=.9); bx.tick_params(colors=S["ink2"],labelsize=9)
import math
xs=list(range(0,4097,32))
# schematic doc-length survival curve: p90 = 2,721
dens=[100*math.exp(-(x/2721.0)**1.7) for x in xs]
bx.fill_between(xs,dens,color="#9dc3ea",alpha=.55)
bx.plot(xs,dens,color=BLU,lw=2)
for pos,sc,col,lab,ha,dx in ((4090,40,RED,"x4","right",-60),(2045,67,AMB,"x8","center",0),
                             (1022,92,BLU,"x16","center",0)):
    bx.axvline(pos,color=col,lw=2,ls="--")
    bx.text(pos+dx,95,lab,color=col,fontsize=10.5,ha=ha,weight="bold",va="top")
    bx.text(pos+dx,86,f"scores {sc}",color="white",fontsize=9,ha=ha,va="top",
            bbox=dict(fc=col,ec="none",pad=2.2))
bx.axvline(2721,color=S["ink2"],lw=1.2)
bx.text(2721,52,"  p90 document\n  length 2,721",fontsize=8.4,color=S["ink2"])
bx.set_xlim(0,4096); bx.set_ylim(0,100)
bx.set_xlabel("perceived position",color=S["ink2"],fontsize=10)
bx.set_ylabel("how often pretraining\nreached this position (%)",color=S["ink2"],fontsize=10)
bx.set_title("Why: past ~2,700 the model has almost no experience",
             color=S["ink"],fontsize=11.5,loc="left",pad=9)
fig.text(.005,.005,"Verified numerically: token 16,360 at x16 and token 1,022 unscaled differ by 1.4e-4 radians "
         "in the interpolated dims -- indistinguishable to the model. Document-length curve is schematic, anchored on the measured p90 of 2,721.",
         fontsize=7.8,color=S["ink2"])
fig.savefig("results_2026-09-16/why_scale_works.png",dpi=150,bbox_inches="tight")
print("  wrote results_2026-09-16/why_scale_works.png")