How can a small model speed up a big one?

"""Keep the copy of the model, the chip and the reply inside ../widget.html equal to the ones in speculative_sim.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 speculative_sim as sim

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


def config():
    """What the page needs to run the same rounds: the big model and the chip, the reply and its wrong guesses, and the settings it offers."""
    return {
        "params": sim.big.PARAMS, "bytesPerWeight": sim.big.FP16_BYTES, "peak": sim.big.A100_PEAK_FLOPS, "bandwidth": sim.big.A100_BANDWIDTH,
        "reply": sim.REPLY, "decoys": sim.DECOYS, "draftStops": sim.DRAFT_STOPS_MS, "defaultDraft": sim.DEFAULT_DRAFT_STOP, "opening": sim.OPENING,
    }


def embed(html):
    """html with the data block rewritten from speculative_sim.py."""
    settings = json.dumps(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 speculative_sim.py")

"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
import speculative_sim as sim

big = sim.big


def main() -> None:
    plain, draft = sim.plain_seconds(), sim.draft_pass().seconds
    print(f"big model {big.PARAMS:,} weights, {big.FP16_BYTES * big.PARAMS / 1e9:.2f} GB; draft model {sim.DRAFT_PARAMS:,} weights, {big.FP16_BYTES * sim.DRAFT_PARAMS / 1e9:.2f} GB")
    print(f"plain big pass {plain * 1000:.3f} ms; draft pass {draft * 1000:.3f} ms = {100 * draft / plain:.1f}% of a big pass (c = {draft / plain:.4f}); ridge {big.ridge():.2f}")
    print()
    print("a big pass over n inputs (the guesses plus one):")
    for n in (1, 3, 5, 9, 100, 153, 154, 200, 1000):
        p = big.prefill(n)
        print(f"  n={n:5,}: {p.seconds * 1000:8.3f} ms, {p.bound}")
    print()
    print("one round at k = 2, 4, 8:", {k: f"{sim.round_seconds(k) * 1000:.3f} ms" for k in (2, 4, 8)})
    print()
    print("tokens per big pass and speed-up at k = 4 (guesses kept with chance a, independently):")
    for a in (0.3, 0.5, 0.6, 0.8, 0.9, 0.95):
        print(f"  a={a}: {sim.expected_tokens(a, 4):.2f} tokens per pass, {sim.speedup(a, 4):.2f} times as fast, simulated {sim.simulate_tokens_per_pass(a, 4, 100_000, seed=4):.2f}")
    print("break-even acceptance:", {k: round(sim.break_even(k), 3) for k in (1, 2, 4, 8)})
    print("best k and its speed-up:", {a: (sim.best_k(a), round(sim.speedup(a, sim.best_k(a)), 2)) for a in (0.5, 0.6, 0.8, 0.9, 0.95)})
    print("speed-up at a = 0.6, 0.8, 0.9 for k = 1..12:")
    for a in (0.6, 0.8, 0.9):
        print(f"  a={a}:", [round(sim.speedup(a, k), 2) for k in range(1, 13)])
    print()
    print("many users sharing the chip (the check covers users * (k + 1) inputs):")
    print("  users when the check stops being free:", {k: sim.users_when_checking_stops_being_free(k) for k in (2, 4, 6, 8)})
    print("  speed-up at a = 0.8, k = 4:", {u: round(sim.speedup(0.8, 4, users=u), 2) for u in (1, 16, 30, 31, 32, 64, 128, 154, 1000)})
    print("  users when it stops paying (a = 0.8):", {k: sim.users_when_it_stops_paying(0.8, k) for k in (2, 4, 8)})
    print()
    print(f"toy pair: draft agrees with the big model at {float(sim.agreement()) * 100:.1f}% of places ({sim.agreement()}); measured against the formula at that chance")
    for k in (1, 2, 4, 8):
        print(f"  k={k}: measured {sim.toy_tokens_per_pass(k):.2f} tokens per pass, formula {sim.expected_tokens(float(sim.agreement()), k):.2f}")
    print(f"  reply: {' '.join(sim.REPLY)}")
    print()
    p, q = sim.HAND_P, sim.HAND_Q
    print("hand example (big model p, small model q):", {x: f"{p[x]} / {q[x]}" for x in p})
    print("  chance a guess x is drawn and kept:", {x: str(q[x] * sim.keep_chance(p[x], q[x])) for x in p}, "total kept", sim.keep_total(p, q))
    print("  redraw after a rejection:", {x: str(v) for x, v in sim.residual(p, q).items()}, "-> output", {x: str(v) for x, v in sim.output_distribution(p, q).items()})
    print()
    print("widget sentences:")
    for pct, k in ((sim.OPENING["percent"], sim.OPENING["k"]), (20, 4), (50, 4), (100, 8)):
        v = sim.view(pct, k, sim.DRAFT_STOPS_MS[sim.DEFAULT_DRAFT_STOP], sim.OPENING["seed"])
        print(f"  {pct}% k={k}: {v['spec']} {v['verdict']} | {v['average']}")
    print(f"  plain: {sim.view(80, 4, sim.DRAFT_STOPS_MS[2], sim.OPENING['seed'])['plain']}")
    print(f"  draft stops (ms): {[round(x, 3) for x in sim.DRAFT_STOPS_MS]}; opening run rounds {sim.plan_run(sim.OPENING['percent'], sim.OPENING['k'], sim.OPENING['seed'])}")


if __name__ == "__main__":
    main()

"""Speculative decoding: a small model guesses a few words, the big model checks them all in one pass.

Four parts, each short:

  1. A toy pair of word models (a big one that looks back two words, a small one that looks back one), built from one
     sentence. `speculative_greedy` runs the round exactly as the papers describe it, and test_speculative_sim.py checks
     that the text it keeps is the big model's own text, word for word.
  2. The rule for sampled text (`keep_chance`, `residual`, `output_distribution`), checked with exact fractions.
  3. Tokens per big pass: a seeded simulation (`simulate_tokens_per_pass`) and the formula it should match (`expected_tokens`).
  4. Time: lesson 23's roofline model for the big model (Llama 3.1 8B on an A100) plus a draft model of 1.24 billion
     weights (Llama 3.2 1B). Every number is a best case at the chip's data-sheet peak, not a measurement.

