What does attention do inside an AI model?

"""Attention on one short sentence, step by step, in plain Python.

The sentence is "The cat sat because it was tired". Each word gets six numbers, and each head gets three small maps (query, key
and value). Every number is picked by hand so that the output is easy to read: nothing here was trained, and a real model learns
its numbers from data and uses thousands per word. The five steps below are the scaled dot-product attention of Vaswani et al.,
"Attention Is All You Need" (arXiv 1706.03762, section 3.2): for each word, score every word before it, hide the words after
it, turn the scores into weights that add up to 1, and take the weighted average of the earlier words' values.

    scores   = q . k / sqrt(d)            one per earlier word, d = how many numbers a query has
    weights  = softmax(scores)            positive, and they add up to 1
    output   = sum of weight * value      a blend of the earlier words' values

Python 3 standard library only. The widget (../widget.html) repeats these steps in JavaScript; test_widget.py runs both and
compares them, and test_attention.py checks this file against NumPy matrix operations.
"""
from __future__ import annotations

import math
import random

# --- the sentence, and six numbers per word: four about the word, then its place in the sentence (the cosine and sine of an angle) ---
TOKENS = ["The", "cat", "sat", "because", "it", "was", "tired"]
ABOUT = ["an animal", "pointing to something", "an action or state", "a small function word"]
CONTENT = [
    [0.0, 0.0, 0.0, 1.0],   # The
    [1.0, 0.0, 0.0, 0.0],   # cat      the only animal
    [0.0, 0.6, 1.0, 0.0],   # sat      an action that needs someone to do it
    [0.0, 0.0, 0.0, 1.0],   # because
    [0.0, 1.0, 0.0, 0.0],   # it       points to something, and carries nothing about what
    [0.0, 0.4, 1.0, 0.0],   # was
    [0.0, 0.8, 1.0, 0.0],   # tired    describes something
]
ANGLE = 0.9                  # radians between neighbouring positions


def position(p: int) -> list[float]:
    return [round(math.cos(ANGLE * p), 2), round(math.sin(ANGLE * p), 2)]


VECTORS = [CONTENT[p] + position(p) for p in range(len(TOKENS))]
D_MODEL = len(VECTORS[0])    # numbers per word: 6
D_K = 2                      # numbers per query and per key: 2
D_V = 4                      # numbers per value: 4


# --- three maps per head. A map is a table with D_MODEL rows and as many columns as the vector it makes. ---
def zeros(rows: int, cols: int) -> list[list[float]]:
    return [[0.0] * cols for _ in range(rows)]


def identity_on_content() -> list[list[float]]:
    value = zeros(D_MODEL, D_V)                  # the value keeps the four numbers about the word and drops the position
    for i in range(D_V):
        value[i][i] = 1.0
    return value


def refers_head(strength: float = 3.0) -> dict:
    """Looks for what a word points to: a word that points to something (column 1 of CONTENT) asks for an animal (column 0)."""
    query, key = zeros(D_MODEL, D_K), zeros(D_MODEL, D_K)
    query[1][0] = strength                       # "pointing to something" becomes the first number of the query
    key[0][0] = 1.0                              # "an animal" becomes the first number of the key
    return {"name": "refers", "wq": query, "wk": key, "wv": identity_on_content()}


def before_head(strength: float = 6.0) -> dict:
    """Looks for the word just before: the query is this word's position turned back by one step, so it matches the key of the word at p - 1."""
    c, s = round(strength * math.cos(ANGLE), 2), round(strength * math.sin(ANGLE), 2)
    query, key = zeros(D_MODEL, D_K), zeros(D_MODEL, D_K)
    query[4][0], query[4][1] = c, -s             # rows 4 and 5 are the position numbers
    query[5][0], query[5][1] = s, c
    key[4][0], key[5][1] = 1.0, 1.0              # the key is the position itself
    return {"name": "before", "wq": query, "wk": key, "wv": identity_on_content()}


def even_head() -> dict:
    """Asks for nothing: every query is all zeros, so every earlier word scores 0 and the weights come out equal."""
    return {"name": "even", "wq": zeros(D_MODEL, D_K), "wk": refers_head()["wk"], "wv": identity_on_content()}


HEADS = {"refers": refers_head(), "before": before_head(), "even": even_head()}
HEAD_LABELS = {"refers": "Refers to", "before": "Word before", "even": "No preference"}


# --- the five steps ---
def times(vector: list[float], table: list[list[float]]) -> list[float]:
    """A vector times a map: one output number per column."""
    return [sum(vector[r] * table[r][c] for r in range(len(vector))) for c in range(len(table[0]))]


def dot(a: list[float], b: list[float]) -> float:
    return sum(x * y for x, y in zip(a, b))


def scores_for(q: list[float], keys: list[list[float]], i: int, scale: bool = True) -> list[float]:
    """Step 2 and 3: q . k / sqrt(d) for every word, with -inf for the words after word i (the causal mask)."""
    divisor = math.sqrt(len(q)) if scale else 1.0
    return [dot(q, keys[j]) / divisor if j <= i else -math.inf for j in range(len(keys))]


