#!/usr/bin/env python3
"""Data and figures for "Is RoPE just a recency bias?".

Every simulated number in the post comes from this file. Running it writes
CSV data to ../data and SVG figures to ../figures, then prints the numbers
the post quotes so they can be checked against the text.

    python make_figures.py               # light figures
    python make_figures.py --theme dark  # dark-background variants (*-dark.svg)

Setup shared by all figures: head dimension 128 (64 rotation pairs), RoPE base
10,000. Pair i rotates by omega_i = base^(-2i/d) radians per token, so pair 0
is the fastest and pair 63 the slowest.
"""
import argparse
import csv
from pathlib import Path

import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
from matplotlib.patches import Patch  # noqa: E402

HEAD_DIM = 128
BASE = 10_000.0
N_PAIRS = HEAD_DIM // 2
OMEGA = BASE ** (-2 * np.arange(N_PAIRS) / HEAD_DIM)

BANDS = [
    ("fast", "Fastest 16 pairs", np.arange(0, 16)),
    ("mid", "Middle 32 pairs", np.arange(16, 48)),
    ("slow", "Slowest 16 pairs", np.arange(48, 64)),
]
OFFSETS = np.array([0, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096])
SEED = 2
N_KEYS = 20_000
COSINES = [1.0, 0.5, 0.0]  # query-key cosine; 1.0 = query equal to key

HERE = Path(__file__).resolve().parent
DATA = HERE.parent / "data"
FIGS = HERE.parent / "figures"

THEMES = {
    "light": dict(bg="#ffffff", text="#1f1f1f", text2="#6b6b6b", grid="#e8e8e6",
                  axis="#bdbdb8", accent="#2a78d6", muted="#c9c8bf", ref="#8a8a85"),
    "dark": dict(bg="#1c1c1c", text="#ececec", text2="#a6a6a6", grid="#333333",
                 axis="#5a5a5a", accent="#5aa0f0", muted="#5f5f5a", ref="#9a9a95"),
}


# ---------------------------------------------------------------- simulations

def match_decay():
    """Figure 1. Match between a query and a key with given cosine, by offset.

    For each band, draw N_KEYS random unit keys restricted to that band's pairs.
    The query is cos * key + sqrt(1 - cos^2) * (unit vector orthogonal to key).
    The key is rotated by omega * offset in each pair (only the relative
    rotation matters), and we record the mean and sd of q . R(offset) k across
    keys. Draw order is fixed so results match the post exactly.
    """
    rng = np.random.default_rng(SEED)
    rows = []
    for slug, label, idx in BANDS:
        m = len(idx)
        for cos in COSINES:
            k = rng.standard_normal((N_KEYS, m, 2))
            k /= np.linalg.norm(k.reshape(N_KEYS, -1), axis=1)[:, None, None]
            z = rng.standard_normal((N_KEYS, m, 2))
            z -= (z * k).sum((1, 2))[:, None, None] * k
            z /= np.linalg.norm(z.reshape(N_KEYS, -1), axis=1)[:, None, None]
            q = cos * k + np.sqrt(1 - cos ** 2) * z
            for off in OFFSETS:
                a = OMEGA[idx] * off
                c, s = np.cos(a), np.sin(a)
                kr0 = c * k[..., 0] - s * k[..., 1]
                kr1 = s * k[..., 0] + c * k[..., 1]
                v = (q[..., 0] * kr0 + q[..., 1] * kr1).sum(-1)
                rows.append(dict(band=label, slug=slug, pairs=m, cosine=cos,
                                 offset=int(off), mean=round(float(v.mean()), 4),
                                 sd=round(float(v.std()), 4)))
    # Chance floor: for unrelated unit vectors in 2m dimensions, q . k has sd 1/sqrt(2m)
    # at every offset, since rotating a random direction gives another random direction.
    floor = [dict(band=label, slug=slug, pairs=len(idx),
                  sd=round(float(1 / np.sqrt(2 * len(idx))), 4)) for slug, label, idx in BANDS]
    return rows, floor