Python 3 standard library only. The widget (../widget.html) repeats `plan_run` and `view` in JavaScript;
test_widget.py runs both and compares them.
"""
from __future__ import annotations

import importlib.util
import random
import sys
from collections import Counter
from dataclasses import dataclass
from decimal import ROUND_HALF_UP, Decimal
from fractions import Fraction
from pathlib import Path

sys.dont_write_bytecode = True                           # loading lesson 23's file below must not leave a __pycache__ in its folder
LESSON_23 = Path(__file__).resolve().parents[2] / "why-ai-waits-before-the-first-word" / "code" / "prefill_sim.py"


def _lesson_23():
    spec = importlib.util.spec_from_file_location("prefill_sim_lesson_23", LESSON_23)
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module                      # dataclasses looks the module up by name
    spec.loader.exec_module(module)
    return module


big = _lesson_23()                                       # the big model and the chip: Llama 3.1 8B on an A100 80GB SXM, 8,030,261,248 weights

# --- the draft model: Llama 3.2 1B. Shape typed from config.json of huggingface.co/unsloth/Llama-3.2-1B, a third-party copy of Meta's
# model (base_model meta-llama/Llama-3.2-1B); the total below is what that page lists for the weight files (safetensors metadata, BF16).
# Read 2026-10-19. Its output table is the input table reused ("tied"), and its vocabulary is the big model's, as a draft model needs.
# Meta's own model card (llama-models, models/llama3_2/MODEL_CARD.md) lists the 1B model as 1.23B parameters with shared embeddings.
D_HIDDEN, D_FFN, D_LAYERS, D_HEADS, D_KV_HEADS, D_HEAD_DIM, D_VOCAB = 2048, 8192, 16, 32, 8, 64, 128256
D_WEIGHT_FILES_PARAMS = 1_235_814_400


def draft_parameter_count() -> int:
    """Count the draft model's weights from its config, as lesson 23 does for the big one, but with one table for input and output."""
    attention = 2 * D_HIDDEN * D_HIDDEN + 2 * D_HIDDEN * D_KV_HEADS * D_HEAD_DIM      # query and output matrices, then key and value matrices
    layer = attention + 3 * D_HIDDEN * D_FFN + 2 * D_HIDDEN                          # feed-forward gate, up and down, then two norm vectors
    return D_LAYERS * layer + D_VOCAB * D_HIDDEN + D_HIDDEN                          # one shared table, final norm


DRAFT_PARAMS = draft_parameter_count()


# --- 1. the toy pair ---------------------------------------------------------------------------

TEXT = "boil the water in the kettle . pour it over the tea bag . wait for three minutes . take the bag out . add a little milk . stir with a spoon . drink it while it is hot ."
WORDS = TEXT.split()
PROMPT = WORDS[:2]                                       # the widget starts a reply after "boil the"


def _followers(order: int) -> dict:
    """For each run of `order` words in the text, how often each word came next. The text is a ring, so the last words are followed by the first."""
    ring = WORDS + WORDS[:order]
    table: dict = {}
    for i in range(order, len(ring)):
        table.setdefault(tuple(ring[i - order:i]), Counter())[ring[i]] += 1
    return table


BIG_TABLE, SMALL_TABLE = _followers(2), _followers(1)


def _ranked(counts: Counter) -> list[str]:
    """Most common follower first; a tie goes to the word that sorts first, so the same input always gives the same answer."""
    return [word for word, _ in sorted(counts.items(), key=lambda item: (-item[1], item[0]))]


def big_ranked(context: list[str]) -> list[str]:
    """The big toy model: looks back two words, and falls back to one word if that pair never occurred."""
    pair = tuple(context[-2:])
    return _ranked(BIG_TABLE[pair]) if pair in BIG_TABLE else _ranked(SMALL_TABLE[pair[-1:]])


def small_ranked(context: list[str]) -> list[str]:
    """The small toy model: looks back one word."""
    return _ranked(SMALL_TABLE[tuple(context[-1:])])


def big_choice(context: list[str]) -> str:
    return big_ranked(context)[0]


def small_choice(context: list[str]) -> str:
    return small_ranked(context)[0]


def greedy(choose, prompt: list[str], count: int) -> list[str]:
    """Plain decoding: one model, always its top choice, one word per pass."""
    text = list(prompt)
    for _ in range(count):
        text.append(choose(text))
    return text[len(prompt):]


@dataclass(frozen=True)
class Round:
    guesses: list[str]     # what the small model wrote, one cheap pass each
    own: list[str]         # what the big model would write after the text and after each guess prefix, all from one pass over k + 1 inputs
    kept: int              # guesses kept: the run of matches from the start
    new: list[str]         # words added to the text this round: the kept guesses, then the big model's own word


def speculative_greedy(prompt: list[str], count: int, k: int, draft=small_choice, target=big_choice) -> tuple[list[str], list[Round]]:
    """Speculative decoding for greedy text. Each round: the draft writes k guesses; the target looks at the last word and all k guesses
    in one pass and says what it would write at each of the k + 1 places; keep guesses up to the first place where they differ;
    there, use the target's word; if every guess was kept, the target's word after the last guess comes free."""
    text, rounds = list(prompt), []
    while len(text) - len(prompt) < count:
        guesses: list[str] = []
        for _ in range(k):
            guesses.append(draft(text + guesses))
        own = [target(text + guesses[:i]) for i in range(k + 1)]
        kept = 0
        while kept < k and guesses[kept] == own[kept]:
            kept += 1
        new = guesses[:kept] + [own[kept]]
        rounds.append(Round(guesses, own, kept, new))
        text += new
    return text[len(prompt):len(prompt) + count], rounds


def agreement() -> Fraction:
    """How often the two toy models pick the same next word, over every place in the ring of text."""
    ring = WORDS + WORDS
    n = len(WORDS)
    same = sum(1 for i in range(n) if big_choice(ring[i:i + 2]) == small_choice(ring[i:i + 2]))
    return Fraction(same, n)


def toy_tokens_per_pass(k: int, count: int = 600) -> float:
    """Words per big pass for the toy pair, over a start at every place in the text. The guesses here are not independent chances,
    so this is the number to hold against `expected_tokens` at the same agreement."""
    ring = WORDS * 2
    words = passes = 0
    for i in range(len(WORDS)):
        _, rounds = speculative_greedy(ring[i:i + 2], count, k)
        words += count
        passes += len(rounds)
    return words / passes


# --- 2. the rule for sampled text ---------------------------------------------------------------
# Leviathan, Kalman and Matias (2023) write p for the big model and q for the draft; Chen et al. (2023) swap the letters. This file uses Leviathan's.

def keep_chance(p_x: Fraction, q_x: Fraction) -> Fraction:
    """The big model keeps a guess x with probability min(1, p(x) / q(x))."""
    return min(Fraction(1), p_x / q_x) if q_x else Fraction(1)


def keep_total(p: dict, q: dict) -> Fraction:
    """The chance that a guess is kept at all: the sum of q(x) * min(1, p(x) / q(x)) over every word, which is the sum of min(p, q)."""
    assert p.keys() == q.keys(), "p and q must give odds for the same tokens"
    return sum((q[x] * keep_chance(p[x], q[x]) for x in q), Fraction(0))


def residual(p: dict, q: dict) -> dict:
    """After a rejection the big model draws from the leftover odds: max(0, p - q), scaled to add up to 1."""
    assert p.keys() == q.keys(), "p and q must give odds for the same tokens"
    left = {x: max(Fraction(0), p[x] - q[x]) for x in p}
    total = sum(left.values())
    return {x: left[x] / total for x in left}