def softmax(scores: list[float]) -> list[float]:
    """Step 4: positive weights that add up to 1. Subtracting the biggest score first changes nothing and avoids overflow; exp(-inf) is 0."""
    top = max(scores)
    sizes = [math.exp(s - top) for s in scores]
    total = sum(sizes)
    return [size / total for size in sizes]


def weighted_average(weights: list[float], values: list[list[float]]) -> list[float]:
    """Step 5: each number of the output is the weights times that number of every value, added up."""
    return [sum(w * v[c] for w, v in zip(weights, values)) for c in range(len(values[0]))]


def attend(vectors: list[list[float]], head: dict, scale: bool = True) -> dict:
    """Run one head over a sentence. Returns the queries, keys, values, scores, weights and outputs, one row per word."""
    queries = [times(x, head["wq"]) for x in vectors]       # step 1: what each word looks for ...
    keys = [times(x, head["wk"]) for x in vectors]          # ... the label each word holds up ...
    values = [times(x, head["wv"]) for x in vectors]        # ... and the note it carries
    scores = [scores_for(queries[i], keys, i, scale) for i in range(len(vectors))]
    weights = [softmax(row) for row in scores]
    outputs = [weighted_average(w, values) for w in weights]
    return {"queries": queries, "keys": keys, "values": values, "scores": scores, "weights": weights, "outputs": outputs}


def multi_head(vectors: list[list[float]], heads: list[dict]) -> list[list[float]]:
    """Several heads side by side: each word's outputs from every head, joined into one longer list."""
    runs = [attend(vectors, head)["outputs"] for head in heads]
    return [sum((run[i] for run in runs), []) for i in range(len(vectors))]


# --- what the page shows ---
def sentence(head: str, i: int) -> str:
    """The one line under the bars: where word i puts its weight."""
    w = attend(VECTORS, HEADS[head])["weights"][i]
    token = TOKENS[i]
    looked = [j for j in range(i + 1) if w[j] > 0]
    if len(looked) > 1 and max(w[:i + 1]) - min(w[:i + 1]) < 1e-9:
        return f"“{token}” spreads its weight evenly: {100 * w[0]:.0f}% on each of {len(looked)} words."
    top = max(range(len(w)), key=lambda j: w[j])
    return f"“{token}” puts {100 * w[top]:.0f}% of its weight on “{TOKENS[top]}”."


def blend_line(head: str, i: int) -> str:
    """The line about what the blend carries: the first number of a word's output says how much "an animal" it holds."""
    out = attend(VECTORS, HEADS[head])["outputs"][i]
    return f"Blended note, “an animal”: {out[0]:.2f}. “{TOKENS[i]}” alone: {CONTENT[i][0]:.2f}."


def percent(w: float) -> str:
    """A weight as a whole percent; a weight that is tiny but not zero shows as <1 so it never reads as zero."""
    return "<1" if 0 < w < 0.005 else f"{100 * w:.0f}"


def view(head: str, i: int) -> dict:
    """Everything the page writes for one head and one selected word. A word after word i is not looked at, so its cells are empty."""
    run = attend(VECTORS, HEADS[head])
    return {
        "weights": [percent(w) if j <= i else "" for j, w in enumerate(run["weights"][i])],
        "scores": [f"{s:.2f}" if j <= i else "" for j, s in enumerate(run["scores"][i])],
        "sentence": sentence(head, i), "blend": blend_line(head, i),
        "grid": [[percent(w) if j <= r else "" for j, w in enumerate(row)] for r, row in enumerate(run["weights"])],
    }


# --- why the division by sqrt(d) ---
def scaling_experiment(lengths=(4, 64, 1024), words=8, trials=100, seed=6) -> list[dict]:
    """Random queries and keys with numbers of mean 0 and spread 1, the assumption in the footnote in section 3.2.1 of the paper. For each
    length d, score one query against `words` keys, with and without the division, and report how spread out the scores are (their
    standard deviation) and how much weight the biggest score gets (the average over trials)."""
    rng = random.Random(seed)
    rows = []
    for d in lengths:
        raw, scaled, top_raw, top_scaled = [], [], [], []
        for _ in range(trials):
            q = [rng.gauss(0, 1) for _ in range(d)]
            s = [dot(q, [rng.gauss(0, 1) for _ in range(d)]) for _ in range(words)]
            raw += s
            scaled += [x / math.sqrt(d) for x in s]
            top_raw.append(max(softmax(s)))
            top_scaled.append(max(softmax([x / math.sqrt(d) for x in s])))
        rows.append({"d": d, "spread_raw": spread(raw), "spread_scaled": spread(scaled),
                     "top_raw": sum(top_raw) / trials, "top_scaled": sum(top_scaled) / trials})
    return rows


def spread(numbers: list[float]) -> float:
    mean = sum(numbers) / len(numbers)
    return math.sqrt(sum((x - mean) ** 2 for x in numbers) / len(numbers))


# --- what attention costs ---
LAYERS = 32     # copied from lesson 23 (prefill_sim.py): Llama 3.1 8B has 32 layers and 32 query heads per layer (they share 8 key-value heads)
HEADS_PER_LAYER = 32


