from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch

OUT = Path(__file__).resolve().parent / 'figures'
OUT.mkdir(parents=True, exist_ok=True)

plt.rcParams.update({
    'font.size': 9,
    'axes.titlesize': 10.5,
    'axes.labelsize': 9,
    'legend.fontsize': 7.6,
    'xtick.labelsize': 8,
    'ytick.labelsize': 8,
    'axes.spines.top': False,
    'axes.spines.right': False,
    'figure.dpi': 160,
})

COLORS = {
    'Pure backbone': '#244c86',
    'Fixed token MLP': '#d97706',
    'Dynamic token MLP': '#c2410c',
    'Fixed context': '#0f766e',
    'Dynamic surprisal': '#7c3aed',
    'Low-surprisal': '#64748b',
    'Position control': '#be185d',
}
ALIASES = {
    'Pure': 'Pure backbone',
    'Fixed context': 'Fixed context',
    'Dynamic MLP': 'Dynamic token MLP',
    'Low surprisal': 'Low-surprisal',
    'Position control': 'Position control',
    'Dynamic low-surprisal context': 'Low-surprisal',
    'Dynamic position control': 'Position control',
    'Dynamic surprisal context': 'Dynamic surprisal',
}

def c(name):
    return COLORS[ALIASES.get(name, name)]

def stylize(ax):
    ax.grid(True, linewidth=0.45, alpha=0.22)
    ax.spines['left'].set_alpha(0.35)
    ax.spines['bottom'].set_alpha(0.35)

steps = np.array([5000, 10000, 25000, 50000], dtype=float)
curves = {
    'Pure backbone': ([3.002453070,2.647377918,2.204987323,1.956480433], [2.991404684,2.634873385,2.197992739,1.951222360], [3.013501456,2.659882452,2.211981907,1.961738506]),
    'Fixed token MLP': ([2.958571643,2.630465133,2.196781493,1.956499330], [2.940964913,2.593920542,2.181474156,1.941896148], [2.976178373,2.667009724,2.212088830,1.971102512]),
    'Dynamic token MLP': ([2.959107563,2.617091001,2.202105529,1.958862610], [2.930136347,2.575824340,2.190297807,1.956367819], [2.988078779,2.658357661,2.213913251,1.961357401]),
    'Fixed context': ([2.953397018,2.627311605,2.195984719,1.956121514], [2.944991852,2.621072583,2.188938321,1.951553406], [2.961802184,2.633550627,2.203031117,1.960689622]),
    'Dynamic surprisal': ([2.968121545,2.620589070,2.197657573,1.958168823], [2.963152746,2.607984633,2.194235838,1.947922547], [2.973090344,2.633193507,2.201079308,1.968415099]),
    'Dynamic low-surprisal context': ([2.967516025,2.626597242,2.194625936,1.958025402], [2.962597977,2.619228091,2.190184656,1.952684739], [2.972434073,2.633966393,2.199067216,1.963366065]),
    'Dynamic position control': ([2.984059952,2.614024814,2.196949189,1.955770301], [2.791574734,2.344378814,2.159204142,1.948191764], [3.176545169,2.883670813,2.234694236,1.963348837]),
}

# Figure 1: learning curves
fig, ax = plt.subplots(figsize=(8.2, 4.8))
markers = {
    'Pure backbone': 'o', 'Fixed token MLP': 's', 'Dynamic token MLP': '^',
    'Fixed context': 'D', 'Dynamic surprisal': 'P',
    'Dynamic low-surprisal context': 'X', 'Dynamic position control': 'v'
}
for label, (mean, low, high) in curves.items():
    mean, low, high = map(np.asarray, (mean, low, high))
    color = c(label)
    ax.fill_between(steps/1000, low, high, color=color, alpha=0.12, linewidth=0)
    ax.plot(steps/1000, mean, marker=markers[label], color=color, linewidth=2.0,
            markersize=5.4, label=label.replace('Dynamic ', 'Dyn. ').replace('context', 'ctx.'))
ax.set_xlabel('Adaptation step (thousands)')
ax.set_ylabel('Complete-development BPB')
ax.set_xticks([5,10,25,50])
ax.set_title('Learning curves over three paired warm starts', pad=10)
stylize(ax)
leg = ax.legend(ncol=2, frameon=True, facecolor='white', edgecolor='#dddddd', loc='upper right')
leg.get_frame().set_alpha(0.95)
ax.text(0.01, 0.98, 'Lower is better; shaded regions are 95% paired-seed intervals.',
        transform=ax.transAxes, va='top', ha='left', fontsize=7.7,
        bbox=dict(boxstyle='round,pad=0.25', facecolor='#f8fafc', edgecolor='#e2e8f0'))