def output_distribution(p: dict, q: dict) -> dict:
    """The exact chance of each word coming out of one check: guessed and kept, or rejected and redrawn."""
    reject = 1 - keep_total(p, q)
    redraw = residual(p, q) if reject else {x: Fraction(0) for x in p}
    return {x: q[x] * keep_chance(p[x], q[x]) + reject * redraw[x] for x in p}


HAND_P = {"A": Fraction(1, 2), "B": Fraction(3, 10), "C": Fraction(1, 5)}       # the big model's odds
HAND_Q = {"A": Fraction(1, 5), "B": Fraction(3, 10), "C": Fraction(1, 2)}       # the small model's odds


# --- 3. tokens per big pass ---------------------------------------------------------------------

def expected_tokens(a: float, k: int) -> float:
    """Words per big pass if each guess is kept with chance a, independently of the others: (1 - a^(k+1)) / (1 - a), and k + 1 when a is 1.
    Leviathan et al. call the independence a simplifying assumption."""
    return k + 1.0 if a == 1 else (1 - a ** (k + 1)) / (1 - a)


def simulate_tokens_per_pass(a: float, k: int, rounds: int, seed: int) -> float:
    """Average words per big pass over seeded rounds: draw k guesses, each kept with chance a, count the run kept from the start, add one."""
    rng = random.Random(seed)
    total = 0
    for _ in range(rounds):
        kept = 0
        while kept < k and rng.random() < a:
            kept += 1
        total += kept + 1
    return total / rounds


# --- 4. time ------------------------------------------------------------------------------------

def draft_pass(users: int = 1):
    """One pass of the draft model over one new word for each of `users` users: 2 FLOPs per weight per word, every weight read once."""
    return big.Pass(flops=2 * DRAFT_PARAMS * users, nbytes=big.FP16_BYTES * DRAFT_PARAMS)


def plain_seconds(users: int = 1) -> float:
    """Plain decoding: one big pass per word (for `users` users at once, one word each)."""
    return big.prefill(users).seconds


def round_seconds(k: int, draft_seconds: float | None = None, users: int = 1) -> float:
    """One round: k draft passes, then one big pass over k + 1 inputs for each user."""
    each = draft_pass(users).seconds if draft_seconds is None else draft_seconds
    return k * each + big.prefill(users * (k + 1)).seconds


def speedup(a: float, k: int, draft_seconds: float | None = None, users: int = 1) -> float:
    """Words per second with the draft model divided by words per second without it: E * plain / round (Leviathan et al., Theorem 3.8)."""
    return expected_tokens(a, k) * plain_seconds(users) / round_seconds(k, draft_seconds, users)


def break_even(k: int, draft_seconds: float | None = None, users: int = 1) -> float:
    """The lowest keep chance at which speculative decoding is no slower than plain decoding, found by bisection."""
    low, high = 0.0, 1.0
    for _ in range(60):
        mid = (low + high) / 2
        low, high = (mid, high) if speedup(mid, k, draft_seconds, users) < 1 else (low, mid)
    return high


MAX_K = 32


def best_k(a: float, draft_seconds: float | None = None, users: int = 1) -> int:
    """The draft length that gives the largest speed-up at keep chance a."""
    return max(range(1, MAX_K + 1), key=lambda k: speedup(a, k, draft_seconds, users))


def users_when_checking_stops_being_free(k: int) -> int:
    """The smallest number of users whose big pass over users * (k + 1) inputs is limited by arithmetic instead of memory."""
    return next(users for users in range(1, 10_000) if big.prefill(users * (k + 1)).bound == "arithmetic")


def users_when_it_stops_paying(a: float, k: int) -> int:
    """The smallest number of users at which speculative decoding is slower than plain decoding."""
    return next(users for users in range(1, 10_000) if speedup(a, k, users=users) < 1)


# --- the widget's run ---------------------------------------------------------------------------

REPLY = greedy(big_choice, PROMPT, 17)                   # the reply both lanes write: the big model's own text, 17 tokens (here a token is a whole word)
DRAFT_STOPS_MS = [0.3, 0.6, draft_pass().seconds * 1000, 2.4, 3.9, 7.9]       # draft pass times the widget offers; the third is the 1B model at the chip's peak
DEFAULT_DRAFT_STOP = 2
OPENING = {"percent": 80, "k": 4, "seed": 60}            # the opening setting; the seed is pinned by test_opening_run_is_typical...: a typical run at 80% whose lanes change as the chance moves
M32 = 0xFFFFFFFF


def decoy(i: int) -> str:
    """A wrong guess for place i of the reply: what the small model would write there if it differs from the big model's word, else the
    big model's second choice, else the next word in alphabetical order. It only gives the widget a believable word to strike out."""
    context = PROMPT + REPLY[:i]
    others = [w for w in small_ranked(context) + big_ranked(context) if w != REPLY[i]]
    if others:
        return others[0]
    vocabulary = sorted(set(WORDS))
    return vocabulary[(vocabulary.index(REPLY[i]) + 1) % len(vocabulary)]


DECOYS = [decoy(i) for i in range(len(REPLY))]


def random_stream(seed: int):
    """mulberry32, a small seeded generator that JavaScript can repeat exactly: whole numbers from 0 to 2^32 - 1."""
    state = seed & M32

    def draw() -> int:
        nonlocal state
        state = (state + 0x6D2B79F5) & M32
        t = ((state ^ (state >> 15)) * (state | 1)) & M32
        t ^= (t + (((t ^ (t >> 7)) * (t | 61)) & M32)) & M32
        return (t ^ (t >> 14)) & M32

    return draw


def plan_run(percent: int, k: int, seed: int) -> list[list[int]]:
    """The rounds of one reply at keep chance percent/100. Each round draws k numbers, keeps the run that falls under the chance, and adds the big
    model's word. Returns [start, kept, rejected, own, words] per round: where it starts, guesses shown as kept, whether a wrong guess was struck out,
    whether the big model's own word was needed, and how many words it added. The last round stops at the end of the reply."""
    draw, start, rounds = random_stream(seed), 0, []
    while start < len(REPLY):
        left = len(REPLY) - start
        run = 0
        draws = [draw() for _ in range(k)]
        while run < k and draws[run] * 100 < percent * 2**32:
            run += 1
        words = min(run + 1, left)
        own = int(run + 1 <= left)
        rounds.append([start, min(run, words), int(own and run < k), own, words])
        start += words
    return rounds


def fmt(x: float, places: int) -> str:
    """Round half up on the exact value, as JavaScript's toFixed does (Python's format rounds ties to even)."""
    return format(Decimal(x).quantize(Decimal(1).scaleb(-places), rounding=ROUND_HALF_UP), "f")