def score_count(n: int) -> int:
    """Scores one head makes over a prompt of n words: word 1 scores 1 word, word 2 scores 2, ... word n scores n."""
    return n * (n + 1) // 2


def config() -> dict:
    """What the page needs to run the same steps: the words, the numbers per word, the maps of each head and the labels."""
    return {"tokens": TOKENS, "about": ABOUT, "vectors": VECTORS, "heads": HEADS, "labels": HEAD_LABELS,
            "opening": {"head": "refers", "token": TOKENS.index("it")}}

"""Keep the copy of the toy's numbers inside ../widget.html equal to the ones in attention.py.

    python3 embed_widget.py          rewrite the data block in widget.html
    python3 embed_widget.py --check  exit 1 if widget.html differs (test_widget.py runs the same check)
"""
import json
import re
import sys
from pathlib import Path

import attention as att

WIDGET = Path(__file__).resolve().parent.parent / "widget.html"
CONFIG_BLOCK = re.compile(r'(<script type="application/json" id="attn-config">)(.*?)(</script>)', re.S)


def embed(html):
    """html with the data block rewritten from attention.py."""
    settings = json.dumps(att.config(), ensure_ascii=False, separators=(", ", ": "))
    return CONFIG_BLOCK.sub(lambda m: m.group(1) + "\n" + settings + "\n" + m.group(3), html, count=1)


if __name__ == "__main__":
    current = WIDGET.read_text(encoding="utf-8")
    updated = embed(current)
    if "--check" in sys.argv:
        sys.exit(0 if updated == current else 1)
    WIDGET.write_text(updated, encoding="utf-8")
    print("widget.html data block now matches attention.py")

"""Print every number the lesson quotes: the toy's weights for "it", the full grids, the scaling experiment and the cost.

    python3 show_numbers.py
"""
import math

import attention as att

IT = att.TOKENS.index("it")


def show_it():
    print('The word "it" under the first head (what a word points to). d = 2 numbers per query, so scores are divided by sqrt(2) = %.4f' % math.sqrt(2))
    run = att.attend(att.VECTORS, att.HEADS["refers"])
    print("  query of it:", run["queries"][IT])
    print("  %-8s %-10s %-10s %s" % ("word", "key", "score", "weight"))
    for j, t in enumerate(att.TOKENS):
        score = "masked" if j > IT else "%.3f" % run["scores"][IT][j]
        print("  %-8s %-10s %-10s %.4f" % (t, run["keys"][j], score, run["weights"][IT][j]))
    print("  blended note of it:", [round(x, 3) for x in run["outputs"][IT]], "  about:", att.ABOUT)


def show_grids():
    for name, head in att.HEADS.items():
        run = att.attend(att.VECTORS, head)
        print("\nWeights in percent, head %r (%s). Row = the word asking, column = the word it looks at." % (name, att.HEAD_LABELS[name]))
        print("  %-8s" % "" + "".join("%8s" % t for t in att.TOKENS))
        for i, t in enumerate(att.TOKENS):
            print("  %-8s" % t + "".join("%8s" % att.percent(w) if j <= i else "%8s" % "." for j, w in enumerate(run["weights"][i])))
        print("  sentence:", att.view(name, IT)["sentence"], "|", att.view(name, IT)["blend"])
    both = att.multi_head(att.VECTORS, [att.HEADS["refers"], att.HEADS["before"]])
    print("\nTwo heads joined, outputs for it:", [round(x, 3) for x in both[IT]])


def show_scaling():
    print("\nScores of one random query against 8 random keys, 100 trials (numbers of mean 0 and spread 1, seed 6)")
    print("  %-6s %-22s %-22s %s" % ("d", "spread of scores", "biggest weight", "sqrt(d)"))
    print("  %-6s %-10s %-10s %-10s %-10s" % ("", "no divide", "divide", "no divide", "divide"))
    for r in att.scaling_experiment():
        print("  %-6d %-10.2f %-10.2f %-10.2f %-10.2f %.1f" % (r["d"], r["spread_raw"], r["spread_scaled"], r["top_raw"], r["top_scaled"], math.sqrt(r["d"])))


def show_cost():
    print("\nScores one head makes over a prompt, and for %d layers x %d heads" % (att.LAYERS, att.HEADS_PER_LAYER))
    for n in (10, 1000, 2000, 8000):
        print("  %6d words: %14s per head, %18s in all" % (n, f"{att.score_count(n):,}", f"{att.score_count(n) * att.LAYERS * att.HEADS_PER_LAYER:,}"))


if __name__ == "__main__":
    show_it()
    show_grids()
    show_scaling()
    show_cost()

"""Guards for attention.py: the toy's weights, the mask, weights that add up to 1, an independent NumPy check, the scaling numbers."""
import math
import random
import re
from pathlib import Path

import numpy as np

import attention as att

TOL = 1e-12                                        # agreement with NumPy, in absolute weight or output units
IT = att.TOKENS.index("it")


