Why does an AI wait a moment before the first word?

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

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


def config():
    """What the page needs to run the same formulas: the model, the chip, the timeline and the slider stops."""
    return {
        "params": sim.PARAMS, "layers": sim.LAYERS, "hidden": sim.HIDDEN, "bytesPerWeight": sim.FP16_BYTES, "kvBytes": sim.KV_BYTES_PER_TOKEN,
        "peak": sim.A100_PEAK_FLOPS, "bandwidth": sim.A100_BANDWIDTH, "axisMs": sim.AXIS_MS, "drawnSteps": sim.DRAWN_STEPS,
        "lengths": sim.prompt_lengths(), "presets": sim.PRESETS, "opening": sim.OPENING_TOKENS,
    }


def embed(html):
    """html with the data block rewritten from prefill_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 prefill_sim.py")

"""A roofline time model of prefill and decode for one request on one GPU.

The model is Llama 3.1 8B with 16-bit weights on an NVIDIA A100 80GB SXM, the chip of lesson 19
(tiled-matmul-roofline). A forward pass over n tokens does 2 FLOPs per parameter per token and reads every
weight once. Its time is at least the larger of two floors, the math floor (FLOPs / peak) and the memory
floor (bytes / bandwidth), exactly as in lesson 19. Both floors assume the chip runs at its data-sheet peak,
so a real server is slower: the numbers here are best cases, not measurements.

    prefill: one pass over all n prompt tokens         time = max(2 P n / peak, 2 P / bandwidth)
    decode:  one pass over the one newest token        time = max(2 P / peak,   2 P / bandwidth) = 2 P / bandwidth

With extras=True the model adds what the short formulas leave out: the attention scores between prompt
tokens (which grow with n squared) and the cache of earlier tokens that attention reads. It still leaves out
network time, queueing behind other users and the gap between peak and real speed.

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

from dataclasses import dataclass

# --- the chip: copied from lesson 19 (test_prefill_sim.py checks the copy against ../../tiled-matmul-roofline/code/tiled_roofline.py)
# NVIDIA A100 Tensor Core GPU data sheet (revision "Jun21"), 80GB SXM column: 312 TFLOPS dense BFLOAT16/FP16, 2,039 GB/s. GB is 10**9 bytes.
A100_PEAK_FLOPS = 312e12
A100_BANDWIDTH = 2039e9
FP16_BYTES = 2

# --- the model: Llama 3.1 8B. Meta's model card (github.com/meta-llama/llama-models, models/llama3_1/MODEL_CARD.md) gives 8B parameters, grouped-query
# attention and a 128k context; the shape below is Meta's models/sku_list.py (dim 4096, 32 layers, 32 heads, 8 key-value heads, vocab 128256; the feed-forward
# width 14336 is its multiplier 1.3 rounded up to a multiple of 1024). The numbers were typed from config.json of huggingface.co/NousResearch/Meta-Llama-3.1-8B,
# a third-party copy whose values match. Read 2026-10-19.
HIDDEN = 4096
FFN = 14336
LAYERS = 32
HEADS = 32
KV_HEADS = 8
VOCAB = 128256
HEAD_DIM = HIDDEN // HEADS
# The card rounds the size to 8B. The exact total, to check the count below against, is what Hugging Face lists for that copy's weight files (safetensors metadata, BF16).
WEIGHT_FILES_PARAMS = 8_030_261_248


def parameter_count() -> int:
    """Count the weights from the config: attention, feed-forward and norm weights of each layer, plus the input and output tables."""
    attention = 2 * HIDDEN * HIDDEN + 2 * HIDDEN * KV_HEADS * HEAD_DIM   # query and output matrices, then key and value matrices
    feed_forward = 3 * HIDDEN * FFN                                       # gate, up and down matrices
    layer = attention + feed_forward + 2 * HIDDEN                         # plus two norm weight vectors
    return LAYERS * layer + 2 * VOCAB * HIDDEN + HIDDEN                   # untied input and output tables, final norm


PARAMS = parameter_count()
EMBED_PARAMS = VOCAB * HIDDEN
KV_BYTES_PER_TOKEN = 2 * LAYERS * KV_HEADS * HEAD_DIM * FP16_BYTES        # keys and values, every layer, 16-bit


# --- one pass -----------------------------------------------------------------------------------

def ridge() -> float:
    """FLOP per byte where the chip's memory slope meets its flat compute roof (lesson 19: 153)."""
    return A100_PEAK_FLOPS / A100_BANDWIDTH