def view(percent: int, k: int, draft_ms: float, seed: int) -> dict:
    """Everything the widget says for a setting: both lanes in words, the draft pass time, the two times it animates with, and the rounds."""
    n, rounds = len(REPLY), plan_run(percent, k, seed)
    plain_ms = plain_seconds() * 1000
    round_ms = k * draft_ms + big.prefill(k + 1).seconds * 1000
    per_token = len(rounds) * round_ms / n
    ratio = plain_ms / per_token
    average = expected_tokens(percent / 100, k) * plain_ms / round_ms
    shown = float(fmt(average, 1))                           # the verdict is read off the number printed below it, so the two cannot disagree
    verdict = ("On average, slower than plain decoding: the draft passes cost more than the kept guesses save." if shown < 1
               else "On average, about the same speed as plain decoding." if shown < 1.1 else "")
    return {
        "plain": f"{n} big passes, 1 token each, {fmt(plain_ms, 2)} ms per token.",
        "spec": f"This reply: {len(rounds)} big passes, {fmt(n / len(rounds), 1)} tokens each, {fmt(per_token, 2)} ms per token, {fmt(ratio, 1)} times as fast.",
        "verdict": verdict,
        "average": f"Averaged over many replies at this setting: {fmt(expected_tokens(percent / 100, k), 2)} tokens per big pass, {fmt(average, 1)} times as fast.",
        "draft": f"{fmt(draft_ms, 1)} ms ({fmt(draft_ms / plain_ms * 100, 0)}%)",
        "plainMs": plain_ms,
        "roundMs": round_ms,
        "rounds": rounds,
    }

"""Guards for speculative_sim.py: the draft model's size, the chip it takes from lesson 23, the toy pair and the proof that its text is unchanged,
the sampling rule, the tokens-per-pass formula, the time model, and the numbers the lesson prints."""
from fractions import Fraction
import random

import pytest

import speculative_sim as sim

big = sim.big
DRAFT_MS = sim.draft_pass().seconds * 1000
RING = sim.WORDS * 3


def test_numbers_taken_from_lesson_23_are_the_ones_this_lesson_prints():
    """Guards a change to lesson 23's model or chip that would leave this lesson's prose stale: the weights, the 7.877 ms step, the ridge and the 153 / 154 flip."""
    assert big.PARAMS == 8_030_261_248 and big.A100_PEAK_FLOPS == 312e12 and big.A100_BANDWIDTH == 2039e9 and big.FP16_BYTES == 2
    assert big.decode_step(1).seconds * 1000 == pytest.approx(7.877, abs=5e-4)
    assert big.ridge() == pytest.approx(153.016, abs=1e-3) and big.first_arithmetic_limited_prompt() == 154


def test_draft_model_count_matches_the_listed_total_and_shares_the_big_models_vocabulary():
    """Guards a wrong config value: counting a separate output table (a second 263 million weights) would give 1.50 billion and a draft pass 21% too slow.
    The shared vocabulary is what lets the big model check the draft's words one for one."""
    assert sim.draft_parameter_count() == sim.DRAFT_PARAMS == sim.D_WEIGHT_FILES_PARAMS == 1_235_814_400
    separate = sim.DRAFT_PARAMS + sim.D_VOCAB * sim.D_HIDDEN                                  # a second table of the same size
    assert round(separate / 1e9, 2) == 1.50 and round((separate / sim.DRAFT_PARAMS - 1) * 100) == 21
    assert sim.D_VOCAB == big.VOCAB == 128_256
    assert sim.DRAFT_PARAMS * big.FP16_BYTES / 1e9 == pytest.approx(2.47, abs=5e-3)


def test_both_toy_models_have_an_answer_after_every_run_of_words_in_the_text():
    """Guards a gap in the tables: the text is a ring, so no context may be missing, and ties must break the same way every time."""
    for i in range(len(sim.WORDS)):
        context = RING[i:i + 2]
        assert sim.big_choice(context) == sim.big_ranked(context)[0] and sim.small_choice(context) == sim.small_ranked(context)[0]
    assert sim.big_choice(["xx", "kettle"]) == sim.small_choice(["kettle"]) == "."          # a pair never seen falls back to the last word alone
    assert sim.REPLY == sim.WORDS[2:2 + len(sim.REPLY)] and len(sim.REPLY) == 17


@pytest.mark.parametrize("k", [1, 2, 3, 4, 6, 8])
def test_speculative_text_equals_the_big_models_own_text(k):
    """The claim the lesson rests on, run: for every start in the text, keeping guesses up to the first disagreement and using the big model's word there
    gives exactly the word-by-word text of the big model. A draft that is always wrong changes only the speed, never the text."""
    for i in range(len(sim.WORDS)):
        prompt = RING[i:i + 2]
        plain = sim.greedy(sim.big_choice, prompt, 50)
        text, rounds = sim.speculative_greedy(prompt, 50, k)
        assert text == plain, (k, prompt)
        for r in rounds:
            assert len(r.guesses) == k and len(r.own) == k + 1 and 1 <= len(r.new) <= k + 1
            assert r.new[:r.kept] == r.guesses[:r.kept] and r.new[-1] == r.own[r.kept]       # kept guesses, then the big model's word
            assert (r.kept == k) == (len(r.new) == k + 1)                                    # all kept: the extra word came free
            if r.kept < k:
                assert r.guesses[r.kept] != r.own[r.kept]                                    # stopped at a real disagreement
        always_wrong = lambda context: next(w for w in sim.WORDS if w != sim.big_choice(context))
        junk, junk_rounds = sim.speculative_greedy(prompt, 50, k, draft=always_wrong)
        assert junk == plain and all(len(r.new) == 1 for r in junk_rounds)


def test_toy_pair_agreement_and_its_words_per_pass_against_the_formula():
    """Guards the lesson's note that real guesses are not independent: the toy pair agrees 68% of the time, the formula for that
    chance gives 2.69 words per pass at k = 4 and the toy measures 2.92. The gap is why the formula is called an assumption."""
    assert sim.agreement() == Fraction(28, 41)
    measured, formula = sim.toy_tokens_per_pass(4), sim.expected_tokens(float(sim.agreement()), 4)
    assert (round(measured, 2), round(formula, 2)) == (2.92, 2.69) and measured > 1.05 * formula


def test_sampling_rule_gives_the_big_models_odds_exactly():
    """Guards the rule's two halves: keep with chance min(1, p/q), redraw from max(0, p - q) scaled. Drop the redraw or use q and the odds change.
    The hand example is the one in the lesson; the loop is the same check on many invented pairs, in exact fractions."""
    p, q = sim.HAND_P, sim.HAND_Q
    assert {x: q[x] * sim.keep_chance(p[x], q[x]) for x in p} == {"A": Fraction(1, 5), "B": Fraction(3, 10), "C": Fraction(1, 5)}
    assert sim.keep_total(p, q) == Fraction(7, 10) and sim.residual(p, q) == {"A": 1, "B": 0, "C": 0}
    assert sim.output_distribution(p, q) == p
    rng = random.Random(7)
    for _ in range(30):
        size = rng.randint(2, 6)
        weights = lambda: [rng.randint(0, 9) + (i == 0) for i in range(size)]               # at least one nonzero
        wp, wq = weights(), weights()
        P = {i: Fraction(w, sum(wp)) for i, w in enumerate(wp)}
        Q = {i: Fraction(w, sum(wq)) for i, w in enumerate(wq)}
        assert sim.output_distribution(P, Q) == P
        assert sim.output_distribution(P, P) == P and sim.keep_total(P, P) == 1
    for rule in (sim.keep_total, sim.residual):                                              # odds for different tokens are a caller's mistake, not a KeyError
        with pytest.raises(AssertionError, match="same tokens"):
            rule({"A": Fraction(1)}, {"B": Fraction(1)})