fig.tight_layout()
for ext in ['pdf','png']:
    fig.savefig(OUT / f'learning_curves.{ext}', bbox_inches='tight')
plt.close(fig)

# Figure 2: endpoint and paired comparisons
arms = ['Position control', 'Fixed context', 'Pure backbone', 'Fixed token MLP', 'Low-surprisal', 'Dynamic surprisal', 'Dynamic token MLP']
means = np.array([1.955770301, 1.956121514, 1.956480433, 1.956499330, 1.958025402, 1.958168823, 1.958862610])
lows = np.array([1.948191764, 1.951553406, 1.951222360, 1.941896148, 1.952684739, 1.947922547, 1.956367819])
highs = np.array([1.963348837, 1.960689622, 1.961738506, 1.971102512, 1.963366065, 1.968415099, 1.961357401])
comparators = ['Pure', 'Fixed context', 'Dynamic MLP', 'Low surprisal', 'Position control']
delta_mean = np.array([0.0016883901, 0.0020473098, -0.0006937866, 0.0001434216, 0.0023985227])
delta_low = np.array([-0.0051039741, -0.0037701613, -0.0085394507, -0.0048416383, -0.0067327012])
delta_high = np.array([0.0084807543, 0.0078647809, 0.0071518774, 0.0051284815, 0.0115297466])
fig, axes = plt.subplots(1, 2, figsize=(8.5, 4.6), gridspec_kw={'width_ratios':[1.08,1]})
for i, arm in enumerate(arms):
    axes[0].errorbar(means[i], i, xerr=np.array([[means[i]-lows[i]], [highs[i]-means[i]]]),
                     fmt='o', color=c(arm), ecolor=c(arm), capsize=3, elinewidth=1.6, markersize=5.2)
axes[0].set_yticks(range(len(arms)), arms)
axes[0].invert_yaxis()
axes[0].set_xlabel('BPB at 50k')
axes[0].set_title('Endpoint ranking')
stylize(axes[0])
axes[1].axvline(0.0, color='#334155', linewidth=1.0, linestyle='--')
for i, comp in enumerate(comparators):
    color = c(comp)
    axes[1].errorbar(delta_mean[i], i,
                     xerr=np.array([[delta_mean[i]-delta_low[i]], [delta_high[i]-delta_mean[i]]]),
                     fmt='o', color=color, ecolor=color, capsize=3, elinewidth=1.6, markersize=5.2)
axes[1].set_yticks(range(len(comparators)), comparators)
axes[1].invert_yaxis()
axes[1].set_xlabel('Candidate minus comparator BPB')
axes[1].set_title('Preregistered paired tests')
stylize(axes[1])
fig.tight_layout(w_pad=2.0)
for ext in ['pdf','png']:
    fig.savefig(OUT / ('endpoint_and_paired.' + ext), bbox_inches='tight')
plt.close(fig)

# Figure 3: mechanism diagnostics
seed_labels = ['11111', '22222', '33333']
rate_mean = np.array([0.03149470465435531, 0.03149201388326202, 0.031489774600649084])
rate_std = np.array([1.1044665430586648e-05, 8.02547709491516e-06, 1.1674269242521738e-05])
rate_cv = np.array([0.0003506832514163397, 0.0002548416600051321, 0.00037073206749092135])
pearson = np.array([0.041068441387792094, -0.26254092638892285, 0.16457092956772099])
fig, axes = plt.subplots(1, 3, figsize=(8.6, 3.5))
x = np.arange(3)
pal = ['#2563eb','#f59e0b','#8b5cf6']
axes[0].bar(x, rate_mean, color=pal, alpha=0.9)
axes[0].errorbar(x, rate_mean, yerr=rate_std, fmt='none', ecolor='#334155', capsize=3, elinewidth=1.0)
axes[0].axhline(0.031496062992125984, linestyle='--', color='#64748b', linewidth=1.0)
axes[0].set_xticks(x, seed_labels)
axes[0].set_title('Predicted chunk rate')
axes[0].set_ylabel('Mean rate')
axes[0].ticklabel_format(axis='y', style='plain', useOffset=False)
stylize(axes[0])
axes[1].bar(x, rate_cv, color=pal, alpha=0.9)
axes[1].axhline(0.01, linestyle='--', color='#64748b', linewidth=1.0)
axes[1].set_xticks(x, seed_labels)
axes[1].set_title('Rate coefficient of variation')
axes[1].set_ylabel('CV')
stylize(axes[1])
axes[2].bar(x, pearson, color=pal, alpha=0.9)
axes[2].axhline(0.10, linestyle='--', color='#64748b', linewidth=1.0)
axes[2].axhline(0.0, color='#334155', linewidth=0.8)
axes[2].set_xticks(x, seed_labels)
axes[2].set_title('Previous surprise → next rate')
axes[2].set_ylabel('Pearson r')
stylize(axes[2])
for ax in axes:
    ax.set_xlabel('Paired seed')