def test_the_toys_weights_for_it_match_the_hand_worked_numbers():
    """Guards a wrong weight: head 1 gives "it" a score of 3 / sqrt(2) on "cat" and 0 on the others, so cat gets e^2.121 / (e^2.121 + 4)."""
    refers = att.attend(att.VECTORS, att.HEADS["refers"])
    scores, weights = refers["scores"][IT], refers["weights"][IT]
    cat = att.TOKENS.index("cat")
    assert scores[:IT + 1] == [0.0, 3 / math.sqrt(2), 0.0, 0.0, 0.0] and scores[IT + 1:] == [-math.inf] * 2
    assert weights[cat] == math.exp(3 / math.sqrt(2)) / (math.exp(3 / math.sqrt(2)) + 4) and round(weights[cat], 3) == 0.676
    assert [round(w, 3) for w in weights] == [0.081, 0.676, 0.081, 0.081, 0.081, 0.0, 0.0]
    for word, size, weight in (("sat", 3.57, 0.641), ("tired", 5.46, 0.476)):                     # the other words the lesson works by hand
        row = att.TOKENS.index(word)
        assert round(math.exp(refers["scores"][row][cat]), 2) == size and round(refers["weights"][row][cat], 3) == weight
    assert refers["weights"][att.TOKENS.index("because")][:4] == [0.25] * 4                      # a query of all zeros: every score is 0
    before = att.attend(att.VECTORS, att.HEADS["before"])["weights"][IT]
    assert max(range(IT + 1), key=lambda j: before[j]) == att.TOKENS.index("because") and round(max(before), 2) == 0.71
    even = att.attend(att.VECTORS, att.HEADS["even"])["weights"][IT]
    assert even[:IT + 1] == [0.2] * 5
    assert att.view("refers", IT)["sentence"] == "“it” puts 68% of its weight on “cat”."
    assert att.view("refers", 0)["sentence"] == "“The” puts 100% of its weight on “The”."
    assert att.view("even", IT)["sentence"] == "“it” spreads its weight evenly: 20% on each of 5 words."
    assert att.view("refers", IT)["blend"] == "Blended note, “an animal”: 0.68. “it” alone: 0.00."
    assert len(att.VECTORS) == len(att.TOKENS) == 7 and {len(v) for v in att.VECTORS} == {att.D_MODEL}


def test_weights_are_positive_add_to_one_and_a_word_never_looks_ahead():
    """Guards the two failures a reader would see: a weight that is negative or the row not adding to 1, and a word whose answer changes
    when a LATER word changes. Runs every head with the division on and off, then rewrites every word after word i and checks rows 0..i."""
    n = len(att.TOKENS)
    for name, head in att.HEADS.items():
        for scale in (True, False):
            run = att.attend(att.VECTORS, head, scale)
            for i in range(n):
                row = run["weights"][i]
                assert abs(sum(row) - 1) < TOL and all(w > 0 for w in row[:i + 1]) and row[i + 1:] == [0.0] * (n - 1 - i), (name, scale, i)
                changed = [v if j <= i else [x + 1.7 for x in v] for j, v in enumerate(att.VECTORS)]
                other = att.attend(changed, head, scale)
                assert other["outputs"][i] == run["outputs"][i] and other["weights"][i] == row, (name, scale, i)


def reference(x, wq, wk, wv, scale=True):
    """Attention as NumPy matrix operations: softmax(Q K^T / sqrt(d) + mask) V, with the mask as -inf above the diagonal."""
    q, k, v = x @ wq, x @ wk, x @ wv
    s = q @ k.T / (math.sqrt(k.shape[1]) if scale else 1.0)
    s = s + np.triu(np.full(s.shape, -np.inf), k=1)
    e = np.exp(s - s.max(axis=1, keepdims=True))
    w = e / e.sum(axis=1, keepdims=True)
    return w, w @ v


def test_matches_numpy_on_the_toy_and_on_random_inputs():
    """Guards a slip in any one step (the division, the mask, the softmax, the average): both routes must agree within 1e-12."""
    rng = random.Random(0)
    cases = [(att.VECTORS, h) for h in att.HEADS.values()]
    for n, d_model, d_k, d_v in [(1, 3, 2, 2), (5, 4, 3, 4), (9, 8, 8, 5), (12, 16, 4, 6)]:
        def table(rows, cols):
            return [[rng.gauss(0, 1) for _ in range(cols)] for _ in range(rows)]
        cases.append((table(n, d_model), {"wq": table(d_model, d_k), "wk": table(d_model, d_k), "wv": table(d_model, d_v)}))
    for vectors, head in cases:
        for scale in (True, False):
            run = att.attend(vectors, head, scale)
            w, out = reference(np.array(vectors), np.array(head["wq"]), np.array(head["wk"]), np.array(head["wv"]), scale)
            assert np.allclose(run["weights"], w, rtol=0, atol=TOL) and np.allclose(run["outputs"], out, rtol=0, atol=TOL)