@pytest.mark.parametrize("a, k", [(0.3, 1), (0.6, 2), (0.6, 4), (0.8, 4), (0.9, 8), (0.5, 12)])
def test_simulated_words_per_pass_match_the_formula(a, k):
    """Guards the formula against the process it describes: seeded rounds, each guess kept with chance a, within 1.5%. Also that the closed form is the sum
    1 + a + ... + a^k (the chance that the word at each place is reached) and ends at 1 and k + 1."""
    assert sim.simulate_tokens_per_pass(a, k, 100_000, seed=k) == pytest.approx(sim.expected_tokens(a, k), rel=0.015)
    assert sim.expected_tokens(a, k) == pytest.approx(sum(a**j for j in range(k + 1)), rel=1e-12)
    assert sim.expected_tokens(0, k) == 1 and sim.expected_tokens(1, k) == k + 1


def test_time_model_numbers_the_lesson_prints():
    """Pins the time and speed-up figures in lesson.md and the figures, so a change to the model cannot leave stale prose."""
    plain = big.decode_step(1).seconds
    assert plain * 1000 == pytest.approx(7.877, abs=5e-4) and DRAFT_MS == pytest.approx(1.212, abs=5e-4)
    assert sim.draft_pass().seconds / plain == pytest.approx(0.1539, abs=5e-5) == pytest.approx(sim.DRAFT_PARAMS / big.PARAMS, abs=5e-5)
    assert all(big.prefill(n).seconds == plain for n in range(1, 154)) and big.prefill(154).seconds > plain      # checking up to 153 inputs costs one plain pass
    assert sim.round_seconds(4) * 1000 == pytest.approx(12.725, abs=5e-4)
    assert [round(sim.expected_tokens(a, 4), 2) for a in (0.6, 0.8, 0.9)] == [2.31, 3.36, 4.10]
    assert [round(sim.speedup(a, 4), 2) for a in (0.3, 0.6, 0.8, 0.9)] == [0.88, 1.43, 2.08, 2.53]
    assert round(sim.round_seconds(4) * 1000 / sim.expected_tokens(0.8, 4), 2) == 3.79                          # ms per word at the opening setting
    assert [round(sim.break_even(k), 3) for k in (2, 4, 8)] == [0.247, 0.386, 0.554]
    assert [sim.best_k(a) for a in (0.5, 0.6, 0.8, 0.9, 0.95)] == [2, 2, 5, 8, 12]
    assert [round(sim.speedup(a, sim.best_k(a)), 2) for a in (0.6, 0.8, 0.9)] == [1.50, 2.08, 2.75]
    assert sim.speedup(0.386, 4) < 1 < sim.speedup(0.387, 4)


def test_speed_up_falls_as_users_are_added_and_ends_below_one():
    """Guards the claim that spare room is what pays for guessing: with several users the check covers users * (k + 1) inputs, so it stays
    free only up to the ridge. At k = 4 that is 31 users; past 154 inputs per word the speed-up is flat at E / (k * 0.154 + k + 1)."""
    assert [sim.users_when_checking_stops_being_free(k) for k in (2, 4, 6, 8)] == [52, 31, 22, 18]      # 6: 22 * 7 = 154 exactly, a tie counts as arithmetic
    assert [round(sim.speedup(0.8, 4, users=u), 2) for u in (1, 30, 32, 64, 128)] == [2.08, 2.08, 2.02, 1.24, 0.70]
    assert sim.users_when_it_stops_paying(0.8, 4) == 85
    limit = sim.expected_tokens(0.8, 4) / (4 * sim.DRAFT_PARAMS / big.PARAMS + 5)
    assert sim.speedup(0.8, 4, users=1000) == pytest.approx(limit, rel=1e-9) and round(limit, 2) == 0.60
    speeds = [sim.speedup(0.8, 4, users=u) for u in range(1, 400)]
    assert all(later <= earlier + 1e-12 for earlier, later in zip(speeds, speeds[1:]))