@dataclass(frozen=True)
class Pass:
    flops: int
    nbytes: int

    @property
    def math_s(self) -> float:
        return self.flops / A100_PEAK_FLOPS

    @property
    def memory_s(self) -> float:
        return self.nbytes / A100_BANDWIDTH

    @property
    def seconds(self) -> float:
        """At least this long: the pass cannot beat either floor, and the two overlap."""
        return max(self.math_s, self.memory_s)

    @property
    def intensity(self) -> float:
        return self.flops / self.nbytes

    @property
    def bound(self) -> str:
        """A tie at the ridge counts as arithmetic-limited (compute-bound), as in lesson 19."""
        return "arithmetic" if self.intensity >= ridge() else "memory"

    @property
    def busy(self) -> float:
        """Share of the chip's peak arithmetic this pass reaches."""
        return self.math_s / self.seconds


def attention_flops_prefill(n: int) -> int:
    """Scores and weighted sums between the n prompt tokens, counting each token against itself and the ones before it only."""
    return 4 * LAYERS * HIDDEN * (n * (n + 1) // 2)


def prefill(n: int, extras: bool = False) -> Pass:
    """One pass over n prompt tokens: 2 FLOPs per parameter per token, every weight read once."""
    flops, nbytes = 2 * PARAMS * n, FP16_BYTES * PARAMS
    if extras:
        flops += attention_flops_prefill(n)
        nbytes += n * KV_BYTES_PER_TOKEN                                  # the cache this pass writes
    return Pass(flops, nbytes)


def decode_step(context: int, extras: bool = False) -> Pass:
    """One pass over the newest token with `context` earlier tokens in the cache."""
    flops, nbytes = 2 * PARAMS, FP16_BYTES * PARAMS
    if extras:
        flops += 4 * LAYERS * HIDDEN * context                            # the new token against every earlier one
        nbytes += context * KV_BYTES_PER_TOKEN                            # the whole cache is read
    return Pass(flops, nbytes)


def first_arithmetic_limited_prompt() -> int:
    """The smallest prompt whose prefill is limited by arithmetic: the whole-model intensity is n FLOP per byte."""
    for n in range(1, MAX_PROMPT + 1):
        if prefill(n).bound == "arithmetic":
            return n
    raise ValueError(f"no prompt up to {MAX_PROMPT} tokens is limited by arithmetic")


def grid_intensity(m: int, side: int = 4096) -> float:
    """FLOP per byte of one side x side weight grid times m tokens, counting the tokens going in and the results coming out as lesson 19 does.
    The whole-model count above leaves those out, which is why lesson 19 reaches the ridge at 166 tokens and this file at 154."""
    return 2 * m * side * side / (FP16_BYTES * (2 * m * side + side * side))


def attention_share_of_linear(n: int) -> float:
    """Attention FLOPs of a prompt of n tokens as a share of the weight FLOPs (extras on)."""
    return attention_flops_prefill(n) / (2 * PARAMS * n)


# --- the widget's view ---------------------------------------------------------------------------

MIN_PROMPT, MAX_PROMPT = 10, 8000
DRAWN_STEPS = 10         # decode steps drawn after the prefill block
AXIS_MS = 600            # fixed length of the timeline, so a longer prompt visibly gives a longer block
PRESETS = [{"name": "A question", "tokens": 20}, {"name": "An article", "tokens": 1000}, {"name": "A long report", "tokens": 8000}]
OPENING_TOKENS = 8000


def prompt_lengths() -> list[int]:
    """Slider stops: 10 to 8,000 tokens, equal steps on a log scale, whole numbers below 100 and two significant figures above,
    plus the two prompts either side of the ridge (153 and 154) so the reader can see the block start to grow."""
    stops: list[int] = []
    for pos in range(101):
        x = MIN_PROMPT * (MAX_PROMPT / MIN_PROMPT) ** (pos / 100)
        n = round(x) if x < 100 else int(float(f"{x:.2g}"))
        if not stops or n != stops[-1]:
            stops.append(n)
    first = first_arithmetic_limited_prompt()
    return sorted(set(stops) | {first - 1, first})


def commas(n: int) -> str:
    return f"{n:,}"


def fmt_ms(seconds: float) -> str:
    ms = seconds * 1000
    return f"{ms:.2f}" if ms < 10 else f"{ms:.1f}" if ms < 100 else f"{ms:.0f}"


def view(n: int, extras: bool = False) -> dict:
    """Everything the widget shows for a prompt of n tokens: the three timeline sizes and the four lines of words."""
    pre, step = prefill(n, extras), decode_step(n, extras)
    steps = int(pre.seconds / step.seconds + 0.5)
    ratio = "about as long as one decode step" if steps <= 1 else f"about as long as {commas(steps)} decode steps"
    return {
        "bound": pre.bound,
        "prefillWidth": f"{pre.seconds * 1000 / AXIS_MS * 100:.3f}%",
        "stepWidth": f"{step.seconds * 1000 / AXIS_MS * 100:.3f}%",
        "ttft": f"Best-case time to the first token: {fmt_ms(pre.seconds)} ms",
        "tpot": f"Then one token every {fmt_ms(step.seconds)} ms, at most {1 / step.seconds:.0f} tokens a second.",
        "why": f"Reading {commas(n)} tokens is limited by {pre.bound}: it takes {ratio}.",
        "busy": f"At best, the chip is busy with arithmetic {int(pre.busy * 100)}% of the time while reading the prompt and {step.busy * 100:.1f}% while writing.",
    }

"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
from prefill_sim import (
    A100_BANDWIDTH, A100_PEAK_FLOPS, AXIS_MS, DRAWN_STEPS, EMBED_PARAMS, FP16_BYTES, HEAD_DIM, HIDDEN, KV_BYTES_PER_TOKEN,
    LAYERS, PARAMS, WEIGHT_FILES_PARAMS, attention_flops_prefill, grid_intensity, attention_share_of_linear, decode_step, first_arithmetic_limited_prompt, fmt_ms,
    prefill, prompt_lengths, ridge, view,
)

GIB = 2**30


def main() -> None:
    step = decode_step(1)
    print(f"parameters from the config {PARAMS:,}; listed for the weight files {WEIGHT_FILES_PARAMS:,} (Meta's card rounds it to 8B)")
    print(f"weights in 16 bit: {FP16_BYTES * PARAMS / 1e9:.2f} GB; chip: {A100_PEAK_FLOPS / 1e12:.0f} TFLOP/s, {A100_BANDWIDTH / 1e9:.0f} GB/s")
    print(f"ridge {ridge():.2f} FLOP per byte; first prompt limited by arithmetic: {first_arithmetic_limited_prompt()} tokens")
    print(f"decode step {step.seconds * 1000:.3f} ms = {1 / step.seconds:.1f} tokens/s at most; arithmetic busy {step.busy * 100:.2f}% (1/ridge = {100 / ridge():.2f}%)")
    print(f"FLOPs per token 2P = {2 * PARAMS:,} ({2 * PARAMS / 1e9:.2f} GFLOP)")
    print()
    print("tokens | math floor ms | memory floor ms | time ms | limited by | busy | steps-long")
    for n in (1, 10, 20, 100, 150, 153, 154, 160, 300, 1000, 3000, 8000):
        p = prefill(n)
        print(f"{n:6,} | {p.math_s * 1000:9.3f} | {p.memory_s * 1000:9.3f} | {p.seconds * 1000:8.3f} | {p.bound:10} | {p.busy * 100:5.1f}% | {p.seconds / step.seconds:6.2f}")
    print()
    print("widget lines at the opening state (8,000 tokens) and at 20 tokens:")
    for n in (8000, 20):
        print(view(n))
    print()
    print(f"timeline: axis {AXIS_MS} ms; longest case = prefill {fmt_ms(prefill(8000, True).seconds)} ms + {DRAWN_STEPS} steps of {fmt_ms(decode_step(8000, True).seconds)} ms"
          f" = {(prefill(8000, True).seconds + DRAWN_STEPS * decode_step(8000, True).seconds) * 1000:.1f} ms")
    print(f"slider stops {len(prompt_lengths())}: {prompt_lengths()[:6]} ... {prompt_lengths()[-3:]}")
    print()
    print("left out of the short formula:")
    for n in (1000, 8000):
        a, b = prefill(n), prefill(n, True)
        print(f"  n={n:,}: attention FLOPs (each token against itself and earlier ones) {attention_flops_prefill(n):.3e}"
              f" = {100 * attention_share_of_linear(n):.1f}% of the weight FLOPs; prefill {a.seconds * 1000:.1f} -> {b.seconds * 1000:.1f} ms;"
              f" decode step {step.seconds * 1000:.3f} -> {decode_step(n, True).seconds * 1000:.3f} ms")
    print(f"  all pairs counted instead of the causal half: attention at 8,000 would be {4 * LAYERS * HIDDEN * 8000**2:.3e} FLOPs")
    n5 = next(n for n in range(1, 100000) if attention_share_of_linear(n) >= 0.05)
    n10 = next(n for n in range(1, 100000) if attention_share_of_linear(n) >= 0.10)
    print(f"  attention reaches 5% of the weight FLOPs at {n5:,} tokens and 10% at {n10:,}")
    print(f"  KV cache: {KV_BYTES_PER_TOKEN:,} bytes per token = 2 x {LAYERS} layers x 8 KV heads x {HEAD_DIM} x {FP16_BYTES} bytes")
    for n in (1000, 8000, 8192):
        print(f"  cache at {n:,} tokens {n * KV_BYTES_PER_TOKEN / 1e9:.3f} GB = {n * KV_BYTES_PER_TOKEN / GIB:.3f} GiB = {100 * n * KV_BYTES_PER_TOKEN / (FP16_BYTES * PARAMS):.1f}% of the weight bytes")
    first_grid = next(m for m in range(1, 1000) if grid_intensity(m) >= ridge())
    print(f"  lesson 19's single grid (4096 x 4096, activations counted): first prompt limited by arithmetic {first_grid}; activation bytes {100 * 2 * first_grid * 4096 * 2 / (2 * 4096 ** 2):.1f}% of the weight bytes there")
    print(f"  input table {EMBED_PARAMS:,} weights = {100 * EMBED_PARAMS / PARAMS:.1f}% of all weights (a lookup reads one row, the short formula counts all)")


if __name__ == "__main__":
    main()

"""Guards for prefill_sim.py: the model's size, the chip it copies from lesson 19, the ridge, the two times, and the numbers the lesson prints."""
import re
from pathlib import Path

import pytest

import prefill_sim as sim

LESSON_19 = Path(__file__).resolve().parents[2] / "tiled-matmul-roofline" / "code" / "tiled_roofline.py"
GIB = 2**30


def test_parameter_count_equals_the_total_listed_for_the_weight_files():
    """Guards a wrong config value or layer formula: tying the output table to the input table would give 7.50 billion, and
    key and value matrices as wide as the query matrix (no grouped heads) would give 8.8 billion. Either changes every time here."""
    assert sim.parameter_count() == sim.WEIGHT_FILES_PARAMS == 8_030_261_248
    assert sim.PARAMS - sim.EMBED_PARAMS == 7_504_924_672


def test_chip_constants_equal_lesson_19s(tmp_path):
    """Guards a copy that drifted: the lesson quotes lesson 19's A100 numbers, so the two files must hold the same ones."""
    source = LESSON_19.read_text(encoding="utf-8")
    for name in ("A100_PEAK_FLOPS", "A100_BANDWIDTH", "FP16_BYTES"):
        theirs = re.search(rf"^{name} = (\S+)$", source, re.M)
        assert theirs, f"{name} not found in lesson 19: update this test"
        assert float(theirs.group(1)) == getattr(sim, name), name


def test_ridge_is_153_and_the_first_prompt_limited_by_arithmetic_is_154():
    """Guards an off-by-one at the ridge: 153.02 is not a whole number, so 153 tokens are still limited by memory and 154 are not."""
    assert sim.ridge() == pytest.approx(153.016, abs=1e-3)
    assert sim.first_arithmetic_limited_prompt() == 154
    assert sim.prefill(153).bound == "memory" and sim.prefill(154).bound == "arithmetic"


def test_a_tie_at_the_ridge_counts_as_arithmetic_limited():
    """Guards the rule lesson 19 states (a tie is compute-bound) against a `>` that would reverse it."""
    exact = sim.Pass(flops=int(sim.A100_PEAK_FLOPS), nbytes=int(sim.A100_BANDWIDTH))
    assert exact.math_s == exact.memory_s and exact.bound == "arithmetic"


def test_lesson_19s_166_is_one_grid_with_its_inputs_and_outputs_counted():
    """Guards the lesson's note on why 166 (lesson 19, one 4096 by 4096 grid) and 154 (this lesson, whole model, weight bytes only) differ:
    adding the activation traffic to the bytes moves the first arithmetic-limited prompt from 154 to 166. Lesson 19 prints 153.55 at 166."""
    first = next(m for m in range(1, 1000) if sim.grid_intensity(m) >= sim.ridge())
    assert first == 166 and round(sim.grid_intensity(166), 2) == 153.55 and sim.grid_intensity(165) < sim.ridge()
    assert round(100 * 2 * 166 * 4096 * 2 / (2 * 4096**2), 1) == 8.1                  # activation bytes as a share of the weight bytes at 166


@pytest.mark.parametrize("n", [1, 7, 100, 153, 154, 1000, 8000])
def test_prefill_intensity_is_the_prompt_length(n):
    """The lesson's key claim, computed: 2 P n FLOPs over 2 P bytes is n FLOP per byte. Guards a missing factor of 2 on either side."""
    assert sim.prefill(n).intensity == n


def test_prefill_takes_one_decode_step_below_the_ridge_and_grows_in_line_above_it():
    """Guards the shape the widget draws: flat at the decode step up to 153 tokens, then proportional to the prompt."""
    step = sim.decode_step(1).seconds
    assert all(sim.prefill(n).seconds == step for n in range(1, 154))
    one_token = 2 * sim.PARAMS / sim.A100_PEAK_FLOPS
    for n in (154, 300, 1000, 8000):
        assert sim.prefill(n).seconds == pytest.approx(n * one_token, rel=1e-12)
    assert sim.prefill(2000).seconds == pytest.approx(2 * sim.prefill(1000).seconds, rel=1e-12)


@pytest.mark.parametrize("extras", [False, True])
@pytest.mark.parametrize("context", [1, 100, 8000, 100_000])
def test_decode_step_is_limited_by_memory_for_every_context(context, extras):
    """Guards the lesson's second claim: one user's decode step never reaches the ridge, with or without the attention cache."""
    step = sim.decode_step(context, extras)
    assert step.bound == "memory" and step.seconds == step.memory_s and step.busy < 0.02


def test_decode_step_matches_lesson_19s_weight_read_time():
    """Guards the decode step drifting from lesson 19's formula, bytes of weights over bandwidth (6.87 ms for 7 billion parameters there)."""
    assert sim.decode_step(1).seconds == sim.PARAMS * sim.FP16_BYTES / sim.A100_BANDWIDTH
    assert 7e9 * 2 / sim.A100_BANDWIDTH * 1000 == pytest.approx(6.866, abs=1e-3)


def test_numbers_the_lesson_prints():
    """Pins the figures in lesson.md and the figures' text, so a change to the model cannot leave stale prose."""
    ms = lambda s: s * 1000
    assert ms(sim.decode_step(1).seconds) == pytest.approx(7.877, abs=5e-4) and 1 / sim.decode_step(1).seconds == pytest.approx(126.96, abs=0.01)
    assert sim.decode_step(1).busy == pytest.approx(1 / sim.ridge(), rel=1e-12)
    assert [round(ms(sim.prefill(n).math_s), 2) for n in (100, 1000)] == [5.15, 51.48]
    assert [round(ms(sim.prefill(n).seconds), 1) for n in (20, 100, 154, 1000, 3000, 8000)] == [7.9, 7.9, 7.9, 51.5, 154.4, 411.8]
    assert round(sim.prefill(100).busy * 100, 1) == 65.4
    assert 2 * sim.PARAMS == 16_060_522_496 and sim.FP16_BYTES * sim.PARAMS / 1e9 == pytest.approx(16.06, abs=5e-3)
    assert sim.EMBED_PARAMS == 525_336_576 and round(100 * sim.EMBED_PARAMS / sim.PARAMS, 1) == 6.5


def test_cache_bytes_per_token_agree_with_lesson_22():
    """Guards the cache size: lesson 22 gives 128 KiB per token and 1 GiB at 8,192 tokens for this model (8 key-value heads, not 32)."""
    assert sim.KV_BYTES_PER_TOKEN == 131_072 == 128 * 1024
    assert 8192 * sim.KV_BYTES_PER_TOKEN == GIB
    assert round(8000 * sim.KV_BYTES_PER_TOKEN / 1e9, 3) == 1.049


@pytest.mark.parametrize("n", [1, 2, 3, 10, 50])
def test_attention_flops_count_each_token_against_itself_and_the_ones_before(n):
    """Guards n * (n + 1) / 2 against n * (n - 1) / 2 or a full n * n: count the pairs one by one."""
    pairs = sum(1 for i in range(n) for j in range(i + 1))
    assert sim.attention_flops_prefill(n) == 4 * sim.LAYERS * sim.HIDDEN * pairs


def test_attention_numbers_the_lesson_prints():
    n5 = next(n for n in range(1, 100_000) if sim.attention_share_of_linear(n) >= 0.05)
    n10 = next(n for n in range(1, 100_000) if sim.attention_share_of_linear(n) >= 0.10)
    assert (n5, n10) == (3063, 6126)
    assert round(100 * sim.attention_share_of_linear(1000), 1) == 1.6 and round(100 * sim.attention_share_of_linear(8000), 1) == 13.1
    assert f"{sim.attention_flops_prefill(8000):.3e}" == "1.678e+13"
    assert [round(sim.prefill(8000, True).seconds * 1000, 1), round(sim.decode_step(8000, True).seconds * 1000, 2)] == [465.6, 8.39]
    assert round(100 * 8000 * sim.KV_BYTES_PER_TOKEN / (sim.FP16_BYTES * sim.PARAMS), 1) == 6.5


def test_extras_never_make_a_pass_faster():
    """Guards a sign error in the extras: they add work and traffic, so each time is at least the short formula's."""
    for n in sim.prompt_lengths():
        assert sim.prefill(n, True).seconds >= sim.prefill(n).seconds and sim.decode_step(n, True).seconds >= sim.decode_step(n).seconds, f"stop {n}"


def test_slider_stops_cover_10_to_8000_with_the_presets_and_the_ridge_neighbours():
    """Guards a slider that skips a preset or the flip: the presets must be stops, and 153 (memory) and 154 (arithmetic) must both be stops."""
    stops = sim.prompt_lengths()
    assert stops[0] == 10 and stops[-1] == 8000 and stops == sorted(set(stops)) and len(stops) >= 90
    assert all(p["tokens"] in stops for p in sim.PRESETS) and sim.OPENING_TOKENS in stops
    assert sum(1 for n in stops if 100 <= n < 154) >= 3 and sum(1 for n in stops if 154 <= n <= 300) >= 5
    assert {153, 154} <= set(stops) and sim.prefill(153).bound == "memory" and sim.prefill(154).bound == "arithmetic"


def test_timeline_axis_holds_the_longest_case():
    """Guards a timeline that runs off its bar: the longest prefill plus the drawn decode steps, extras on, must fit."""
    longest = sim.prefill(sim.MAX_PROMPT, True).seconds + sim.DRAWN_STEPS * sim.decode_step(sim.MAX_PROMPT, True).seconds
    assert longest * 1000 < sim.AXIS_MS


def test_view_for_the_two_states_the_lesson_names():
    long_ = sim.view(8000)
    assert long_["ttft"] == "Best-case time to the first token: 412 ms"
    assert long_["tpot"] == "Then one token every 7.88 ms, at most 127 tokens a second."
    assert long_["why"] == "Reading 8,000 tokens is limited by arithmetic: it takes about as long as 52 decode steps."
    assert long_["busy"] == "At best, the chip is busy with arithmetic 100% of the time while reading the prompt and 0.7% while writing." and long_["bound"] == "arithmetic"
    short = sim.view(20)
    assert short["ttft"] == "Best-case time to the first token: 7.88 ms"
    assert short["why"] == "Reading 20 tokens is limited by memory: it takes about as long as one decode step."
    assert short["busy"] == "At best, the chip is busy with arithmetic 13% of the time while reading the prompt and 0.7% while writing." and short["bound"] == "memory"
    assert sim.view(153)["busy"].startswith("At best, the chip is busy with arithmetic 99% ")      # memory-bound by a hair: never shown as 100%


def test_view_draws_the_prefill_block_as_one_step_below_the_ridge_and_wider_above():
    """Guards the picture the widget exists for: at or below 153 tokens the prefill block is exactly one decode step wide."""
    width = lambda text: float(text.rstrip("%"))
    for n in (10, 20, 100, 150, 153):
        v = sim.view(n)
        assert v["prefillWidth"] == v["stepWidth"], n
    assert width(sim.view(1000)["prefillWidth"]) > 6 * width(sim.view(1000)["stepWidth"])
    widths = [width(sim.view(n)["prefillWidth"]) for n in sim.prompt_lengths()]
    assert widths == sorted(widths) and widths[-1] > 50 * widths[0]

"""Guards for ../widget.html: its data block, its size, that its formulas give the same times and words as prefill_sim.py, and what a reader sees."""
import json
import re
import subprocess

import pytest

import prefill_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="pfw-root">'):HTML.index("<script")]
STOPS = sim.prompt_lengths()
EXTRA_LENGTHS = [1, 2, 153, 154, 155, 5000, 6126, 100_000]       # beyond the slider, where a formula slip would show

# Reads the model out of the widget between its marker comments, runs it with Node, and compares every time and every line of words
# with the ones Python wrote into the fixture, bit for bit (the seconds are compared as exact numbers).
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="pfw-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 {pass: pass, prefill, decodeStep, view, commas, fmtMs};')(cfg);
const bad = [];
let checks = 0;
const same = (a, b) => JSON.stringify(a) === JSON.stringify(b);
fx.cases.forEach((c) => {
  const p = api.prefill(c.n, c.extras), d = api.decodeStep(c.n, c.extras);
  checks += 3;
  if (p.seconds !== c.prefill.seconds || p.bound !== c.prefill.bound || p.busy !== c.prefill.busy) bad.push(['prefill', c.n, c.extras, p, c.prefill]);
  if (d.seconds !== c.decode.seconds || d.bound !== c.decode.bound || d.busy !== c.decode.busy) bad.push(['decode', c.n, c.extras, d, c.decode]);
  if (!same(api.view(c.n, c.extras), c.view)) bad.push(['view', c.n, c.extras, api.view(c.n, c.extras), c.view]);
});
checks++;                                          // a tie at the ridge counts as arithmetic, as in prefill_sim.py
if (api.pass(cfg.peak, cfg.bandwidth).bound !== 'arithmetic') bad.push(['tie', api.pass(cfg.peak, cfg.bandwidth)]);
fx.commas.forEach(([n, text]) => { checks++; if (api.commas(n) !== text) bad.push(['commas', n, api.commas(n), text]); });
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, moves the slider, presses the buttons and the checkbox,
# 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="pfw-config">([\s\S]*?)<\/script>/)[1];
const code = html.match(/<script>([\s\S]*?)<\/script>/)[1];

function el(tag) {
  const e = { tag, children: [], listeners: {}, value: '', checked: false, 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['pfw-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, 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]); };
function look(page, v, what, n) {                    // compares the page with one expected view
  want(what + ': length', page['pfw-n'].textContent, n);
  ['ttft', 'tpot', 'why', 'busy'].forEach((k) => want(what + ': ' + k, page['pfw-' + k].textContent, v[k]));
  const blocks = page['pfw-bar'].children;
  want(what + ': block count', blocks.length, fx.drawnSteps + 1);
  want(what + ': prefill width', blocks[0].style.width, v.prefillWidth);
  want(what + ': prefill colour', blocks[0].attrs['data-bound'], v.bound);
  want(what + ': first word mark', page['pfw-mark'].style.left, v.prefillWidth);
  want(what + ': step widths', blocks.slice(1).map((b) => b.style.width), blocks.slice(1).map(() => v.stepWidth));
  want(what + ': step colours', blocks.slice(1).map((b) => b.attrs['data-bound']), blocks.slice(1).map(() => 'memory'));
}

const first = open();
look(first, fx.opening, 'opening state', fx.openingText);
want('slider range', [first['pfw-len'].max, first['pfw-len'].value], [String(fx.stops.length - 1), String(fx.stops.indexOf(fx.openingTokens))]);
want('axis labels', [first['pfw-mid'].textContent, first['pfw-end'].textContent], fx.axis);
want('nothing spoken yet', first['pfw-live'].textContent, '');
const buttons = first['pfw-btns'].children;
want('preset labels', buttons.map((b) => b.textContent), fx.presetLabels);
want('preset pressed at the opening', buttons.map((b) => b.attrs['aria-pressed']), fx.presets.map((p) => String(p.tokens === fx.openingTokens)));

// every slider stop, with the checkbox off and on, by dragging (input event) and by releasing (change event)
[false, true].forEach((extras) => {
  const page = open();
  page['pfw-extra'].checked = extras;
  fx.stops.forEach((n, i) => {
    page['pfw-len'].value = String(i);
    page['pfw-live'].textContent = '';
    fire(page['pfw-len'], 'input');
    const v = fx.views[String(extras)][String(n)];
    look(page, v, 'stop ' + n + ' extras ' + extras, fx.texts[String(n)]);
    want('dragging stays silent at ' + n, page['pfw-live'].textContent, '');
    want('slider reads out', page['pfw-len'].attrs['aria-valuetext'], fx.texts[String(n)] + ' tokens');
    fire(page['pfw-len'], 'change');
    want('release speaks at ' + n, page['pfw-live'].textContent, fx.texts[String(n)] + ' tokens. ' + v.ttft + '. ' + v.tpot);
  });
});

// the three presets move the slider and mark themselves pressed
const press = open();
fx.presets.forEach((p, k) => {
  fire(press['pfw-btns'].children[k]);
  const n = p.tokens;
  look(press, fx.views.false[String(n)], 'preset ' + p.name, fx.texts[String(n)]);
  want('preset ' + p.name + ' slider', press['pfw-len'].value, String(fx.stops.indexOf(n)));
  want('preset ' + p.name + ' pressed', press['pfw-btns'].children.map((b) => b.attrs['aria-pressed']), fx.presets.map((q, j) => String(j === k)));
  want('preset ' + p.name + ' spoken', press['pfw-live'].textContent, fx.texts[String(n)] + ' tokens. ' + fx.views.false[String(n)].ttft + '. ' + fx.views.false[String(n)].tpot);
});
// ticking the box recomputes the current length and speaks
press['pfw-extra'].checked = true; fire(press['pfw-extra'], 'change');
const longest = String(fx.presets[2].tokens), ticked = fx.views.true[longest];
look(press, ticked, 'box ticked at the longest prompt', fx.texts[longest]);
want('box ticked speaks', press['pfw-live'].textContent, fx.texts[longest] + ' tokens. ' + ticked.ttft + '. ' + ticked.tpot);

// 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({ stops: fx.stops.length, bad: bad.slice(0, 5) }));
process.exit(bad.length ? 1 : 0);
"""


def pass_record(p):
    return {"seconds": p.seconds, "bound": p.bound, "busy": p.busy}


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 slider stop must be put back."""
    drifted = HTML.replace('"params": 8030261248', '"params": 7000000000', 1).replace('"axisMs": 600', '"axisMs": 500', 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 formulas, so a misspelt id passed there and threw in the browser."""
    script = HTML[HTML.index("<script>"):]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", script)) | {f"pfw-{k}" for k in ("ttft", "tpot", "why", "busy")}
    assert "$('pfw-' + k)" in script                      # the four text lines are looked up by suffix; they are added above
    missing = sorted(i for i in lookups if i != "pfw-config" and f'id="{i}"' not in ROOT)
    assert not missing, missing
    assert 'id="pfw-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: the slider and the preset buttons (built by the script) count, and the checkbox stays 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)) + len(sim.PRESETS) <= 5
    assert "<details><summary>More settings</summary>" in ROOT and "<details open" not in ROOT
    assert re.search(r"<details><summary>More settings</summary>.*?type=\"checkbox\"", 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)", "--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_more_settings_note_quotes_what_the_code_computes():
    """Guards a stale sentence: the details text says the extras add about 13% to the first word and 6.5% to each later token at 8,000 tokens."""
    n = sim.MAX_PROMPT
    first = 100 * (sim.prefill(n, True).seconds / sim.prefill(n).seconds - 1)
    later = 100 * (sim.decode_step(n, True).seconds / sim.decode_step(n).seconds - 1)
    assert f"{first:.0f}%" == "13%" and f"{later:.1f}%" == "6.5%"
    assert "about 13% later and each later token about 6.5% later" in HTML and "At 8,000 tokens" in HTML


def test_widget_formulas_give_the_times_and_words_of_the_python_model(tmp_path):
    """Guards the widget's JavaScript drifting from prefill_sim.py: the lesson prints Python's numbers and the reader sees the
    widget's, so the two must agree to the last bit on every slider stop, with the extras on and off, and on lengths beyond the slider."""
    cases = []
    for extras in (False, True):
        for n in STOPS + EXTRA_LENGTHS:
            cases.append({"n": n, "extras": extras, "prefill": pass_record(sim.prefill(n, extras)), "decode": pass_record(sim.decode_step(n, extras)),
                          "view": sim.view(n, extras)})
    commas = [[n, sim.commas(n)] for n in (0, 7, 999, 1000, 8000, 12345, 1234567)]
    result = run_node(tmp_path, PARITY, {"cases": cases, "commas": commas}, lambda d: [str(d / "fixture.json"), str(WIDGET)])
    assert result["bad"] == [] and result["checks"] == 3 * len(cases) + len(commas) + 1
    assert {c["view"]["bound"] for c in cases} == {"memory", "arithmetic"}


def test_widget_page_shows_what_the_lesson_names(tmp_path):
    """Guards the part a reader sees: the opening state (8,000 tokens, 412 ms, a wide green block), every slider stop, the
    presets, the checkbox, what is spoken on release, the theme message, and the words that keep a best case from reading as a
    measurement (the visible caveat line, the closed panel, the legend named by limit, the live region)."""
    views = {str(extras).lower(): {str(n): sim.view(n, extras) for n in STOPS} for extras in (False, True)}
    opening = views["false"]["8000"]
    assert opening["ttft"] == "Best-case time to the first token: 412 ms" and opening["bound"] == "arithmetic"
    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
    assert "Left of the line: prefill, the whole prompt in one pass. Right of it: decode, one token per step." in outside
    leaves_out = re.search(r"<details><summary>What this toy leaves out</summary>(.*?)</details>", ROOT, re.S).group(1)
    for said in ("Network time, waiting in line behind other users", "Real servers are slower than these times", "Llama 3.1 8B with 16-bit weights on one A100"):
        assert said in leaves_out, said
    assert re.findall(r'<li><i data-bound="(\w+)"></i><span>(.*?)</span></li>', ROOT) == [("memory", "Limited by memory"), ("arithmetic", "Limited by arithmetic")]
    assert 'class="pfw-sr" id="pfw-live" role="status" aria-live="polite" aria-atomic="true"' in ROOT      # what a screen reader hears
    assert ".pfw-root .pfw-btns button { flex: 1 1 auto; white-space: nowrap; }" in HTML      # a button wraps to the next row, never its label
    assert views["false"]["20"]["bound"] == "memory" and views["false"]["20"]["prefillWidth"] == views["false"]["20"]["stepWidth"]
    ids = sorted(set(re.findall(r'id="(pfw-[\w-]+)"', HTML)))
    fixture = {"ids": ids, "stops": STOPS, "views": views, "texts": {str(n): sim.commas(n) for n in STOPS}, "opening": opening,
               "openingText": "8,000", "openingTokens": sim.OPENING_TOKENS, "drawnSteps": sim.DRAWN_STEPS, "presets": sim.PRESETS,
               "presetLabels": [f"{p['name']} ({sim.commas(p['tokens'])})" for p in sim.PRESETS], "axis": [f"{sim.AXIS_MS // 2} ms", f"{sim.AXIS_MS} ms"]}
    result = run_node(tmp_path, PAGE, fixture, lambda d: [str(WIDGET), str(d / "fixture.json")])
    assert result == {"stops": len(STOPS), "bad": []}