def test_two_heads_are_joined_side_by_side_without_mixing():
    """Guards heads sharing numbers: each word's joined output is head 1's output followed by head 2's, each as when run alone."""
    both = att.multi_head(att.VECTORS, [att.HEADS["refers"], att.HEADS["before"]])
    one, two = att.attend(att.VECTORS, att.HEADS["refers"])["outputs"], att.attend(att.VECTORS, att.HEADS["before"])["outputs"]
    assert both == [a + b for a, b in zip(one, two)] and len(both[IT]) == 2 * att.D_V


def test_scaling_experiment_gives_the_numbers_the_lesson_prints():
    """Guards the lesson's claim: without the division the scores spread out in step with sqrt(d) and the biggest word takes almost all
    the weight; with it the spread stays near 1 at every length."""
    rows = {r["d"]: r for r in att.scaling_experiment()}
    assert {d: round(r["spread_raw"], 1) for d, r in rows.items()} == {4: 2.2, 64: 8.1, 1024: 31.6}
    assert {d: round(r["spread_scaled"], 2) for d, r in rows.items()} == {4: 1.08, 64: 1.01, 1024: 0.99}
    assert {d: round(r["top_raw"], 2) for d, r in rows.items()} == {4: 0.53, 64: 0.90, 1024: 0.96}
    assert {d: round(r["top_scaled"], 2) for d, r in rows.items()} == {4: 0.34, 64: 0.36, 1024: 0.36}
    assert all(abs(r["spread_raw"] / math.sqrt(d) - 1) < 0.1 for d, r in rows.items())


def test_cost_counts_match_a_brute_force_count_and_lesson_23():
    """Guards the n-squared-over-2 claim and the copied model shape: the formula equals the number of unmasked cells, the three printed
    figures are right, and the layer and head counts are the ones in lesson 23's prefill_sim.py."""
    for n in range(1, 25):
        assert att.score_count(n) == sum(1 for i in range(n) for j in range(n) if j <= i)
    assert [att.score_count(n) for n in (1000, 8000)] == [500_500, 32_004_000]
    assert att.score_count(8000) * att.LAYERS * att.HEADS_PER_LAYER == 32_772_096_000
    other = (Path(__file__).resolve().parents[2] / "why-ai-waits-before-the-first-word" / "code" / "prefill_sim.py").read_text(encoding="utf-8")
    assert int(re.search(r"^LAYERS = (\d+)", other, re.M).group(1)) == att.LAYERS
    assert int(re.search(r"^HEADS = (\d+)", other, re.M).group(1)) == att.HEADS_PER_LAYER

"""Guards for ../widget.html: its data block, its size, that its JavaScript gives the same weights and words as attention.py, and what a reader sees."""
import html
import json
import math
import re
import subprocess

import attention as att
from embed_widget import CONFIG_BLOCK, WIDGET, embed

HTML = WIDGET.read_text(encoding="utf-8")
ROOT = HTML[HTML.index('<div class="attn-root">'):HTML.index("<script")]
TOLERANCE = 1e-12                                   # the JavaScript and Python weights and outputs may differ by at most this much

# Reads the model out of the widget between its marker comments, runs it with Node, and compares the weights and outputs with the numbers
# Python wrote into the fixture (within 1e-12) and the words and percentages as exact strings, for every head and every word.
PARITY = r"""
import fs from 'node:fs';
const [fixture, widget, tolerance] = process.argv.slice(2);
const fx = JSON.parse(fs.readFileSync(fixture, 'utf8'));
const html = fs.readFileSync(widget, 'utf8');
const cfg = JSON.parse(html.match(/<script type="application\/json" id="attn-config">([\s\S]*?)<\/script>/)[1]);
const js = html.slice(html.indexOf('/* ---------- the model ('), html.indexOf('/* ---------- display'));
const api = new Function('cfg', "'use strict';" + js + '; return {attend: attend, view: view};')(cfg);
const tol = Number(tolerance), bad = [];
let checks = 0, worst = 0;
const near = (a, b) => { const d = Math.abs(a - b); worst = Math.max(worst, d); return d <= tol; };
Object.keys(fx.heads).forEach((name) => {
  const run = api.attend(name), want = fx.heads[name];
  want.weights.forEach((row, i) => row.forEach((w, j) => { checks++; if (!near(run.weights[i][j], w)) bad.push(['weight', name, i, j, run.weights[i][j], w]); }));
  want.outputs.forEach((row, i) => row.forEach((o, c) => { checks++; if (!near(run.outputs[i][c], o)) bad.push(['output', name, i, c, run.outputs[i][c], o]); }));
  want.scores.forEach((row, i) => row.forEach((s, j) => {
    checks++;
    const got = run.scores[i][j];
    if (s === null ? got !== -Infinity : !near(got, s)) bad.push(['score', name, i, j, got, s]);
  }));
  want.views.forEach((v, i) => {
    checks++;
    if (JSON.stringify({ ...api.view(name, i), top: undefined }) !== JSON.stringify({ ...v, top: undefined })) bad.push(['view', name, i, api.view(name, i), v]);
  });
});
console.log(JSON.stringify({ checks, worst, bad: bad.slice(0, 4) }));
process.exit(bad.length ? 1 : 0);
"""