def test_plan_run_covers_the_reply_once_and_follows_the_chance():
    """Guards the widget's run: every word written once, rounds that stop at the end, none kept at 0%, all kept at 100%."""
    n = len(sim.REPLY)
    for percent in (0, 20, 50, 80, 100):
        for k in (1, 4, 8):
            for seed in range(1, 6):
                rounds = sim.plan_run(percent, k, seed)
                assert sum(r[4] for r in rounds) == n and [r[0] for r in rounds] == [sum(x[4] for x in rounds[:i]) for i in range(len(rounds))]
                for start, kept, rejected, own, words in rounds:
                    assert 0 <= kept <= min(k, words) and 1 <= words <= k + 1 and (own == 1 or words == kept)
    assert all(r[1] == 0 for r in sim.plan_run(0, 4, 1)) and len(sim.plan_run(0, 4, 1)) == n
    assert len(sim.plan_run(100, 4, 1)) == -(-n // 5) and all(r[2] == 0 for r in sim.plan_run(100, 4, 1))


def test_opening_run_is_typical_and_shows_each_kind_of_round():
    """Guards the first screen: the opening setting (80%, k = 4, seed 60) must show a rejection, a round with every guess kept and its free word, and a
    speed-up near the average (2.1 here), so the reader does not meet a lucky or unlucky run first. It must also move: a seed whose rejection draws all
    sat above 0.95 left the lanes frozen from 80% to 95%, so the chance slider looked dead. Here the speed-up changes at least once on that stretch."""
    view = sim.view(sim.OPENING["percent"], sim.OPENING["k"], DRAFT_MS, sim.OPENING["seed"])
    rounds = view["rounds"]
    speed = lambda v: v["plainMs"] * len(sim.REPLY) / (len(v["rounds"]) * v["roundMs"])
    assert speed(view) == pytest.approx(sim.speedup(0.8, 4), rel=0.05)
    assert len({round(speed(sim.view(pct, 4, DRAFT_MS, sim.OPENING["seed"])), 6) for pct in (80, 85, 90, 95)}) > 1
    assert any(r[2] for r in rounds) and any(r[1] == sim.OPENING["k"] and r[3] for r in rounds) and any(r[1] == 0 for r in rounds)
    assert view["plain"] == "17 big passes, 1 token each, 7.88 ms per token."
    assert view["spec"] == "This reply: 5 big passes, 3.4 tokens each, 3.74 ms per token, 2.1 times as fast."
    assert view["average"] == "Averaged over many replies at this setting: 3.36 tokens per big pass, 2.1 times as fast." and view["verdict"] == ""


def test_view_says_so_when_guessing_loses_and_rounds_ties_up():
    """Guards the low-speed message and the rounding: half up as JavaScript's toFixed does (2.25 is 2.3, not Python's 2.2). The verdict follows the averaged
    speed-up printed below it, whatever the cause: with 95% of guesses kept and the draft pass as slow as the big one it is still slower, and it must not
    blame the guesses. With every guess kept the same slow draft is only 0.2% behind, which prints as 1.0 times, so the verdict says about the same."""
    assert sim.fmt(2.25, 1) == "2.3" and sim.fmt(0.125, 2) == "0.13" and sim.fmt(7.8766, 2) == "7.88" and sim.fmt(12.0, 1) == "12.0"
    low = sim.view(20, 4, DRAFT_MS, 1)
    assert low["verdict"].startswith("On average, slower than plain decoding") and low["spec"].endswith("times as fast.")
    slow_draft = sim.view(95, 4, sim.DRAFT_STOPS_MS[-1], 1)
    assert slow_draft["verdict"] == "On average, slower than plain decoding: the draft passes cost more than the kept guesses save." and "0.9 times as fast" in slow_draft["average"]
    assert sim.view(100, 4, sim.DRAFT_STOPS_MS[-1], 1)["verdict"] == "On average, about the same speed as plain decoding."
    assert sim.view(100, 8, DRAFT_MS, 1)["verdict"] == "" and sim.DRAFT_STOPS_MS[sim.DEFAULT_DRAFT_STOP] == DRAFT_MS

"""Guards for ../widget.html: its data block, its size, that its random run and sentences are the ones speculative_sim.py writes, and what a reader sees."""
import json
import re
import subprocess

import pytest

import speculative_sim as sim
from embed_widget import CONFIG_BLOCK, WIDGET, config, embed

HTML = WIDGET.read_text(encoding="utf-8")
ROOT = HTML[HTML.index('<div class="sdw-root">'):HTML.index("<script")]
SLOWER = "On average, slower than plain decoding: the draft passes cost more than the kept guesses save."
GRID = [(pct, k, d) for pct in range(0, 101, 5) for k in range(1, 9) for d in range(len(sim.DRAFT_STOPS_MS))]


def settings_view(pct, k, draft_index, seed):
    return sim.view(pct, k, sim.DRAFT_STOPS_MS[draft_index], seed)


# Reads the model out of the widget between its marker comments, runs it with Node, and compares the random numbers, the rounds and every
# sentence with the ones Python wrote into the fixture (the two times as numbers, the rest as exact text).
PARITY = r"""
import fs from 'node:fs';
const [fixture, widget] = 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="sdw-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 {stream, planRun, view, expectedTokens};')(cfg);
const bad = [];
let checks = 0;
const same = (a, b) => JSON.stringify(a) === JSON.stringify(b);
const near = (a, b) => Math.abs(a - b) <= 1e-12 * Math.abs(b);
fx.streams.forEach((s) => {
  const draw = api.stream(s.seed);
  const got = s.draws.map(() => draw());
  checks++;
  if (!same(got, s.draws)) bad.push(['stream', s.seed, got.slice(0, 3), s.draws.slice(0, 3)]);
});
fx.cases.forEach((c) => {
  const v = api.view(c.pct, c.k, cfg.draftStops[c.d], c.seed);
  checks += 3;
  if (!same(v.rounds, c.view.rounds)) bad.push(['rounds', c.pct, c.k, c.seed, v.rounds, c.view.rounds]);
  if (!['plain', 'spec', 'verdict', 'average', 'draft'].every((key) => v[key] === c.view[key])) bad.push(['text', c.pct, c.k, c.d, v, c.view]);
  if (!near(v.plainMs, c.view.plainMs) || !near(v.roundMs, c.view.roundMs)) bad.push(['times', c.pct, c.k, c.d, v.plainMs, v.roundMs]);
});
fx.expected.forEach(([a, k, e]) => { checks++; if (!near(api.expectedTokens(a, k), e)) bad.push(['expected', a, k, api.expectedTokens(a, k), e]); });
console.log(JSON.stringify({ checks, bad: bad.slice(0, 4) }));
process.exit(bad.length ? 1 : 0);
"""

# Runs the page script against a small stand-in for the DOM, with its timer and reduced-motion setting replaced, moves the sliders, presses the buttons
# and reads back what a reader sees: the words in each lane, the sentences, what is spoken, what Replay shows frame by frame.
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="sdw-config">([\s\S]*?)<\/script>/)[1];
const code = html.match(/<script>([\s\S]*?)<\/script>/)[1];