def chance_crossing(rows, floor, slug, cos=1.0):
    """First offset >= 1 where the mean match falls inside the chance band."""
    sd = next(f["sd"] for f in floor if f["slug"] == slug)
    for r in rows:
        if r["slug"] == slug and r["cosine"] == cos and r["offset"] >= 1 and abs(r["mean"]) < sd:
            return r["offset"]
    return None


def rotations_per_pair(contexts=(4096, 32768)):
    """Figure 2. Full turns each pair makes across a context of C tokens."""
    rows = []
    for ctx in contexts:
        turns = OMEGA * ctx / (2 * np.pi)
        for i, t in enumerate(turns):
            rows.append(dict(context=ctx, pair=i, turns=float(f"{t:.6g}")))
    return rows


def softmax(x):
    e = np.exp(x - x.max())
    return e / e.sum()


def head_vs_alibi(target=3, n_fast=8, scale=4.0, slope=0.5, max_offset=32):
    """Figure 3. Attention weights over offsets 0..max_offset.

    RoPE head: query aligned with the key `target` tokens back, using only the
    n_fast fastest pairs, so logit(offset) = scale * sum_i cos(omega_i * (offset - target)).
    ALiBi: logit(offset) = -slope * offset. Content is held constant.
    """
    off = np.arange(max_offset + 1)
    rope = softmax(scale * np.cos(np.outer(off - target, OMEGA[:n_fast])).sum(1))
    alibi = softmax(-slope * off)
    rows = [dict(head="RoPE head, 8 fastest pairs", slug="rope", offset=int(o),
                 weight=round(float(w), 6)) for o, w in zip(off, rope)]
    rows += [dict(head="ALiBi, slope 0.5", slug="alibi", offset=int(o),
                  weight=round(float(w), 6)) for o, w in zip(off, alibi)]
    return rows


# Results quoted from papers (not simulated), kept here so the data folder is complete.
DROPE_DROP_POINT = [
    dict(rope_removed_at="Step 0 (NoPE throughout)", validation_perplexity=23.77),
    dict(rope_removed_at="Step 8K", validation_perplexity=22.42),
    dict(rope_removed_at="Step 14K", validation_perplexity=21.73),
    dict(rope_removed_at="Never (RoPE throughout)", validation_perplexity=21.72),
]
DROPE_SOURCE = ("Gelberg et al. 2025, Table 6: 500M model, 16B tokens (16K steps); "
                "https://arxiv.org/abs/2512.12167")
PE_PERPLEXITY = [
    dict(encoding="RoPE", wikipedia_ppl=18.36, fineweb_ppl=22.44),
    dict(encoding="ALiBi", wikipedia_ppl=18.55, fineweb_ppl=22.18),
    dict(encoding="NoPE", wikipedia_ppl=19.86, fineweb_ppl=23.79),
]
PE_SOURCE = ("Bayesian Attention Mechanism paper, Table 2: 120M models trained at 512 tokens; "
             "https://arxiv.org/pdf/2505.22842")


# ---------------------------------------------------------------- output helpers