# Runs the page script against a small stand-in for the DOM, presses every word and every head button, and reads back what a reader sees.
PAGE = r"""
import fs from 'node:fs';
const [widget, fixture] = process.argv.slice(2);
const html = fs.readFileSync(widget, 'utf8');
const fx = JSON.parse(fs.readFileSync(fixture, 'utf8'));
const config = html.match(/<script type="application\/json" id="attn-config">([\s\S]*?)<\/script>/)[1];
const code = html.match(/<script>([\s\S]*?)<\/script>/)[1];

function el(tag) {
  const e = { tag, children: [], listeners: {}, style: {}, attrs: {}, textContent: '' };
  e.setAttribute = (k, v) => { e.attrs[k] = v; };
  e.addEventListener = (type, f) => (e.listeners[type] = e.listeners[type] || []).push(f);
  e.appendChild = (c) => (e.children.push(c), c);
  return e;
}
function open() {
  const page = { root: el('html'), messages: [], parent: { other: true } };
  fx.ids.forEach((id) => { page[id] = el('x'); });
  page['attn-config'].textContent = config;
  const doc = { getElementById: (id) => page[id], createElement: el, documentElement: page.root };
  const win = { addEventListener: (type, f) => page.messages.push(f) };
  new Function('document', 'window', 'parent', code)(doc, win, page.parent);
  return page;
}
const fire = (e) => (e.listeners.click || []).forEach((f) => f({ target: e }));
const bad = [];
const want = (what, got, expected) => { if (JSON.stringify(got) !== JSON.stringify(expected)) bad.push([what, got, expected]); };
const pct = (c) => (c === '<1' ? 0.5 : Number(c));
const mix = (td) => { const m = /var\(--green\) (\d+)%/.exec(td.style.background || ''); return m ? Number(m[1]) : NaN; };
function look(page, head, i, what) {                 // compares the page with Python's view of one head and one selected word
  const v = fx.views[head][i], rows = page['attn-bars'].children;
  want(what + ': sentence', page['attn-sentence'].textContent, v.sentence);
  want(what + ': blend', page['attn-blend'].textContent, v.blend);
  want(what + ': scores line', page['attn-scores'].textContent, v.scoresLine);
  want(what + ': word pressed', page['attn-tokens'].children.map((b) => b.attrs['aria-pressed']), fx.tokens.map((t, k) => String(k === i)));
  want(what + ': head pressed', page['attn-heads'].children.map((b) => b.attrs['aria-pressed']), fx.names.map((n) => String(n === head)));
  want(what + ': percent column', rows.map((li) => li.children[2].textContent), v.weights.map((w) => (w === '' ? '' : w + '%')));
  want(what + ': bar widths', rows.map((li) => li.children[1].children[0].style.width), v.widths);
  want(what + ': later words say so', rows.map((li) => [li.attrs['data-later'], li.children[1].children[1].style.display]), v.weights.map((w) => [String(w === ''), w === '' ? '' : 'none']));
  want(what + ': the biggest is marked', rows.map((li) => li.attrs['data-top']), v.weights.map((w, j) => String(j === v.top)));
  const grid = page['attn-grid'].children;
  want(what + ': grid size', grid.length, fx.tokens.length + 1);
  want(what + ': grid numbers', grid.slice(1).map((tr) => tr.children.slice(1).map((td) => td.textContent)), v.grid);
  // shaded only where the word looks; in a row whose weights differ the biggest cell is mixed with more green than the smallest, and no mix passes 50%
  want(what + ': grid shading only where looked, darker for a bigger weight, never above half green', grid.slice(1).map((tr, r) => {
    const cells = tr.children.slice(1), shares = v.grid[r].map((c) => (c === '' ? NaN : pct(c))), seen = shares.map((s, j) => j).filter((j) => !Number.isNaN(shares[j]));
    const mixes = cells.map((td) => ((td.style.background || '') !== '' ? mix(td) : NaN));
    const big = seen.reduce((a, j) => (shares[j] > shares[a] ? j : a), seen[0]), small = seen.reduce((a, j) => (shares[j] < shares[a] ? j : a), seen[0]);
    return [cells.map((td) => (td.style.background || '') !== ''), shares[big] > shares[small] ? mixes[big] > mixes[small] : true, Math.max(...seen.map((j) => mixes[j])) <= 50];
  }), v.grid.map((r) => [r.map((c) => c !== ''), true, true]));
  want(what + ': grid row marked', grid.slice(1).map((tr) => tr.attrs['data-selected']), fx.tokens.map((t, k) => String(k === i)));
}

const first = open();
look(first, fx.opening.head, fx.opening.token, 'opening state');
want('nothing spoken yet', first['attn-live'].textContent, '');
want('word buttons', first['attn-tokens'].children.map((b) => b.textContent), fx.tokens);
want('head buttons', first['attn-heads'].children.map((b) => b.textContent), fx.names.map((n) => fx.labels[n]));
want('grid headers', first['attn-grid'].children[0].children.slice(1).map((th) => th.children[0].textContent), fx.tokens);

// every head with every word, by pressing the buttons
const page = open();
fx.names.forEach((head, h) => {
  fire(page['attn-heads'].children[h]);
  fx.tokens.forEach((t, i) => {
    fire(page['attn-tokens'].children[i]);
    look(page, head, i, head + ' / ' + t);
    want(head + ' / ' + t + ': spoken', page['attn-live'].textContent, fx.views[head][i].sentence + ' ' + fx.views[head][i].blend);
  });
});

// the host page theme: only the parent window may set it, and only to light or dark
const themed = open();
const send = (data, source) => themed.messages.forEach((f) => f({ data, source }));
send({ type: 'spicyshark:theme', theme: 'dark' }, { other: false });
send({ type: 'spicyshark:theme', theme: 'purple' }, themed.parent);
send({ type: 'other', theme: 'dark' }, themed.parent);
send(null, themed.parent);
want('ignored messages', themed.root.attrs, {});
send({ type: 'spicyshark:theme', theme: 'dark' }, themed.parent);
want('theme from the parent', themed.root.attrs, { 'data-theme': 'dark' });
send({ type: 'spicyshark:theme', theme: 'light' }, themed.parent);
want('theme back to light', themed.root.attrs, { 'data-theme': 'light' });

console.log(JSON.stringify({ heads: fx.names.length, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


def run_node(tmp_path, script, fixture, args):
    (tmp_path / "fixture.json").write_text(json.dumps(fixture), encoding="utf-8")
    (tmp_path / "check.mjs").write_text(script, encoding="utf-8")
    run = subprocess.run(["node", str(tmp_path / "check.mjs"), *args(tmp_path)], capture_output=True, text=True, timeout=120)
    assert run.returncode == 0, run.stdout + run.stderr
    return json.loads(run.stdout)


def share(text):
    return 0.5 if text == "<1" else float(text)


def test_widget_data_block_matches_the_python_numbers():
    """Guards a toy edit that was not copied into the widget (run embed_widget.py), and an embed that skips a drifted block."""
    assert embed(HTML) == HTML
    drifted = HTML.replace('"opening": {"head": "refers"', '"opening": {"head": "even"', 1).replace('"tokens": ["The", "cat"', '"tokens": ["A", "cat"', 1)
    assert drifted != HTML and embed(drifted) == HTML


def test_every_element_the_script_looks_up_exists_inside_the_root():
    """Guards a widget that never starts: the parity run covers only the formulas, so a misspelt id passed there and threw in the browser."""
    script = HTML[HTML.index("<script>"):]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", script))
    missing = sorted(i for i in lookups if i != "attn-config" and f'id="{i}"' not in ROOT)
    assert not missing, missing
    assert 'id="attn-config"' in HTML


def test_widget_is_small_and_asks_for_nothing_online():
    assert len(HTML.splitlines()) <= 260 and len(HTML.encode("utf-8")) < 40 * 1024
    page = CONFIG_BLOCK.sub("", HTML)
    for banned in ("http://", "https://", "fetch(", "XMLHttpRequest", "WebSocket", "sendBeacon", "import("):
        assert banned not in page, banned
    assert " src=" not in page and "@import" not in page
    assert "<form" not in HTML and 'type="submit"' not in HTML          # a form cannot submit inside the lesson page's sandboxed frame
    assert "transition" not in HTML and "animation" not in HTML         # nothing moves, so prefers-reduced-motion has nothing to switch off


def test_widget_shows_at_most_five_controls_and_hides_the_rest():
    """Guards the first screen at 390 px: the row of words counts as one control, each head button as one, and the grid sits in a closed details element."""
    outside = re.sub(r"<details>.*?</details>", "", ROOT, flags=re.S)
    assert len(re.findall(r"<(?:select|button|input|textarea)\b", outside)) + 1 + len(att.HEADS) <= 5
    assert "<details><summary>More settings</summary>" in ROOT and "<details open" not in ROOT
    assert re.search(r"<details><summary>More settings</summary>.*?id=\"attn-grid\"", ROOT, re.S)


def test_widget_keeps_the_site_theme_listener_and_palette():
    assert "d.type !== 'spicyshark:theme'" in HTML and "e.source !== parent" in HTML and "parent === window" in HTML
    for pair in ("--bg: light-dark(#ffffff, #0a0f17)", "--panel: light-dark(#ebf0f6, #131b28)", "--ink: light-dark(#0a1626, #f4f7fb)",
                 "--mute: light-dark(#44566b, #a3b1c4)", "--ctl: light-dark(#66768a, #7a8a9f)", "--ink-on: light-dark(#f6f8fb, #05070c)",
                 "--green: light-dark(#00855F, #009E73)", ":root[data-theme=light]", ":root[data-theme=dark]"):
        assert pair in HTML, pair
    for banned in ("#0072B2", "#56B4E9", "#D55E00"):
        assert banned.lower() not in HTML.lower()
    assert not set(re.findall(r"#[0-9a-fA-F]{3,8}\b", HTML)) - {"#ffffff", "#0a0f17", "#ebf0f6", "#131b28", "#0a1626", "#f4f7fb", "#44566b", "#a3b1c4",
                                                               "#66768a", "#7a8a9f", "#f6f8fb", "#05070c", "#00855F", "#009E73"}


def test_widget_formulas_give_the_weights_and_words_of_the_python_toy(tmp_path):
    """Guards the widget's JavaScript drifting from attention.py: the lesson prints Python's numbers and the reader sees the widget's, so
    the two must agree within 1e-12 on every weight, score and output, and exactly on every percentage and sentence, for every head and word."""
    heads = {}
    for name, head in att.HEADS.items():
        run = att.attend(att.VECTORS, head)
        heads[name] = {"weights": run["weights"], "outputs": run["outputs"],
                       "scores": [[None if s == -math.inf else s for s in row] for row in run["scores"]],
                       "views": [att.view(name, i) for i in range(len(att.TOKENS))]}
    result = run_node(tmp_path, PARITY, {"heads": heads}, lambda d: [str(d / "fixture.json"), str(WIDGET), str(TOLERANCE)])
    n, h = len(att.TOKENS), len(att.HEADS)
    assert result["bad"] == [] and result["checks"] == h * (2 * n * n + n * att.D_V + n) and result["worst"] <= TOLERANCE


def test_widget_page_shows_what_the_lesson_names(tmp_path):
    """Guards the part a reader sees: it opens on "it" looking for what it points to with 68% on "cat", every head and every word give
    the bars, sentence, blend line and grid Python computes, later words say they are not looked at, the buttons speak, and the grid figure
    prints the same percentages."""
    opening = att.view("refers", att.TOKENS.index("it"))
    assert opening["sentence"] == "“it” puts 68% of its weight on “cat”." and json.loads(CONFIG_BLOCK.search(HTML).group(2))["opening"] == att.config()["opening"]
    visible = re.sub(r"<details>.*?</details>", "", ROOT, flags=re.S)                                                                 # what shows without opening anything
    for line in ("A toy: every number is picked by hand, not learned.", "A weight is a share of attention."):
        assert line in visible, line
    assert 'class="attn-sr" id="attn-live" role="status" aria-live="polite" aria-atomic="true"' in ROOT                               # what a screen reader hears
    assert "grid-template-columns: repeat(auto-fit, minmax(var(--min), 1fr))" in HTML and ".attn-root .attn-row button { white-space: nowrap;" in HTML   # an even grid: a button wraps whole, never its label, never stretched
    labels = re.findall(r'class="attn-label"[^>]*>([^<]+)<', ROOT) + list(att.HEAD_LABELS.values())
    assert len(labels) == 2 + len(att.HEADS) and not [t for t in labels if {w.lower() for w in t.split()} & {w.lower() for w in att.TOKENS}]   # no label reads as a word of the sentence ("it")
    lesson, legend = (WIDGET.parent / "lesson.md").read_text(encoding="utf-8"), (WIDGET.parent / "figs" / "two-heads.svg").read_text(encoding="utf-8")
    for key in ("refers", "before"):                                                                                                  # figure 3 and its caption print the widget's own labels
        label = att.HEAD_LABELS[key]
        assert f">{label}</text>" in legend and f"labelled {label}," in legend and f'"{label}"' in lesson, key        # drawn legend, the figure's description, the caption
    named = set(re.findall(r'In the head (?:called )?"([^"]+)"', lesson))                                                             # and so does every sentence that names a head
    assert named and named <= set(att.HEAD_LABELS.values()), named
    grid_svg = (WIDGET.parent / "figs" / "grid.svg").read_text(encoding="utf-8")                                                     # the figure prints the percentages Python computes
    assert [html.unescape(t) for t in re.findall(r'text-anchor="middle"[^>]*>([^<]+)</text>', grid_svg)] == [c for row in att.view("refers", len(att.TOKENS) - 1)["grid"] for c in row if c]
    views = {}
    for name in att.HEADS:
        views[name] = []
        for i in range(len(att.TOKENS)):
            v = att.view(name, i)
            top = -1 if (i > 0 and len(set(v["weights"][:i + 1])) == 1) else max(range(i + 1), key=lambda j: (share(v["weights"][j]), -j))
            v["top"] = top
            v["widths"] = [f"{share(w):g}%" if w != "" else "0%" for w in v["weights"]]
            v["scoresLine"] = (f"Scores before softmax for “{att.TOKENS[i]}”: "
                               + ", ".join(f"{att.TOKENS[j]} {s}" for j, s in enumerate(v["scores"]) if s != "") + ".")
            views[name].append(v)
    ids = sorted(set(re.findall(r"\$\('([\w-]+)'\)", HTML[HTML.index("<script>"):])))
    fixture = {"ids": ids, "tokens": att.TOKENS, "names": list(att.HEADS), "labels": att.HEAD_LABELS, "views": views, "opening": att.config()["opening"]}
    result = run_node(tmp_path, PAGE, fixture, lambda d: [str(WIDGET), str(d / "fixture.json")])
    assert result == {"heads": len(att.HEADS), "bad": []}