fig.tight_layout(w_pad=2.0)
for ext in ['pdf','png']:
    fig.savefig(OUT / ('mechanism_diagnostics.' + ext), bbox_inches='tight')
plt.close(fig)

# Figure 4: latency-quality (legend used instead of direct labels to avoid overlap)
latency = {
    'Pure backbone': (12.325685, 1.956480433),
    'Fixed token MLP': (27.534581, 1.956499330),
    'Dynamic token MLP': (31.590368, 1.958862610),
    'Fixed context': (26.735189, 1.956121514),
    'Dynamic surprisal': (30.431712, 1.958168823),
    'Low-surprisal': (31.353984, 1.958025402),
    'Position control': (31.854304, 1.955770301),
}
fig, ax = plt.subplots(figsize=(8.1, 4.6))
markers2 = {
    'Pure backbone': 'o', 'Fixed token MLP': 's', 'Dynamic token MLP': '^',
    'Fixed context': 'D', 'Dynamic surprisal': 'P', 'Low-surprisal': 'X', 'Position control': 'v'
}
for label, (xv, yv) in latency.items():
    ax.scatter([xv], [yv], s=60, color=c(label), marker=markers2[label],
               edgecolors='white', linewidths=0.7, zorder=3,
               label=label.replace('Dynamic ','Dyn. ').replace('context','ctx.'))
ax.set_xlabel('H100 batch-64 p50 latency (ms)')
ax.set_ylabel('Complete-development BPB at 50k')
ax.set_title('Quality and latency do not improve together')
stylize(ax)
ax.set_xlim(10.5, 33.5)
ax.text(0.01, 0.98,
        'Lower and farther left is better. Labels are moved to the legend\n'
        'to avoid overprinting among tightly clustered points.',
        transform=ax.transAxes, va='top', ha='left', fontsize=7.4,
        bbox=dict(boxstyle='round,pad=0.25', facecolor='#f8fafc', edgecolor='#e2e8f0'))
leg = ax.legend(loc='center left', bbox_to_anchor=(1.02, 0.5), frameon=True, facecolor='white', edgecolor='#dddddd')
leg.get_frame().set_alpha(0.96)
fig.tight_layout()
for ext in ['pdf','png']:
    fig.savefig(OUT / ('latency_quality.' + ext), bbox_inches='tight')
plt.close(fig)

# Figure 5: system overview
fig, ax = plt.subplots(figsize=(8.4, 4.9))
ax.set_xlim(0, 14)
ax.set_ylim(0, 8)
ax.axis('off')
ax.text(7, 7.62, 'Causal token-routing system evaluated in this study', ha='center', va='center', fontsize=12.5, fontweight='bold', color='#1e293b')

def box(x, y, w, h, text, fc, ec='#cbd5e1', fontsize=8.3, ls='-', lw=1.1):
    p = FancyBboxPatch((x, y), w, h, boxstyle='round,pad=0.06,rounding_size=0.10',
                       facecolor=fc, edgecolor=ec, linewidth=lw, linestyle=ls)
    ax.add_patch(p)
    ax.text(x+w/2, y+h/2, text, ha='center', va='center', fontsize=fontsize, color='#0f172a')
    return p

def arrow(x1,y1,x2,y2,color='#475569',ls='-',lw=1.2):
    ax.add_patch(FancyArrowPatch((x1,y1),(x2,y2),arrowstyle='-|>',mutation_scale=11,linewidth=lw,linestyle=ls,color=color))

box(0.35, 3.0, 1.7, 1.2, 'Byte sequence\nlength 128', '#ecfeff', ec='#67e8f9')
box(2.45, 2.75, 2.7, 1.8, 'Backbone blocks 1-6\nAttention / Mamba-2\nalternating', '#eff6ff', ec='#93c5fd')
box(5.55, 5.15, 2.45, 1.2, 'Tied midpoint head\ncausal surprisal $S_t$', '#fef3c7', ec='#fbbf24')
box(5.55, 2.65, 2.45, 1.5, 'Budget network\nprevious chunk only\n705 parameters', '#fae8ff', ec='#e879f9')
box(8.4, 3.1, 2.6, 1.65, 'Causal quota\nrank positions +\ninteger chunk budget', '#f3e8ff', ec='#c084fc')
box(8.4, 5.58, 2.6, 1.05, 'Ranking controls\nhigh / low / position', '#f8fafc', ec='#94a3b8', ls='--')
box(11.25, 3.1, 2.45, 1.65, 'Sparse actuator\nselected-query context\nor token MLP', '#fff7ed', ec='#fdba74')
box(3.7, 0.52, 7.1, 1.05, 'Backbone blocks 7-12 → tied output head → next-byte loss', '#eff6ff', ec='#93c5fd')
arrow(2.05,3.6,2.45,3.6,color='#0891b2')
arrow(5.15,3.65,5.55,5.7,color='#2563eb')
arrow(6.78,5.15,6.78,4.2,color='#ca8a04')
arrow(8.0,3.45,8.4,3.9,color='#a21caf')
arrow(8.0,5.75,8.4,4.42,color='#a21caf')
arrow(11.0,3.9,11.25,3.9,color='#ea580c')
arrow(12.48,3.1,9.85,1.57,color='#ea580c')
arrow(5.12,3.08,4.95,1.57,color='#64748b')
ax.text(6.8, 4.63, 'detached observed\nsurprisal', ha='center', va='center', fontsize=7.4, color='#92400e')
ax.text(9.72, 2.58, 'hard selected mask', ha='center', va='center', fontsize=7.4, color='#6b21a8')
ax.text(7, 0.08, 'Ranking, budget prediction, and sparse refinement are explicit and separately controlled.',
        ha='center', va='bottom', fontsize=8.3, color='#334155')