def write_csv(path, rows):
    """Plain CSV, header row first; column meanings are in the package README."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)


def fmt_offset(v):
    return f"{v // 1024}K" if v >= 1024 else f"{v}"


def style(t):
    plt.rcParams.update({
        "font.family": "sans-serif",
        "font.sans-serif": ["Inter", "Helvetica Neue", "Helvetica", "Arial", "DejaVu Sans"],
        "font.size": 10,
        "text.color": t["text"],
        "axes.labelcolor": t["text2"],
        "axes.edgecolor": t["axis"],
        "axes.facecolor": t["bg"],
        "figure.facecolor": t["bg"],
        "savefig.facecolor": t["bg"],
        "xtick.color": t["text2"],
        "ytick.color": t["text2"],
        "xtick.major.size": 0,
        "ytick.major.size": 0,
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.spines.left": False,
        "axes.spines.bottom": False,
        "svg.fonttype": "path",  # set to "none" to keep text as editable <text>
        "svg.hashsalt": "rope-recency",
    })


def header(fig, t, title, subtitle):
    fig.text(0.03, 0.955, title, fontsize=13, fontweight="bold", color=t["text"], va="top")
    fig.text(0.03, 0.885, subtitle, fontsize=9.5, color=t["text2"], va="top")


def save(fig, name, theme, alt):
    FIGS.mkdir(parents=True, exist_ok=True)
    suffix = "" if theme == "light" else f"-{theme}"
    path = FIGS / f"{name}{suffix}.svg"
    fig.savefig(path, format="svg", metadata={"Title": alt, "Date": None})
    plt.close(fig)
    return path


# ---------------------------------------------------------------- figures

def fig_match_decay(rows, floor, t, theme):
    cross = {slug: chance_crossing(rows, floor, slug) for slug, _, _ in BANDS}
    title = ("Fast, middle and slow pairs lose a matching key at offsets "
             f"{fmt_offset(cross['fast'])}, {fmt_offset(cross['mid'])} and {fmt_offset(cross['slow'])}")
    fig, axes = plt.subplots(1, 3, figsize=(7.6, 3.95), sharey=True)
    fig.subplots_adjust(left=0.07, right=0.98, top=0.63, bottom=0.15, wspace=0.12)
    header(fig, t, title, f"Query equal to key, with random content ({N_KEYS:,} keys per band; "
                          f"head dim {HEAD_DIM}, RoPE base {BASE:,.0f})")
    fig.legend(handles=[Patch(facecolor=t["muted"], alpha=0.5, label="Unrelated query and key, ±1 sd"),
                        Patch(facecolor=t["accent"], alpha=0.3,
                              label="Matching key: mean (line) and ±1 sd across keys")],
               loc="upper left", bbox_to_anchor=(0.025, 0.85), ncol=2, frameon=False,
               fontsize=9, labelcolor=t["text2"], handlelength=1.2, columnspacing=1.6)
    for ax, (slug, label, idx) in zip(axes, BANDS):
        pts = [r for r in rows if r["slug"] == slug and r["cosine"] == 1.0 and r["offset"] >= 1]
        x = np.array([r["offset"] for r in pts])
        mu = np.array([r["mean"] for r in pts])
        sd = np.array([r["sd"] for r in pts])
        fsd = next(f["sd"] for f in floor if f["slug"] == slug)
        ax.axhspan(-fsd, fsd, color=t["muted"], alpha=0.35, lw=0)
        for yv in (0.5, 1.0):
            ax.axhline(yv, color=t["grid"], lw=0.8, zorder=0)
        ax.axhline(0, color=t["axis"], lw=0.8, zorder=0)
        ax.fill_between(x, mu - sd, mu + sd, color=t["accent"], alpha=0.18, lw=0)
        ax.plot(x, mu, color=t["accent"], lw=2, marker="o", ms=3)
        c = cross[slug]
        ax.scatter([c], [mu[list(x).index(c)]], s=90, facecolors="none", edgecolors=t["text"], lw=1.3, zorder=5)
        ax.set_xscale("log", base=2)
        ax.set_xlim(0.75, 6200)
        ax.set_xticks([1, 8, 64, 512, 4096], ["1", "8", "64", "512", "4K"])
        ax.set_ylim(-0.5, 1.1)
        ax.set_yticks([0, 0.5, 1.0], ["0.0", "0.5", "1.0"])
        ax.text(0, 1.13, label, transform=ax.transAxes, fontsize=10, fontweight="bold",
                color=t["text"], va="bottom")
        ax.text(0, 1.03, f"Reaches chance at offset {fmt_offset(c)}", transform=ax.transAxes,
                fontsize=9, color=t["text2"], va="bottom")
    fig.supxlabel("Token offset (log scale)", fontsize=9.5, color=t["text2"], y=0.02)
    alt = ("Match between a query and an identical key by token offset, per RoPE frequency band, "
           "against the chance band for unrelated pairs.")
    return save(fig, "fig1-match-decay", theme, alt), cross


def fig_rotations(rows, t, theme, limit=0.25):
    # Past a quarter turn, cos(angle) < 0: the pair works against a matching key.
    n_under = {ctx: sum(1 for r in rows if r["context"] == ctx and r["turns"] < limit)
               for ctx in sorted({r["context"] for r in rows})}
    c_lo, c_hi = sorted(n_under)
    hi_count = "none do" if n_under[c_hi] == 0 else f"{n_under[c_hi]} do"
    title = (f"{n_under[c_lo]} of {N_PAIRS} pairs stay under a quarter turn in a {c_lo // 1024}K context; "
             f"at {c_hi // 1024}K, {hi_count}")
    fig, ax = plt.subplots(figsize=(7.6, 3.95))
    fig.subplots_adjust(left=0.09, right=0.83, top=0.78, bottom=0.17)
    header(fig, t, title, f"Full rotations of each RoPE pair across the context (head dim {HEAD_DIM}, "
                          f"base {BASE:,.0f})")
    ax.set_yscale("log")
    ax.set_ylim(0.01, 10_000)
    ax.set_yticks([0.01, 0.1, 1, 10, 100, 1000, 10000],
                  ["0.01", "0.1", "1", "10", "100", "1,000", "10,000"])
    ax.minorticks_off()
    ax.grid(axis="y", color=t["grid"], lw=0.8)
    ax.set_axisbelow(True)
    ax.set_xlim(-1, 64)
    ax.set_xticks([])
    for xv in (15.5, 47.5):
        ax.axvline(xv, color=t["axis"], lw=0.8, ls=(0, (2, 3)))
    for xc, lab in ((7.5, "Fastest 16"), (31.5, "Middle 32"), (55.5, "Slowest 16")):
        ax.text(xc, -0.05, lab, transform=ax.get_xaxis_transform(), ha="center", va="top",
                fontsize=9, color=t["text2"])
    ax.text(31.5, -0.15, "RoPE pair, fastest to slowest", transform=ax.get_xaxis_transform(),
            ha="center", va="top", fontsize=9.5, color=t["text2"])
    ax.axhline(limit, color=t["ref"], lw=1.3, ls=(0, (6, 4)))
    ax.text(0, limit * 1.25, "Quarter turn: a matching key's term falls to zero", fontsize=9,
            color=t["text2"], va="bottom")
    fig.text(0.03, 0.80, "Full turns (log scale)", fontsize=9, color=t["text2"], ha="left", va="bottom")
    # Name above the line end for the upper series, below it for the lower one.
    label_dy = {c_hi: (14, 1), c_lo: (-3, -17)}
    for ctx, alpha in ((c_lo, 1.0), (c_hi, 0.45)):
        pts = [r for r in rows if r["context"] == ctx]
        x = [r["pair"] for r in pts]
        y = [r["turns"] for r in pts]
        ax.plot(x, y, color=t["accent"], alpha=alpha, lw=2, marker="o", ms=2.5)
        dy_name, dy_sub = label_dy[ctx]
        ax.annotate(f"{ctx // 1024}K context", (x[-1], y[-1]), xytext=(10, dy_name),
                    textcoords="offset points", fontsize=10, color=t["text"], va="center",
                    annotation_clip=False)
        ax.annotate(f"{n_under[ctx]} under 1/4 turn", (x[-1], y[-1]), xytext=(10, dy_sub),
                    textcoords="offset points", fontsize=9, color=t["text2"], va="center",
                    annotation_clip=False)
    alt = ("Full rotations of each of the 64 RoPE pairs across a 4K and a 32K context, on a log scale, "
           "against a quarter-turn line.")
    return save(fig, "fig2-rotations-per-pair", theme, alt), n_under


def fig_head_vs_alibi(rows, t, theme, show=10):
    fig, axes = plt.subplots(1, 2, figsize=(7.6, 3.4), sharey=True)
    fig.subplots_adjust(left=0.07, right=0.98, top=0.71, bottom=0.17, wspace=0.12)
    header(fig, t, "RoPE can aim attention at offset 3; ALiBi always peaks at offset 0",
           "Attention weight by offset: toy head on the 8 fastest RoPE pairs (logit scale 4) "
           "vs ALiBi (slope 0.5)")
    peaks = {}
    for ax, slug, color in ((axes[0], "rope", t["accent"]), (axes[1], "alibi", t["muted"])):
        pts = [r for r in rows if r["slug"] == slug and r["offset"] <= show]
        x = [r["offset"] for r in pts]
        y = [r["weight"] for r in pts]
        for yv in (0.5, 1.0):
            ax.axhline(yv, color=t["grid"], lw=0.8, zorder=0)
        ax.axhline(0, color=t["axis"], lw=0.8, zorder=0)
        ax.bar(x, y, width=0.7, color=color, zorder=2)
        peak = max(pts, key=lambda r: r["weight"])
        peaks[slug] = peak
        ax.text(peak["offset"], peak["weight"] + 0.03, f"{peak['weight']:.3f}", ha="center",
                va="bottom", fontsize=9.5, color=t["text"])
        ax.set_xticks(range(show + 1))
        ax.set_ylim(0, 1.12)
        ax.set_yticks([0, 0.5, 1.0], ["0.0", "0.5", "1.0"])
        ax.set_title(pts[0]["head"], loc="left", fontsize=10, fontweight="bold", color=t["text"], pad=6)
    fig.supxlabel("Token offset (0 = current token)", fontsize=9.5, color=t["text2"], y=0.02)
    alt = "Attention weight by token offset for a toy RoPE positional head and for ALiBi."
    return save(fig, "fig3-rope-head-vs-alibi", theme, alt), peaks


# ---------------------------------------------------------------- main

def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--theme", choices=sorted(THEMES), default="light")
    args = ap.parse_args()
    t = THEMES[args.theme]
    style(t)

    decay, floor = match_decay()
    turns = rotations_per_pair()
    heads = head_vs_alibi()

    write_csv(DATA / "fig1_match_decay.csv", decay)
    write_csv(DATA / "fig1_chance_floor.csv", floor)
    write_csv(DATA / "fig2_rotations_per_pair.csv", turns)
    write_csv(DATA / "fig3_attention_weights.csv", heads)
    write_csv(DATA / "table_drope_drop_point.csv",
              [dict(r, source=DROPE_SOURCE) for r in DROPE_DROP_POINT])
    write_csv(DATA / "table_pe_perplexity.csv",
              [dict(r, source=PE_SOURCE) for r in PE_PERPLEXITY])

    p1, cross = fig_match_decay(decay, floor, t, args.theme)
    p2, n_under = fig_rotations(turns, t, args.theme)
    p3, peaks = fig_head_vs_alibi(heads, t, args.theme)

    # Numbers the post quotes, for checking against the text.
    half = {slug: chance_crossing(decay, floor, slug, cos=0.5) for slug, _, _ in BANDS}
    slowest_4k = next(r["turns"] for r in turns if r["context"] == 4096 and r["pair"] == N_PAIRS - 1)
    print("Wrote", p1.name, p2.name, p3.name, "and 6 CSVs")
    print("Chance floor sd per band:", {f["slug"]: f["sd"] for f in floor})
    print("Matching key reaches chance at:", cross, "| cosine 0.5:", half)
    print("Pairs under a quarter turn:", n_under, "| slowest pair at 4K:", slowest_4k, "turns")
    print("Peaks:", {k: (v["offset"], v["weight"]) for k, v in peaks.items()})


if __name__ == "__main__":
    main()