function el(tag) {
  const e = { tag, children: [], listeners: {}, value: '', hidden: false, attrs: {}, text: '' };
  Object.defineProperty(e, 'textContent', { get: () => e.text, set: (v) => { e.text = v; e.children.length = 0; } });
  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(reduced) {
  const page = { root: el('html'), messages: [], parent: { other: true }, timers: [], cleared: [] };
  fx.ids.forEach((id) => { page[id] = el('x'); });
  page['sdw-config'].textContent = config;
  const doc = { getElementById: (id) => page[id], createElement: el, documentElement: page.root };
  const win = { addEventListener: (type, f) => page.messages.push(f), matchMedia: (q) => ({ matches: reduced && q === '(prefers-reduced-motion: reduce)' }) };
  const setTimer = (f, ms) => (page.timers.push({ f, ms }), page.timers.length);
  new Function('document', 'window', 'parent', 'setInterval', 'clearInterval', code)(doc, win, page.parent, setTimer, (id) => page.cleared.push(id));
  return page;
}
const fire = (e, type) => (e.listeners[type || '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 guesses = (k) => k + (k === 1 ? ' guess' : ' guesses');
const speech = (pct, k, v) => pct + '% of guesses kept, ' + guesses(k) + ' per round. ' + v.spec + (v.verdict ? ' ' + v.verdict : '');

function passes(page) { return page['sdw-spec'].children; }
function look(page, v, what, pct, k) {              // compares the page with one expected view
  ['plain', 'spec', 'verdict', 'average'].forEach((key) => want(what + ': ' + key, page['sdw-' + (key === 'plain' ? 'plain-sum' : key === 'spec' ? 'spec-sum' : key)].textContent, v[key]));
  want(what + ': percent label', page['sdw-pct-n'].textContent, pct + '%');
  want(what + ': guesses label', page['sdw-k-n'].textContent, String(k));
  want(what + ': draft label', page['sdw-draft-n'].textContent, v.draft);
  want(what + ': slider readouts', [page['sdw-pct'].attrs['aria-valuetext'], page['sdw-k'].attrs['aria-valuetext'], page['sdw-draft'].attrs['aria-valuetext']],
       [pct + ' percent', guesses(k), v.draft + ' of a big pass']);
  const plain = page['sdw-plain'].children;
  want(what + ': plain lane words', plain.map((c) => [c.tag, c.attrs['data-kind'], c.textContent]), fx.reply.map((w) => ['strong', 'own', w]));
  want(what + ': pass numbers', passes(page).map((g) => g.attrs['data-pass']), v.rounds.map((r, j) => String(j + 1)));
  const expected = v.rounds.map((r) => {
    const chips = [];
    for (let i = 0; i < r[1]; i++) chips.push(['span', 'kept', fx.reply[r[0] + i]]);
    if (r[2]) chips.push(['del', 'rejected', fx.decoys[r[0] + r[1]]]);
    if (r[3]) chips.push(['strong', 'own', fx.reply[r[0] + r[1]]]);
    return chips;
  });
  want(what + ': passes', passes(page).map((g) => g.children.map((c) => [c.tag, c.attrs['data-kind'], c.textContent])), expected);
  const spelled = [];                                // the kept guesses and the big model's own words, in order, are the reply: nothing lost, nothing added
  passes(page).forEach((g) => g.children.forEach((c) => { if (c.attrs['data-kind'] !== 'rejected') spelled.push(c.textContent); }));
  want(what + ': the reply is spelled out', spelled, fx.reply);
  want(what + ': all shown', plain.concat(passes(page)).map((c) => c.attrs['data-wait']).filter((w) => w !== 'no'), []);
}

const first = open(false);
const opening = fx.opening;
look(first, opening.view, 'opening state', opening.pct, opening.k);
want('slider values', [first['sdw-pct'].value, first['sdw-k'].value, first['sdw-draft'].value], [String(opening.pct), String(opening.k), String(fx.defaultDraft)]);
want('nothing spoken yet', first['sdw-live'].textContent, '');
want('replay shown', first['sdw-replay'].hidden, false);
want('nothing animating at the start', first.timers.length, 0);

// every combination: dragging stays silent, releasing speaks
fx.combos.forEach((c) => {
  const page = open(false);
  page['sdw-pct'].value = String(c.pct); page['sdw-k'].value = String(c.k); page['sdw-draft'].value = String(c.d);
  fire(page['sdw-pct'], 'input');
  look(page, c.view, 'setting ' + [c.pct, c.k, c.d], c.pct, c.k);
  want('dragging stays silent at ' + [c.pct, c.k, c.d], page['sdw-live'].textContent, '');
  fire(page['sdw-k'], 'change');
  want('release speaks at ' + [c.pct, c.k, c.d], page['sdw-live'].textContent, speech(c.pct, c.k, c.view));
});

// New run draws the next seed and speaks
fx.newRuns.forEach((c) => {
  const page = open(false);
  page['sdw-pct'].value = String(c.pct); page['sdw-k'].value = String(c.k); page['sdw-draft'].value = String(c.d);
  fire(page['sdw-new']);
  look(page, c.view, 'new run at ' + [c.pct, c.k, c.d], c.pct, c.k);
  want('new run speaks', page['sdw-live'].textContent, speech(c.pct, c.k, c.view));
});

// Replay: everything hidden, then each pass appears when the clock passes its end; the last frame shows all and announces once
const replaying = open(false);
fire(replaying['sdw-replay']);
const lanes = () => replaying['sdw-plain'].children.map((c) => [c, (replaying['sdw-plain'].children.indexOf(c) + 1) * opening.view.plainMs])
  .concat(passes(replaying).map((g, j) => [g, (j + 1) * opening.view.roundMs]));
const end = Math.max(fx.reply.length * opening.view.plainMs, opening.view.rounds.length * opening.view.roundMs);
want('replay starts hidden', lanes().map(([c]) => c.attrs['data-wait']).filter((w) => w !== 'yes'), []);
want('one timer at 40 ms', replaying.timers.map((t) => t.ms), [40]);
const tick = replaying.timers[0].f;
let frame = 0;
[1, 25, 50, 75, 99].forEach((stop) => {
  while (frame < stop) { tick(); frame++; }
  const now = end * frame / 100;
  want('replay frame ' + frame, lanes().map(([c, at]) => c.attrs['data-wait'] === (at <= now ? 'no' : 'yes')).filter((ok) => !ok), []);
  want('silent while replaying at ' + frame, replaying['sdw-live'].textContent, '');
});
const shownAt = (stop) => {                          // the shorter lane is finished while the plain one is still writing
  const now = end * stop / 100;
  return [opening.view.rounds.length * opening.view.roundMs <= now, fx.reply.length * opening.view.plainMs <= now];
};
want('the draft lane finishes before the plain one', shownAt(75), [true, false]);
tick(); frame++;
want('last frame shows everything', lanes().map(([c]) => c.attrs['data-wait']).filter((w) => w !== 'no'), []);
want('last frame announces', replaying['sdw-live'].textContent, 'Replay finished. ' + speech(opening.pct, opening.k, opening.view));
want('timer stopped', replaying.cleared.includes(1), true);

// moving a slider during a replay stops it and shows the finished state
const interrupted = open(false);
fire(interrupted['sdw-replay']);
interrupted['sdw-pct'].value = '50';
fire(interrupted['sdw-pct'], 'input');
want('slider stops the replay', interrupted.cleared.includes(1), true);
want('slider shows the finished state', interrupted['sdw-plain'].children.concat(passes(interrupted)).map((c) => c.attrs['data-wait']).filter((w) => w !== 'no'), []);

// reduced motion: no Replay button, finished state, nothing animating
const calm = open(true);
want('reduced motion hides Replay', calm['sdw-replay'].hidden, true);
want('reduced motion starts nothing', calm.timers.length, 0);
look(calm, opening.view, 'reduced motion opening', opening.pct, opening.k);

// the host page theme: only the parent window may set it, and only to light or dark
const themed = open(false);
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({ combos: fx.combos.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 test_widget_data_block_matches_the_python_constants():
    """Guards a constant edit that was not copied into the widget (run embed_widget.py)."""
    assert embed(HTML) == HTML


def test_embed_rewrites_a_widget_whose_data_block_drifted():
    """Guards embed() skipping the block: a hand-edited parameter count or opening draft stop must be put back."""
    drifted = HTML.replace('"params": 8030261248', '"params": 7000000000', 1).replace('"defaultDraft": 2', '"defaultDraft": 0', 1)
    assert drifted != HTML and embed(drifted) == HTML
    assert json.loads(CONFIG_BLOCK.search(HTML).group(2)) == json.loads(json.dumps(config()))


def test_every_element_the_script_looks_up_exists_inside_the_root():
    """Guards a widget that never starts: the parity run covers only the model, so a misspelt id passed there and threw in the browser."""
    script = HTML[HTML.index("<script>"):]
    assert "$('sdw-' + name)" in script                   # the four sentences are looked up by suffix; they are added below
    lookups = set(re.findall(r"'(sdw-[\w-]+)'", script)) | {f"sdw-{n}" for n in ("plain-sum", "spec-sum", "verdict", "average")}
    missing = sorted(i for i in lookups if i != "sdw-config" and f'id="{i}"' not in ROOT)
    assert not missing, missing
    assert 'id="sdw-config"' in HTML


def test_widget_is_small_and_asks_for_nothing_online():
    assert len(HTML.splitlines()) <= 300 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


def test_widget_has_no_form_element():
    """Guards a dead button: the lesson page loads widget.html in a sandboxed frame, where a form cannot submit."""
    assert "<form" not in HTML and 'type="submit"' not in HTML


def test_widget_shows_at_most_five_controls_and_hides_the_rest():
    """Guards the first screen at 390 px: two sliders and Replay stay in view; the draft slider and New run sit 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)) <= 5
    more = re.search(r"<details><summary>More settings</summary>(.*?)</details>", ROOT, re.S).group(1)
    assert 'id="sdw-draft"' in more and 'id="sdw-new"' in more and "<details open" not in ROOT


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)", "--orange: light-dark(#b36200, #E69F00)", ":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", "#b36200", "#E69F00"}


def test_the_visible_notes_quote_what_the_code_computes():
    """Guards a stale sentence: the More settings note names the 1.24-billion-weight draft model, its 1.2 ms pass and the 7.88 ms big pass, and the
    left-out list says about 2.5 GB. Also the words that keep a best case from reading as a measurement, and the key that names each kind of chip."""
    view = settings_view(80, 4, sim.DEFAULT_DRAFT_STOP, sim.OPENING["seed"])
    assert f"{sim.DRAFT_PARAMS / 1e9:.2f}-billion-weight" in HTML and view["draft"].startswith("1.2 ms") and "The opening stop, 1.2 ms," in HTML
    assert view["plain"].endswith("7.88 ms per token.") and "takes 7.88 ms" in HTML and f"about {sim.DRAFT_PARAMS * 2 / 1e9:.1f} GB" in HTML
    outside = re.sub(r"<details>.*?</details>", "", ROOT, flags=re.S)
    assert "A model at the chip&rsquo;s peak speed, not a measurement." in outside                  # visible without opening anything
    leaves_out = re.search(r"<details><summary>What this toy leaves out</summary>(.*?)</details>", ROOT, re.S).group(1)
    for said in ("Real guesses are not independent chances", "its own memory", "usually runs further below the chip", "the check stops being free",
                 "Llama 3.1 8B (big) and Llama 3.2 1B (draft) at 16 bits on one A100"):
        assert said in leaves_out, said
    assert re.findall(r'<li><\w+ data-kind="(\w+)">token</\w+><span>(.*?)</span></li>', ROOT) == [
        ("kept", "Guess kept"), ("rejected", "Guess rejected"), ("own", "The big model&rsquo;s own")]
    assert 'class="sdw-sr" id="sdw-live" role="status" aria-live="polite" aria-atomic="true"' in ROOT      # what a screen reader hears
    assert ".sdw-root .sdw-btns button { flex: 1 1 auto; white-space: nowrap; }" in HTML      # a button wraps to the next row, never its label
    assert "[data-kind=rejected]" in HTML and "line-through" in HTML                                          # a rejected guess is struck, not only orange


def test_widget_model_gives_the_random_numbers_rounds_and_sentences_of_the_python_model(tmp_path):
    """Guards the widget's JavaScript drifting from speculative_sim.py: the lesson prints Python's numbers and the reader sees the widget's, so the
    random numbers, the rounds, every sentence and the two times must agree on every acceptance, guesses and draft setting the sliders offer."""
    cases = [{"pct": pct, "k": k, "d": d, "seed": 1 + i % 7, "view": settings_view(pct, k, d, 1 + i % 7)} for i, (pct, k, d) in enumerate(GRID)]
    seeds = (1, 2, 3, 12345, 2**32 - 1)
    streams = []
    for seed in seeds:
        draw = sim.random_stream(seed)
        streams.append({"seed": seed, "draws": [draw() for _ in range(50)]})
    expected = [[a, k, sim.expected_tokens(a, k)] for a in (0, 0.05, 0.5, 0.8, 0.95, 1) for k in (1, 4, 8)]
    result = run_node(tmp_path, PARITY, {"cases": cases, "streams": streams, "expected": expected}, lambda d: [str(d / "fixture.json"), str(WIDGET)])
    assert result["bad"] == [] and result["checks"] == 3 * len(cases) + len(streams) + len(expected)
    assert {c["view"]["verdict"] for c in cases} >= {"", "On average, about the same speed as plain decoding.", SLOWER}
    for c in cases:                                       # the verdict follows the averaged speed-up printed below it, not the one seeded run
        shown = float(re.search(r"([\d.]+) times as fast\.$", c["view"]["average"]).group(1))
        assert c["view"]["verdict"] == (SLOWER if shown < 1 else "On average, about the same speed as plain decoding." if shown == 1 else ""), c


def test_widget_page_shows_the_words_the_lesson_describes_and_replays_them(tmp_path):
    """Guards the part a reader sees: the opening state (80%, 4 guesses, 5 big passes against 17), the words in each lane spelling the reply, every kind
    of verdict, what is spoken on release and after New run, Replay frame by frame (the draft lane finishes first), a slider stopping it, reduced
    motion hiding it, and the theme message."""
    opening = settings_view(80, 4, sim.DEFAULT_DRAFT_STOP, sim.OPENING["seed"])
    assert sim.OPENING == {"percent": 80, "k": 4, "seed": 60} and opening["spec"].startswith("This reply: 5 big passes, 3.4 tokens each")
    assert len(opening["rounds"]) * opening["roundMs"] < len(sim.REPLY) * opening["plainMs"]
    verdicts = {}
    for pct, k, d in GRID:
        verdicts.setdefault(settings_view(pct, k, d, sim.OPENING["seed"])["verdict"], (pct, k, d))
    picks = [(pct, k, d) for pct in (0, 20, 50, 80, 100) for k in (1, 4, 8) for d in (0, 2, 5)] + list(verdicts.values())
    combos = [{"pct": pct, "k": k, "d": d, "view": settings_view(pct, k, d, sim.OPENING["seed"])} for pct, k, d in dict.fromkeys(picks)]
    new_runs = [{"pct": pct, "k": k, "d": 2, "view": settings_view(pct, k, 2, sim.OPENING["seed"] + 1)} for pct, k in ((80, 4), (20, 4), (100, 1))]
    ids = sorted(set(re.findall(r'id="(sdw-[\w-]+)"', HTML)))
    fixture = {"ids": ids, "reply": sim.REPLY, "decoys": sim.DECOYS, "defaultDraft": sim.DEFAULT_DRAFT_STOP, "combos": combos, "newRuns": new_runs,
               "opening": {"pct": 80, "k": 4, "view": opening}}
    result = run_node(tmp_path, PAGE, fixture, lambda d: [str(WIDGET), str(d / "fixture.json")])
    assert result == {"combos": len(combos), "bad": []}
    assert len(verdicts) == 3