fig.tight_layout()
for ext in ['pdf','png']:
    fig.savefig(OUT / ('system_overview.' + ext), bbox_inches='tight')
plt.close(fig)

# Figure 6: protocol (wider layout and smaller top labels to avoid overlap)
fig, ax = plt.subplots(figsize=(9.0, 4.35))
ax.set_xlim(0,100)
ax.set_ylim(0,10.8)
ax.axis('off')
ax.text(50, 10.05, 'Data separation, long-horizon evaluation, and sealed-test gate', ha='center', fontsize=11.5, fontweight='bold', color='#1e293b')
segments = [
    (0,80,'Train / adaptation\n0–80 MB','#dbeafe','#60a5fa',7.7),
    (80,85,'Reused dev.\n80–85 MB','#fef3c7','#f59e0b',7.2),
    (85,90,'Prior dev.\nforbidden','#fee2e2','#f87171',6.6),
    (90,95,'Prior dev.\nforbidden','#fee2e2','#f87171',6.6),
    (95,100,'Sealed final\nnot opened','#ede9fe','#a78bfa',6.6),
]
for start_s,end_s,label,fc,ec,fs in segments:
    p = FancyBboxPatch((start_s, 6.65), end_s-start_s, 1.45, boxstyle='round,pad=0.03,rounding_size=0.05', facecolor=fc, edgecolor=ec, linewidth=1.0)
    ax.add_patch(p)
    ax.text((start_s+end_s)/2, 7.37, label, ha='center', va='center', fontsize=fs, color='#0f172a')
ax.text(0, 6.1, 'enwik8 byte ranges', ha='left', fontsize=8, color='#475569')
x_positions = [8, 27, 49, 70, 88]
labels = ['Shared warm start\nper seed','5k full-dev eval','10k full-dev eval','25k full-dev eval','50k endpoint']
for xp, label in zip(x_positions, labels):
    ax.scatter([xp], [3.7], s=44, color='#2563eb', zorder=3)
    ax.text(xp, 2.88, label, ha='center', va='top', fontsize=7.55)
for a,b in zip(x_positions[:-1], x_positions[1:]):
    arrow(a+1.4,3.7,b-1.4,3.7,color='#64748b')
ax.text(50, 4.88, '7 arms × 3 paired seeds = 21 training runs; 4 checkpoints each = 84 complete-development evaluations', ha='center', fontsize=8.25, color='#334155')
gate = FancyBboxPatch((78.8, 0.68), 19.4, 1.55, boxstyle='round,pad=0.05,rounding_size=0.1', facecolor='#eff6ff', edgecolor='#60a5fa', linewidth=1.15)
ax.add_patch(gate)
ax.text(88.5, 1.45, 'Quality gate AND\nmechanism gate', ha='center', va='center', fontsize=8.0, color='#1d4ed8')
arrow(88.7, 2.68, 88.7, 2.26, color='#64748b')
ax.text(61.5, 1.45, 'Gate failed → sealed final test remained unopened', ha='center', fontsize=8.2, fontweight='bold', color='#7c2d12')
arrow(78.6, 1.45, 72.0, 1.45, color='#f97316', ls='--')
ax.text(99.0, 8.75,
        'Top bands are dataset partitions; bottom timeline shows the\n'
        'preregistered evaluation landmarks and the fail-closed holdout policy.',
        ha='right', va='top', fontsize=7.0, color='#475569')
fig.tight_layout()
for ext in ['pdf','png']:
    fig.savefig(OUT / ('study_protocol.' + ext), bbox_inches='tight')
plt.close(fig)

print(f'Wrote figures to {OUT}')