How does an AI server answer many people at once?

"""A step-by-step scheduler for one GPU that serves many users, on top of the time model of lesson 23 (why-ai-waits-before-the-first-word).

The model is Llama 3.1 8B with 16-bit weights on an NVIDIA A100 80GB SXM. One step of the server is one pass of the model over every token in it:
the next token of each user who is already writing, plus the whole prompt of each user who joins in that step. The pass reads every weight once,
whatever the number of users, and reads the attention cache of each user who is writing. Its time is the larger of two floors, as in lessons 19 and 23:

    memory floor = (2 P + cache bytes per token * sum of the writing users' contexts) / bandwidth
    math floor   = 2 P * (writing users + joining prompt tokens) / peak          (2 FLOPs per weight per token)

Both floors assume the chip runs at its data-sheet peak, so these are best cases, not measurements. The scheduler runs two policies over the same
requests, one step at a time: fixed groups ("fixed") and continuous batching ("continuous"). The workload is seeded and invented.

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

import math
import random
from dataclasses import dataclass
from typing import Callable, NamedTuple

# --- the chip and the model: copied from lesson 23 (prefill_sim.py), which copied the chip from lesson 19.
# test_batching_sim.py checks every value against ../../why-ai-waits-before-the-first-word/code/prefill_sim.py.
A100_PEAK_FLOPS = 312e12          # NVIDIA A100 data sheet, revision "Jun21", 80GB SXM column: 312 TFLOPS dense FP16
A100_BANDWIDTH = 2039e9           # the same column: 2,039 GB/s. GB is 10**9 bytes.
GPU_MEMORY_BYTES = 80e9           # the same column: "80GB HBM2e"; the sheet does not say which GB, read here as 80 billion bytes
FP16_BYTES = 2
HIDDEN, FFN, LAYERS, HEADS, KV_HEADS, VOCAB = 4096, 14336, 32, 32, 8, 128256     # Llama 3.1 8B, config.json
HEAD_DIM = HIDDEN // HEADS


def parameter_count() -> int:
    attention = 2 * HIDDEN * HIDDEN + 2 * HIDDEN * KV_HEADS * HEAD_DIM
    layer = attention + 3 * HIDDEN * FFN + 2 * HIDDEN
    return LAYERS * layer + 2 * VOCAB * HIDDEN + HIDDEN


PARAMS = parameter_count()
KV_BYTES_PER_TOKEN = 2 * LAYERS * KV_HEADS * HEAD_DIM * FP16_BYTES        # lesson 22: 128 KiB for each token of context


# --- one step ---------------------------------------------------------------------------------

@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:
        return max(self.math_s, self.memory_s)

    @property
    def bound(self) -> str:
        """A tie counts as arithmetic-limited, as in lessons 19 and 23."""
        return "arithmetic" if self.flops / self.nbytes >= A100_PEAK_FLOPS / A100_BANDWIDTH else "memory"

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


def batch_step(contexts: list[int], joining_tokens: int = 0) -> Pass:
    """One step for users who are writing (`contexts` holds how many tokens each one has in its cache) and for prompts that join (`joining_tokens` in all).
    Left out: the attention arithmetic, the cache a joining prompt writes and the attention among prompt tokens (lesson 23's "extras"), all small next to the rest."""
    tokens = len(contexts) + joining_tokens
    return Pass(2 * PARAMS * tokens, FP16_BYTES * PARAMS + KV_BYTES_PER_TOKEN * sum(contexts))


def step_seconds(contexts: list[int], joining_tokens: int = 0) -> float:
    return batch_step(contexts, joining_tokens).seconds


def users_that_fit(context: int) -> int:
    """How many users, each with `context` tokens in the cache, fit in the chip's memory beside the weights (best case: all the rest is cache)."""
    return int((GPU_MEMORY_BYTES - FP16_BYTES * PARAMS) // (KV_BYTES_PER_TOKEN * context))


# --- the scheduler ----------------------------------------------------------------------------

class Request(NamedTuple):
    arrival: float       # seconds since the first request
    prompt: int          # tokens in the prompt
    reply: int           # tokens in the reply, at least 1


class Step(NamedTuple):
    start: float
    end: float
    slots: list          # for each slot, the id of the request running in it during this step, or -1
    joined: list         # ids that joined at the start of this step (prefill happens in it)
    finished: list       # ids whose last token came out at the end of this step


class Run(NamedTuple):
    steps: list
    first_token: list    # for each request, the time its first token came out
    finished: list       # for each request, the time its last token came out
    join_step: list      # for each request, the index of the step it joined in


def simulate(requests: list[Request], slots: int, policy: str, cost: Callable[[list[int], int], float] = step_seconds) -> Run:
    """Run `requests` (sorted by arrival) through `slots` slots, one step at a time. A request's id is its index.

    fixed:      new requests are taken only when every slot is empty, so a group runs until its longest reply ends and finished users' slots sit empty.
    continuous: at the start of every step, each empty slot takes the next waiting request.
    Either way, an idle server waits for the next arrival. A request that joins has its prompt read in the step it joins, and its first token comes
    out at the end of that step; each later step writes one more token. `cost(contexts, joining_tokens)` is the step time."""
    assert policy in ("fixed", "continuous")
    n = len(requests)
    held, made = [-1] * slots, [0] * n
    first, last, steps, join_step = [0.0] * n, [0.0] * n, [], [0] * n
    now, nxt, done = 0.0, 0, 0
    while done < n:
        joined = []
        if policy == "continuous" or all(h < 0 for h in held):
            if all(h < 0 for h in held) and requests[nxt].arrival > now:
                now = requests[nxt].arrival                         # nothing is running: wait for the next request
            for s in range(slots):
                if held[s] < 0 and nxt < n and requests[nxt].arrival <= now:
                    held[s] = nxt
                    joined.append(nxt)
                    join_step[nxt] = len(steps)
                    nxt += 1
        writing = [h for h in held if h >= 0 and h not in joined]
        dt = cost([requests[i].prompt + made[i] for i in writing], sum(requests[i].prompt for i in joined))
        start, now = now, now + dt
        snapshot, finished = list(held), []
        for s, h in enumerate(held):
            if h < 0:
                continue
            made[h] += 1
            if h in joined:
                first[h] = now
            if made[h] == requests[h].reply:
                last[h] = now
                held[s] = -1
                finished.append(h)
        steps.append(Step(start, now, snapshot, joined, finished))
        done += len(finished)
    return Run(steps, first, last, join_step)


def report(run: Run, requests: list[Request], slots: int) -> dict:
    """What the lesson and the widget say about a run. Sums are plain left-to-right loops so that JavaScript gives the same bits.
    `slots_in_use` is the share of slot-seconds in use while the server was running steps (waiting for the next request is not counted)."""
    waits = [run.first_token[i] - requests[i].arrival for i in range(len(requests))]
    total_wait = tokens = 0.0
    for w in waits:
        total_wait += w
    for r in requests:
        tokens += r.reply
    used = busy = 0.0
    for st in run.steps:
        used += sum(1 for h in st.slots if h >= 0) * (st.end - st.start)
        busy += st.end - st.start
    seconds = run.steps[-1].end
    return {"replies": len(requests), "seconds": seconds, "tokens_per_second": tokens / seconds, "average_wait": total_wait / len(requests),
            "slowest_wait": max(waits), "slots_in_use": used / (slots * busy), "steps": len(run.steps)}


def snapshot(run: Run, requests: list[Request], t: float) -> dict:
    """What one lane shows at simulated time t (seconds): for each slot how much of its reply is written (0 to 1, or -1 when empty),
    how many requests have arrived and wait for a slot, and how many replies are finished."""
    k = len(run.steps) - 1
    while k >= 0 and run.steps[k].start > t:
        k -= 1
    shares = [-1.0] * len(run.steps[0].slots)
    if k >= 0 and t < run.steps[k].end:
        for s, h in enumerate(run.steps[k].slots):
            if h >= 0:
                shares[s] = (k - run.join_step[h] + 1) / requests[h].reply
    waiting = sum(1 for i, r in enumerate(requests) if r.arrival <= t < run.steps[run.join_step[i]].start)
    return {"slots": shares, "in_use": sum(1 for x in shares if x >= 0), "waiting": waiting, "done": sum(1 for f in run.finished if f <= t)}


# --- the invented workload ----------------------------------------------------------------------
# A seeded draw, not a measurement. Each request has a prompt of 30 to 600 tokens (equal steps on a log scale, mean about 190) and a reply whose length
# depends on the "mix" setting. Requests arrive at random, with the gaps between them drawn once; a traffic setting only scales the gaps.

SEED = 2026
N_REQUESTS = 48
RATES = [0.5, 1, 2, 4, 8, 16]                      # new requests per second, the traffic slider's stops
SLOT_CHOICES = [4, 8, 16, 32]
MIXES = {"same": (60, 60), "mixed": (10, 300), "very mixed": (5, 600)}   # reply length from low to high, equal steps on a log scale


def draw_workload(seed: int = SEED, n: int = N_REQUESTS) -> dict:
    rng = random.Random(seed)
    gaps = [rng.expovariate(1.0) for _ in range(n)]                   # seconds between requests at one request a second
    prompts = [round(30 * (600 / 30) ** rng.random()) for _ in range(n)]
    where = [rng.random() for _ in range(n)]
    replies = {name: [round(lo * (hi / lo) ** x) for x in where] for name, (lo, hi) in MIXES.items()}
    return {"gaps": gaps, "prompts": prompts, "replies": replies}


WORKLOAD = draw_workload()


def make_requests(rate: float, mix: str, workload: dict = WORKLOAD) -> list[Request]:
    """The same requests at any traffic: only the gaps scale. The first request arrives at time 0."""
    out, t = [], 0.0
    for gap, prompt, reply in zip(workload["gaps"], workload["prompts"], workload["replies"][mix]):
        out.append(Request(t, prompt, reply))
        t += gap / rate
    return out


# --- what the widget starts with, and the words it shows (widget.html repeats these in JavaScript; test_widget.py compares them) --------

OPENING = {"rate": 8, "slots": 8, "mix": "mixed", "share": 0.4}     # the opening state: heavy-ish traffic, 40% of the way through the longer run
QUEUE_CAP = 16                                                       # little squares drawn for waiting requests; more are counted as "+N"
TIME_TICKS = 1000                                                    # the time slider has this many positions
PLAY_SECONDS = 10                                                    # Play takes about this long for the longer run


def mix_names(workload: dict = WORKLOAD) -> dict:
    names = {"same": "Same length", "mixed": "Mixed", "very mixed": "Very mixed"}
    out = {}
    for key, replies in workload["replies"].items():
        lo, hi = min(replies), max(replies)
        out[key] = f"{names[key]} ({lo} tokens)" if lo == hi else f"{names[key]} ({lo} to {hi} tokens)"
    return out


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


def fmt_seconds(s: float) -> str:
    return f"{s:.2f}" if s < 10 else f"{s:.1f}"


def rate_text(rate: float) -> str:
    return f"{rate:g}"


def time_at(tick: int, longest_end: float) -> float:
    """Simulated seconds at a position of the time slider."""
    return tick / TIME_TICKS * longest_end


def now_text(t: float, longest_end: float) -> str:
    return f"{fmt_seconds(t)} s of {fmt_seconds(longest_end)} s"


def state_line(snap: dict, slots: int, total: int) -> str:
    return f"In use {snap['in_use']} of {slots}. {snap['waiting']} waiting. {snap['done']} of {total} done."


def summary_line(rep: dict) -> str:
    tokens = commas(int(f"{rep['tokens_per_second']:.0f}"))
    return (f"Whole run: done after {fmt_seconds(rep['seconds'])} s, {tokens} tokens a second. "
            f"First word after {fmt_seconds(rep['average_wait'])} s on average, {fmt_seconds(rep['slowest_wait'])} s at worst. "
            f"Slots in use: {math.floor(rep['slots_in_use'] * 100 + 0.5)}% on average.")

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

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


def config():
    """What the page needs to run the same scheduler: the model, the chip, the drawn requests and the settings it offers."""
    return {
        "params": sim.PARAMS, "bytesPerWeight": sim.FP16_BYTES, "kvBytes": sim.KV_BYTES_PER_TOKEN, "peak": sim.A100_PEAK_FLOPS, "bandwidth": sim.A100_BANDWIDTH,
        "gaps": sim.WORKLOAD["gaps"], "prompts": sim.WORKLOAD["prompts"], "replies": sim.WORKLOAD["replies"],
        "rates": sim.RATES, "slotChoices": sim.SLOT_CHOICES, "mixNames": sim.mix_names(), "opening": sim.OPENING,
        "queueCap": sim.QUEUE_CAP, "timeTicks": sim.TIME_TICKS, "playSeconds": sim.PLAY_SECONDS,
    }


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

"""Print every number the lesson quotes. Run: python3 show_numbers.py"""
import statistics

from batching_sim import (
    A100_BANDWIDTH, A100_PEAK_FLOPS, FP16_BYTES, GPU_MEMORY_BYTES, KV_BYTES_PER_TOKEN, MIXES, N_REQUESTS, OPENING, PARAMS, RATES, SEED, SLOT_CHOICES,
    WORKLOAD, Request, batch_step, make_requests, mix_names, now_text, report, simulate, snapshot, state_line, summary_line, time_at, users_that_fit,
)


def toy() -> None:
    requests = [Request(0.0, 1, r) for r in (2, 6, 3, 1, 4, 2)]
    for policy in ("fixed", "continuous"):
        run = simulate(requests, 4, policy, cost=lambda contexts, joining: 1.0)
        print(f"  {policy}: {len(run.steps)} steps; first tokens at {[run.first_token[i] for i in range(6)]}; last tokens at {run.finished}")
        for k, st in enumerate(run.steps):
            print(f"    step {k + 1}: slots {st.slots}")


def main() -> None:
    one = batch_step([1000])
    print(f"parameters {PARAMS:,}; weights {FP16_BYTES * PARAMS / 1e9:.2f} GB; chip {A100_PEAK_FLOPS / 1e12:.0f} TFLOP/s, {A100_BANDWIDTH / 1e9:.0f} GB/s, {GPU_MEMORY_BYTES / 1e9:.0f} GB")
    print(f"cache per token {KV_BYTES_PER_TOKEN:,} bytes ({KV_BYTES_PER_TOKEN / 1024:.0f} KiB); weights alone take {FP16_BYTES * PARAMS / A100_BANDWIDTH * 1000:.3f} ms to read")
    print(f"arithmetic per user {2 * PARAMS / A100_PEAK_FLOPS * 1e6:.1f} us; cache reading per user at 1,000 tokens {1000 * KV_BYTES_PER_TOKEN / A100_BANDWIDTH * 1e6:.1f} us, at 100 tokens {100 * KV_BYTES_PER_TOKEN / A100_BANDWIDTH * 1e6:.1f} us")
    print()
    print("one step, every user with the same context")
    print("context | users | step ms | limited by | server tokens/s | per user tokens/s | cache GB read | busy")
    for context in (1000, 100):
        for users in (1, 8, 64, 175, 256, 487):
            if users > users_that_fit(context):
                continue
            p = batch_step([context] * users)
            print(f"{context:7,} | {users:5} | {p.seconds * 1000:7.2f} | {p.bound:10} | {users / p.seconds:15,.0f} | {1 / p.seconds:17.1f} | {users * context * KV_BYTES_PER_TOKEN / 1e9:6.2f} | {p.busy * 100:4.1f}%")
    print(f"  one user at 1,000 tokens: {one.seconds * 1000:.2f} ms (weights alone {FP16_BYTES * PARAMS / A100_BANDWIDTH * 1000:.2f} ms, so the cache adds {(one.seconds - FP16_BYTES * PARAMS / A100_BANDWIDTH) * 1000:.2f} ms)")
    first = next(n for n in range(1, 10000) if batch_step([100] * n).bound == "arithmetic")
    print(f"  at 100 tokens of context the step turns arithmetic-limited from {first} users; at 1,000 tokens never up to {users_that_fit(1000)}")
    print(f"  64 users against 1 user at 1,000 tokens: step time x{batch_step([1000] * 64).seconds / one.seconds:.2f}, server tokens/s x{64 * one.seconds / batch_step([1000] * 64).seconds:.1f}")
    print()
    print("users that fit beside the weights (best case: all the rest is cache)")
    for context in (100, 1000, 2000, 8192):
        print(f"  context {context:6,}: cache {context * KV_BYTES_PER_TOKEN / 1e9:6.3f} GB each -> {users_that_fit(context):,} users; room for cache {(GPU_MEMORY_BYTES - FP16_BYTES * PARAMS) / 1e9:.2f} GB")
    print(f"  4,000 tokens each: {users_that_fit(4000)} users")
    print()
    print("the worked table, and the steps that fill the memory")
    for contexts in ([1000], [1000] * 8, [1000] * 64):
        p = batch_step(contexts)
        print(f"  {len(contexts):3} users at 1,000 tokens: math {p.math_s * 1000:.4f} ms, memory {p.memory_s * 1000:.4f} ms, step {p.seconds * 1000:.2f} ms,"
              f" {len(contexts) / p.seconds:,.1f} tokens/s in all, {1 / p.seconds:.2f} per user, {p.nbytes / 1e9:.2f} GB read")
    for users, context in ((users_that_fit(1000), 1000), (users_that_fit(2000), 2000)):
        p = batch_step([context] * users)
        print(f"  memory full, {users} users at {context:,} tokens: step {p.seconds * 1000:.2f} ms, {p.nbytes / 1e9:.2f} GB read")
    joiner = batch_step([300] * 7, 577)
    print(f"  a 577-token prompt joins 7 users at 300 tokens: step {joiner.seconds * 1000:.2f} ms ({joiner.bound}) against {batch_step([300] * 7).seconds * 1000:.2f} ms without it")
    print()
    print("toy example: six requests at time zero, four slots, unit steps, replies 2 6 3 1 4 2")
    toy()
    print()
    prompts, replies = [r.prompt for r in make_requests(1, "mixed")], {m: [r.reply for r in make_requests(1, m)] for m in MIXES}
    print(f"workload: seed {SEED}, {N_REQUESTS} requests; prompts {min(prompts)} to {max(prompts)} (mean {statistics.fmean(prompts):.1f})")
    for m in MIXES:
        print(f"  {mix_names()[m]}: replies {min(replies[m])} to {max(replies[m])}, mean {statistics.fmean(replies[m]):.1f}, total {sum(replies[m]):,}")
    print(f"  first arrivals at traffic 1: {[round(r.arrival, 2) for r in make_requests(1, 'mixed')[:5]]}; gaps sum {sum(WORKLOAD['gaps']):.1f} s at one a second")
    slow = make_requests(0.5, "mixed")
    print(f"  at 0.5 a second the last request arrives at {slow[-1].arrival:.1f} s and its reply is {slow[-1].reply} tokens")
    print()
    print("results (8 slots unless noted): seconds | tokens/s | average wait | slowest wait | slots in use")
    for mix in MIXES:
        for slots in (8,):
            for rate in RATES:
                requests = make_requests(rate, mix)
                a, b = (report(simulate(requests, slots, p), requests, slots) for p in ("fixed", "continuous"))
                print(f"  {mix:10} S={slots:2} rate {rate:4}: fixed {a['seconds']:6.1f} s {a['tokens_per_second']:5.0f}/s wait {a['average_wait']:5.2f}/{a['slowest_wait']:5.2f} use {a['slots_in_use'] * 100:3.0f}%"
                      f" | continuous {b['seconds']:6.1f} s {b['tokens_per_second']:5.0f}/s wait {b['average_wait']:5.2f}/{b['slowest_wait']:5.2f} use {b['slots_in_use'] * 100:3.0f}%")
    never_worse = all(
        report(simulate(make_requests(r, m), s, "continuous"), make_requests(r, m), s)[k] <= report(simulate(make_requests(r, m), s, "fixed"), make_requests(r, m), s)[k]
        for m in MIXES for s in SLOT_CHOICES for r in RATES for k in ("seconds", "average_wait", "slowest_wait"))
    print(f"  continuous never slower or longer-waiting than fixed in all {len(MIXES) * len(SLOT_CHOICES) * len(RATES)} settings: {never_worse}")
    print()
    print("opening state of the widget")
    requests = make_requests(OPENING["rate"], OPENING["mix"])
    runs = [simulate(requests, OPENING["slots"], p) for p in ("fixed", "continuous")]
    reps = [report(run, requests, OPENING["slots"]) for run in runs]
    end = max(rep["seconds"] for rep in reps)
    t = time_at(round(OPENING["share"] * 1000), end)
    print(f"  {now_text(t, end)}")
    for p, run, rep in zip(("fixed", "continuous"), runs, reps):
        print(f"  {p}: {state_line(snapshot(run, requests, t), OPENING['slots'], N_REQUESTS)}")
        print(f"    {summary_line(rep)}")


if __name__ == "__main__":
    main()

"""Guards for batching_sim.py: the chip and model it copies from lesson 23, the one-step time, the two schedulers, the invented workload, and the numbers the lesson prints."""
import importlib.util
import sys
from pathlib import Path

import pytest

import batching_sim as sim

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"
GIB = 2**30
GROUPS = [(rate, mix, slots) for mix in sim.MIXES for slots in sim.SLOT_CHOICES for rate in sim.RATES]


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


def toy(policy):
    """Four slots, six requests that all arrive at once with replies of 2, 6, 3, 1, 4 and 2 tokens, every step taking one second."""
    requests = [sim.Request(0.0, 10, reply) for reply in (2, 6, 3, 1, 4, 2)]
    return requests, sim.simulate(requests, 4, policy, cost=lambda contexts, joining: 1.0)


def test_model_and_chip_equal_lesson_23s_and_a_lone_user_costs_what_it_did_there():
    """Guards a copy that drifted: this lesson quotes lesson 23's 7.88 ms step and 128 KiB per token, so one user alone must still cost exactly that."""
    ours, theirs = sim, lesson_23()
    for name in ("PARAMS", "KV_BYTES_PER_TOKEN", "A100_PEAK_FLOPS", "A100_BANDWIDTH", "FP16_BYTES"):
        assert getattr(ours, name) == getattr(theirs, name), name
    assert ours.batch_step([0]).seconds == theirs.decode_step(1).seconds == ours.FP16_BYTES * ours.PARAMS / ours.A100_BANDWIDTH
    for n in (1, 100, 154, 600):                          # a prompt that joins an empty server is lesson 23's prefill
        assert ours.batch_step([], n).seconds == theirs.prefill(n).seconds
    for context in (1, 1000, 8000):                       # a user alone reads its cache and nothing else (lesson 23, with the cache counted)
        assert ours.batch_step([context]).memory_s == theirs.decode_step(context, True).memory_s


def test_numbers_the_lesson_prints_for_one_step():
    """Pins the step times, rates and cache sizes in lesson.md and the figures, so a change to the model cannot leave stale prose."""
    step = lambda users, context: sim.batch_step([context] * users)
    ms = lambda users, context: round(step(users, context).seconds * 1000, 2)
    assert [ms(1, 1000), ms(64, 1000), ms(256, 1000)] == [7.94, 11.99, 24.33]
    assert [round(users / step(users, 1000).seconds) for users in (1, 64, 256)] == [126, 5337, 10521]
    assert [round(1 / step(users, 1000).seconds, 1) for users in (1, 64, 256)] == [125.9, 83.4, 41.1]
    assert [round(users * 1000 * sim.KV_BYTES_PER_TOKEN / 1e9, 2) for users in (1, 64, 256)] == [0.13, 8.39, 33.55]
    assert round(sim.FP16_BYTES * sim.PARAMS / sim.A100_BANDWIDTH * 1000, 2) == 7.88 and ms(64, 100) == 8.29
    assert all(step(users, 1000).bound == "memory" for users in (1, 64, 256, 487))
    assert round(sim.batch_step([1000] * 64).busy * 100, 1) == 27.5 and sim.batch_step([1000]).busy < 0.01     # the share of the step the arithmetic is busy


def test_each_user_adds_more_reading_than_arithmetic_at_a_thousand_tokens_but_not_at_a_hundred():
    """Guards the lesson's claim about why the chip stays memory-limited: reading one user's cache costs 64 us at 1,000 tokens and
    51 us of arithmetic is all a user adds, so the step never turns arithmetic-limited. At 100 tokens it does, from 175 users on."""
    us = lambda seconds: round(seconds * 1e6, 1)
    math_per_user = sim.batch_step([0]).math_s
    assert us(math_per_user) == 51.5 and us(sim.KV_BYTES_PER_TOKEN * 1000 / sim.A100_BANDWIDTH) == 64.3
    assert not any(sim.batch_step([1000] * u).bound == "arithmetic" for u in range(1, sim.users_that_fit(1000) + 1))
    assert sim.batch_step([100] * 174).bound == "memory" and sim.batch_step([100] * 175).bound == "arithmetic"


def test_users_that_fit_in_memory_beside_the_weights():
    """Guards the memory cap the lesson states (best case, all the rest of the 80 GB is cache): 487, 243 and 59 users."""
    assert [sim.users_that_fit(c) for c in (1000, 2000, 8192)] == [487, 243, 59]
    assert (sim.GPU_MEMORY_BYTES - sim.FP16_BYTES * sim.PARAMS) / 1e9 == pytest.approx(63.94, abs=5e-3) and 8192 * sim.KV_BYTES_PER_TOKEN == GIB


@pytest.mark.parametrize("policy", ["fixed", "continuous"])
def test_toy_example_slot_by_slot(policy):
    """Guards the hand-worked example in the lesson (figure 1 and the Go-deeper table): which request sits in which slot at each step."""
    requests, run = toy(policy)
    slots = [step.slots for step in run.steps]
    if policy == "fixed":
        assert slots == [[0, 1, 2, 3], [0, 1, 2, -1], [-1, 1, 2, -1], [-1, 1, -1, -1], [-1, 1, -1, -1], [-1, 1, -1, -1],
                         [4, 5, -1, -1], [4, 5, -1, -1], [4, -1, -1, -1], [4, -1, -1, -1]]
        assert [run.first_token[i] for i in (4, 5)] == [7.0, 7.0] and run.steps[-1].end == 10.0
    else:
        assert slots == [[0, 1, 2, 3], [0, 1, 2, 4], [5, 1, 2, 4], [5, 1, -1, 4], [-1, 1, -1, 4], [-1, 1, -1, -1]]
        assert [run.first_token[i] for i in (4, 5)] == [2.0, 3.0] and run.steps[-1].end == 6.0
    assert run.finished[1] == 6.0 and [s.finished for s in run.steps][:3] == [[3], [0], [2]]


def test_snapshot_reads_a_lane_in_the_middle_of_a_step():
    """Guards the widget's picture: half way through step 0 each slot shows the share of its reply written (a request of 2, 6, 3 and 1 tokens),
    and two requests are counted as waiting until the step they join starts."""
    requests, run = toy("continuous")
    assert sim.snapshot(run, requests, 0.5) == {"slots": [0.5, 1 / 6, 1 / 3, 1.0], "in_use": 4, "waiting": 2, "done": 0}
    assert sim.snapshot(run, requests, 1.5)["slots"] == [1.0, 2 / 6, 2 / 3, 0.25] and sim.snapshot(run, requests, 1.5)["done"] == 1
    assert sim.snapshot(run, requests, 5.9) == {"slots": [-1.0, 6 / 6, -1.0, -1.0], "in_use": 1, "waiting": 0, "done": 5}
    assert sim.snapshot(run, requests, 6.0) == {"slots": [-1.0] * 4, "in_use": 0, "waiting": 0, "done": 6}


def test_an_idle_server_waits_for_the_next_request_and_a_lone_prompt_costs_lesson_23s_prefill():
    """Guards the clock jump: a request that arrives at 10 s finds nothing running, so its first token comes one joining step later, not at 0 s plus anything."""
    prefill = lesson_23().prefill
    requests = [sim.Request(0.0, 100, 3), sim.Request(10.0, 600, 2)]
    for policy in ("fixed", "continuous"):
        run = sim.simulate(requests, 2, policy)
        assert run.steps[-2].start == 10.0 and run.first_token[1] - 10.0 == pytest.approx(prefill(600).seconds, rel=1e-9)
        assert run.first_token[0] == prefill(100).seconds and run.steps[-1].end == run.finished[1]


@pytest.mark.parametrize("policy", ["fixed", "continuous"])
def test_every_setting_obeys_the_rules_of_its_policy(policy):
    """Guards scheduler slips a reader would see as a bug in the picture, in every traffic, reply and slot setting: a request in two places, a reply with
    the wrong number of steps, a start before arrival, a first token out of order, a fixed group that takes a newcomer early, or a continuous lane
    with an empty slot while someone waits."""
    for rate, mix, slots in GROUPS:
        requests = sim.make_requests(rate, mix)
        run = sim.simulate(requests, slots, policy)
        where = {i: [] for i in range(len(requests))}
        for k, step in enumerate(run.steps):
            assert len(step.slots) == slots and len(set(h for h in step.slots if h >= 0)) == sum(1 for h in step.slots if h >= 0), (rate, mix, slots, k)
            assert step.end > step.start and (k == 0 or step.start >= run.steps[k - 1].end)
            for s, h in enumerate(step.slots):
                if h >= 0:
                    where[h].append((k, s))
            writing = [h for h in step.slots if h >= 0 and h not in step.joined]
            if policy == "fixed" and step.joined:
                assert not writing, (rate, mix, slots, k)                       # a new group starts only into an empty lane
            if policy == "continuous" and -1 in step.slots:
                assert all(run.join_step[i] <= k for i, r in enumerate(requests) if r.arrival <= step.start), (rate, mix, slots, k)
        for i, r in enumerate(requests):
            assert [k for k, _ in where[i]] == list(range(run.join_step[i], run.join_step[i] + r.reply)), (rate, mix, slots, i)
            assert len({s for _, s in where[i]}) == 1                          # a reply stays in its slot
            assert run.steps[run.join_step[i]].start >= r.arrival and run.first_token[i] == run.steps[run.join_step[i]].end
            assert run.first_token[i] <= run.finished[i] == run.steps[run.join_step[i] + r.reply - 1].end


def test_equal_replies_leave_no_slot_empty_inside_a_fixed_group_so_fixed_only_waits():
    """Guards the lesson's caveat that with same-length replies fixed groups waste no slots (a group finishes all at once) and lose only by making newcomers wait."""
    for rate in sim.RATES:
        requests = sim.make_requests(rate, "same")
        run = sim.simulate(requests, 8, "fixed")
        assert all(len(step.finished) in (0, sum(1 for h in step.slots if h >= 0)) for step in run.steps), rate


def test_continuous_is_never_slower_and_never_makes_anyone_wait_longer_in_any_setting_here():
    """Guards the lesson's claim that continuous batching is at least as good on the same requests: the same replies finish no later and the first word comes no later,
    on average and at worst, in every one of the 72 settings (it is not a law: a different workload could break it)."""
    for rate, mix, slots in GROUPS:
        requests = sim.make_requests(rate, mix)
        fixed, continuous = (sim.report(sim.simulate(requests, slots, p), requests, slots) for p in ("fixed", "continuous"))
        for key in ("seconds", "average_wait", "slowest_wait"):
            assert continuous[key] <= fixed[key] * (1 + 1e-12), (rate, mix, slots, key)


def test_numbers_the_lesson_prints_for_the_two_schedulers():
    """Pins the results in lesson.md and figure 4 for 8 slots and mixed replies: light traffic (0.5 a second), the opening state (8) and heavy traffic (16)."""
    def results(rate, mix="mixed", slots=8):
        requests = sim.make_requests(rate, mix)
        return [sim.report(sim.simulate(requests, slots, p), requests, slots) for p in ("fixed", "continuous")]
    line = lambda rep: (round(rep["seconds"], 1), round(rep["tokens_per_second"]), round(rep["average_wait"], 2), round(rep["slowest_wait"], 2), round(rep["slots_in_use"] * 100))
    assert [line(r) for r in results(0.5)] == [(103.9, 47, 0.19, 1.91, 13), (103.9, 47, 0.01, 0.03, 14)]
    assert [line(r) for r in results(8)] == [(14.1, 349, 3.43, 6.18, 38), (8.4, 582, 0.03, 0.33, 61)]
    assert [line(r) for r in results(16)] == [(14.1, 349, 5.02, 9.1, 38), (6.3, 777, 0.53, 1.27, 80)]
    assert [line(r)[:4] for r in results(8, "same")] == [(7.2, 400, 0.31, 0.54), (6.9, 420, 0.02, 0.04)]


def test_workload_is_the_seeded_draw_the_lesson_describes():
    """Guards a changed draw: the lesson says 48 requests, prompts of 30 to 600 tokens averaging about 190, and three reply mixes of 60, 10 to 268 and 5 to 513 tokens."""
    w = sim.WORKLOAD
    assert sim.draw_workload() == w and len(w["gaps"]) == len(w["prompts"]) == 48 and all(len(r) == 48 for r in w["replies"].values())
    assert min(w["prompts"]) >= 30 and max(w["prompts"]) <= 600 and round(sum(w["prompts"]) / 48) == 191
    assert [(min(r), max(r)) for r in w["replies"].values()] == [(60, 60), (10, 268), (5, 513)]
    assert [round(sum(r) / 48, 1) for r in w["replies"].values()] == [60.0, 102.4, 157.4]
    assert sim.mix_names() == {"same": "Same length (60 tokens)", "mixed": "Mixed (10 to 268 tokens)", "very mixed": "Very mixed (5 to 513 tokens)"}


def test_traffic_only_scales_the_gaps_between_arrivals():
    """Guards the widget's promise that the slider changes only the traffic: the same prompts and replies arrive at every rate, 8 times as fast at 8 times the rate."""
    slow, fast = sim.make_requests(1, "mixed"), sim.make_requests(8, "mixed")
    assert slow[0].arrival == fast[0].arrival == 0.0
    assert [(r.prompt, r.reply) for r in slow] == [(r.prompt, r.reply) for r in fast]
    assert all(s.arrival / 8 == pytest.approx(f.arrival, rel=1e-12) for s, f in zip(slow, fast))
    assert all(a.arrival <= b.arrival for a, b in zip(slow, slow[1:]))


def test_the_opening_state_shows_the_point_at_once():
    """Guards the first screen: at the opening setting, 40% in, the fixed lane has most slots empty while 22 wait, and the continuous lane has most slots full and nobody waiting."""
    o = sim.OPENING
    requests = sim.make_requests(o["rate"], o["mix"])
    runs = [sim.simulate(requests, o["slots"], p) for p in ("fixed", "continuous")]
    t = sim.time_at(round(o["share"] * sim.TIME_TICKS), max(r.steps[-1].end for r in runs))
    fixed, continuous = (sim.snapshot(r, requests, t) for r in runs)
    assert (fixed["in_use"], fixed["waiting"], fixed["done"]) == (2, 22, 15) and (continuous["in_use"], continuous["waiting"], continuous["done"]) == (7, 0, 32)
    assert sim.state_line(fixed, 8, 48) == "In use 2 of 8. 22 waiting. 15 of 48 done."
    assert sim.fmt_seconds(t) == "5.63"

"""Guards for ../widget.html: its data block, its size, that its scheduler gives the same steps and words as batching_sim.py, and what a reader sees."""
import json
import math
import re
import subprocess
from functools import lru_cache

import batching_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="bw-root">'):HTML.index("<script")]
POLICIES = ("fixed", "continuous")
SETTINGS = [(rate, mix, slots) for mix in sim.MIXES for slots in sim.SLOT_CHOICES for rate in sim.RATES]
TICKS = [0, 1, 3, 250, 400, 500, 777, 999, 1000]

# Reads the scheduler out of the widget between its marker comments, runs it with Node, and compares every step, time, count and 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="bw-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 {makeRequests, simulate, report, snapshot, commas, fmtSeconds, timeAt, nowText, stateLine, summaryLine, stepSeconds};')(cfg);
const bad = [];
let checks = 0;
const same = (what, got, want) => { checks++; if (JSON.stringify(got) !== JSON.stringify(want)) bad.push([what, JSON.stringify(got).slice(0, 200), JSON.stringify(want).slice(0, 200)]); };
fx.cases.forEach((c) => {
  const what = c.rate + '/' + c.mix + '/' + c.slots + '/' + c.policy;
  const reqs = api.makeRequests(c.rate, c.mix);
  const run = api.simulate(reqs, c.slots, c.policy);
  same(what + ' arrivals', reqs.map((r) => r.arrival), c.arrivals);
  same(what + ' steps', run.steps.map((s) => [s.start, s.end, s.slots, s.joined, s.finished]), c.steps);
  same(what + ' first tokens', run.first, c.first);
  same(what + ' finish times', run.finished, c.finished);
  same(what + ' join steps', run.joinStep, c.joinStep);
  const rep = api.report(run, reqs, c.slots);
  same(what + ' report', [rep.seconds, rep.tokensPerSecond, rep.averageWait, rep.slowestWait, rep.slotsInUse], c.report);
  same(what + ' summary', api.summaryLine(rep), c.summary);
  c.snaps.forEach((s) => {
    const snap = api.snapshot(run, reqs, s.t);
    same(what + ' snapshot at ' + s.t, [snap.slots, snap.inUse, snap.waiting, snap.done], s.snap);
    same(what + ' state line at ' + s.t, api.stateLine(snap, c.slots, reqs.length), s.line);
  });
});
fx.words.forEach((w) => same('words ' + JSON.stringify(w.args), [api.commas(w.args[0]), api.fmtSeconds(w.args[1]), api.nowText(w.args[1], w.args[2]), api.timeAt(w.args[3], w.args[2])], w.want));
same('cost of one step', [api.stepSeconds(0, 0, 100), api.stepSeconds(1, 1000, 0), api.stepSeconds(64, 64000, 0)], fx.steps);
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 sliders, presses Play, changes the settings 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="bw-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, hidden: false, style: {}, attrs: {}, textContent: '' };
  e.setAttribute = (k, v) => { e.attrs[k] = v; };
  e.getAttribute = (k) => e.attrs[k];
  e.addEventListener = (type, f) => (e.listeners[type] = e.listeners[type] || []).push(f);
  e.appendChild = (c) => (e.children.push(c), c);
  e.replaceChildren = (...c) => { e.children = c; };
  return e;
}
function open(reduced) {
  const page = { root: el('html'), messages: [], frames: [], parent: { other: true } };
  fx.ids.forEach((id) => { page[id] = el('x'); });
  page['bw-config'].textContent = config;
  const doc = { getElementById: (id) => page[id], createElement: el, documentElement: page.root };
  const win = { addEventListener: (type, f) => page.messages.push(f), matchMedia: () => ({ matches: !!reduced }), requestAnimationFrame: (f) => page.frames.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 = [];
let checks = 0;
const want = (what, got, expected) => { checks++; if (JSON.stringify(got) !== JSON.stringify(expected)) bad.push([what, JSON.stringify(got).slice(0, 160), JSON.stringify(expected).slice(0, 160)]); };
function look(page, v, what) {                       // compares the page with one expected view
  want(what + ': traffic words', page['bw-rate'].textContent, v.rate);
  want(what + ': time words', page['bw-now'].textContent, v.now);
  ['fixed', 'continuous'].forEach((name, k) => {
    const lane = v.lanes[k], cells = page['bw-' + name + '-cells'].children;
    want(what + ': ' + name + ' cells', cells.map((c) => [c.attrs['data-state'], c.children[0].style.width]), lane.cells.map((c) => [c.state, c.width]));
    want(what + ': ' + name + ' queue', page['bw-' + name + '-queue'].children.map((q) => q.attrs['data-on']), lane.squares);
    want(what + ': ' + name + ' more', page['bw-' + name + '-more'].textContent, lane.more);
    want(what + ': ' + name + ' state', page['bw-' + name + '-state'].textContent, lane.state);
    want(what + ': ' + name + ' summary', page['bw-' + name + '-summary'].textContent, lane.summary);
  });
}
const key = (rate, mix, slots, tick) => [rate, mix, slots, tick].join('|');
const spoken = (page) => page['bw-live'].textContent;

const first = open(false);
look(first, fx.views[key(8, 'mixed', 8, 400)], 'opening state');
want('opening controls', [first['bw-traffic'].value, first['bw-traffic'].max, first['bw-time'].value, first['bw-time'].max, first['bw-slots'].value, first['bw-mix'].value],
  [String(fx.rates.indexOf(8)), String(fx.rates.length - 1), '400', '1000', '8', 'mixed']);
want('option lists', [first['bw-slots'].children.map((o) => o.textContent), first['bw-mix'].children.map((o) => o.textContent)], [fx.slotChoices.map(String), fx.mixLabels]);
want('nothing spoken yet', spoken(first), '');
want('play is shown', first['bw-play'].hidden, false);

// every traffic stop, by dragging (input) and by releasing (change); the time position stays where it was
const traffic = open(false);
fx.rates.forEach((rate, i) => {
  traffic['bw-traffic'].value = String(i);
  const heard = spoken(traffic);
  fire(traffic['bw-traffic'], 'input');
  look(traffic, fx.views[key(rate, 'mixed', 8, 400)], 'traffic ' + rate);
  want('dragging stays silent at ' + rate, spoken(traffic), heard);
  want('slider reads out at ' + rate, traffic['bw-traffic'].attrs['aria-valuetext'], rate + ' new requests a second');
  fire(traffic['bw-traffic'], 'change');
  want('release speaks at ' + rate, spoken(traffic), fx.said.traffic[String(rate)]);
});

// the time slider at both ends and in between, and the words it speaks on release
const clock = open(false);
fx.ticks.forEach((tick) => {
  clock['bw-time'].value = String(tick);
  fire(clock['bw-time'], 'input');
  look(clock, fx.views[key(8, 'mixed', 8, tick)], 'tick ' + tick);
  fire(clock['bw-time'], 'change');
  want('time speaks at ' + tick, spoken(clock), fx.said.time[String(tick)]);
});

// the two settings under More settings
const more = open(false);
fx.slotChoices.forEach((n) => {
  more['bw-slots'].value = String(n);
  fire(more['bw-slots'], 'change');
  look(more, fx.views[key(8, 'mixed', n, 400)], 'slots ' + n);
  want('slots speak at ' + n, spoken(more), fx.said.slots[String(n)]);
});
more['bw-slots'].value = '8';
fx.mixes.forEach((mix) => {
  more['bw-mix'].value = mix;
  fire(more['bw-mix'], 'change');
  look(more, fx.views[key(8, mix, 8, 400)], 'mix ' + mix);
  want('mix speaks for ' + mix, spoken(more), fx.said.mix[mix]);
});

// Play: the first press plays the run from time 0; it moves the clock by the real time that passed, stops itself at the end, restarts from the start,
// goes on from where the clock is after Pause or a moved slider, and Pause and the slider stop it
const player = open(false);
fire(player['bw-play']);
want('play starts', [player['bw-play'].textContent, spoken(player), player.frames.length], ['Pause', 'Playing.', 1]);
let frames = player.frames.splice(0); frames[0](0);
look(player, fx.views[key(8, 'mixed', 8, 0)], 'first Play starts at time 0, not at the opening 40%');
frames = player.frames.splice(0); frames[0](1000);          // one second of real time is a tenth of the longer run
look(player, fx.views[key(8, 'mixed', 8, 100)], 'one second played');
want('time slider follows the clock', player['bw-time'].value, '100');
fire(player['bw-play']);
want('pause', [player['bw-play'].textContent, spoken(player)], ['Play', fx.said.time['100']]);
player.frames.splice(0);
fire(player['bw-play']);
frames = player.frames.splice(0); frames[0](5000);
look(player, fx.views[key(8, 'mixed', 8, 100)], 'Play after Pause goes on from the pause');
frames = player.frames.splice(0); frames[0](15000);
look(player, fx.views[key(8, 'mixed', 8, 1000)], 'end of the run');
want('stops itself at the end', [player['bw-play'].textContent, player.frames.length, spoken(player)], ['Play', 0, fx.said.finished]);
fire(player['bw-play']);
frames = player.frames.splice(0); frames[0](20000);
look(player, fx.views[key(8, 'mixed', 8, 0)], 'Play at the end starts over');
want('playing again', player['bw-play'].textContent, 'Pause');
player['bw-time'].value = '250'; fire(player['bw-time'], 'input');
want('moving the slider stops Play', [player['bw-play'].textContent, player['bw-time'].value], ['Play', '250']);
const dragged = open(false);                                  // a reader who put the clock back on the opening spot is still going on from there
dragged['bw-time'].value = '400'; fire(dragged['bw-time'], 'input');
fire(dragged['bw-play']);
frames = dragged.frames.splice(0); frames[0](0); frames = dragged.frames.splice(0); frames[0](1000);
look(dragged, fx.views[key(8, 'mixed', 8, 500)], 'Play after a moved slider goes on from the slider');

// reduced motion: no Play, nothing moves, and the page opens on the finished run
const calm = open(true);
want('reduced motion hides Play', calm['bw-play'].hidden, true);
look(calm, fx.views[key(8, 'mixed', 8, 1000)], 'reduced motion opens on the finished state');
want('reduced motion time position', calm['bw-time'].value, '1000');

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


@lru_cache(maxsize=None)
def runs(rate, mix, slots):
    requests = sim.make_requests(rate, mix)
    simulated = [sim.simulate(requests, slots, policy) for policy in POLICIES]
    return requests, simulated, [sim.report(run, requests, slots) for run in simulated]


def view(rate, mix, slots, tick):
    """What the page should show after it has settled on these settings and this time position."""
    requests, simulated, reports = runs(rate, mix, slots)
    end = max(rep["seconds"] for rep in reports)
    t = sim.time_at(tick, end)
    lanes = []
    for run, rep in zip(simulated, reports):
        snap = sim.snapshot(run, requests, t)
        waiting = snap["waiting"]
        lanes.append({
            "cells": [{"state": "use" if x >= 0 else "wait" if waiting else "empty", "width": f"{math.floor(x * 100)}%" if x >= 0 else "0%"} for x in snap["slots"]],
            "squares": ["1" if i < waiting else "0" for i in range(sim.QUEUE_CAP)],
            "more": f"+{waiting - sim.QUEUE_CAP}" if waiting > sim.QUEUE_CAP else "",
            "state": sim.state_line(snap, slots, len(requests)), "summary": sim.summary_line(rep),
        })
    return {"rate": sim.rate_text(rate), "now": sim.now_text(t, end), "lanes": lanes}


def said_after_a_run_setting(rate, mix, slots):
    _, _, reports = runs(rate, mix, slots)
    return f"{sim.rate_text(rate)} new requests a second. Fixed groups. {sim.summary_line(reports[0])} Continuous batching. {sim.summary_line(reports[1])}"


def said_at(tick, rate=8, mix="mixed", slots=8):
    v = view(rate, mix, slots, tick)
    return f"At {v['now']}. Fixed groups: {v['lanes'][0]['state']} Continuous batching: {v['lanes'][1]['state']}"


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")
    result = subprocess.run(["node", str(tmp_path / "check.mjs"), *args(tmp_path)], capture_output=True, text=True, timeout=300)
    assert result.returncode == 0, result.stdout + result.stderr
    return json.loads(result.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 a changed reply length must be put back."""
    drifted = HTML.replace('"params": 8030261248', '"params": 7000000000', 1).replace('"mixed": [', '"mixed": [999, ', 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 scheduler, so a misspelt id passed there and threw in the browser."""
    script = HTML[HTML.index("<script>"):]
    lookups = set(re.findall(r"\$\('([\w-]+)'\)", script))
    for suffix in ("cells", "queue", "more", "state", "summary"):                    # looked up per lane as $('bw-' + name + '-suffix')
        assert f"$('bw-' + name + '-{suffix}')" in script or f"$('bw-' + lane.name + '-{suffix}')" in script, suffix
        lookups |= {f"bw-{lane}-{suffix}" for lane in POLICIES}
    missing = sorted(i for i in lookups if i != "bw-config" and f'id="{i}"' not in ROOT)
    assert not missing, missing
    assert 'id="bw-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 two sliders and Play stay in view, and the two selects wait 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)) == 3
    assert "<details><summary>More settings</summary>" in ROOT and "<details open" not in ROOT
    assert re.search(r"<details><summary>More settings</summary>.*?<select id=\"bw-slots\">.*?<select id=\"bw-mix\">", 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)",
                 "--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", "#00855F", "#009E73", "#b36200", "#E69F00"}


def test_widget_scheduler_gives_the_steps_and_words_of_the_python_model(tmp_path):
    """Guards the widget's JavaScript drifting from batching_sim.py: the lesson prints Python's numbers and the reader sees the widget's, so every step of every run
    (all 72 settings, both policies) must agree to the last bit, and so must the totals, the snapshots at nine moments and the lines of words."""
    cases = []
    for rate, mix, slots in SETTINGS:
        requests, simulated, reports = runs(rate, mix, slots)
        end = max(rep["seconds"] for rep in reports)
        for policy, run, rep in zip(POLICIES, simulated, reports):
            snaps = []
            for tick in TICKS:
                t = sim.time_at(tick, end)
                snap = sim.snapshot(run, requests, t)
                snaps.append({"t": t, "snap": [snap["slots"], snap["in_use"], snap["waiting"], snap["done"]], "line": sim.state_line(snap, slots, len(requests))})
            cases.append({"rate": rate, "mix": mix, "slots": slots, "policy": policy, "arrivals": [r.arrival for r in requests],
                          "steps": [[s.start, s.end, s.slots, s.joined, s.finished] for s in run.steps], "first": run.first_token, "finished": run.finished,
                          "joinStep": run.join_step, "snaps": snaps, "summary": sim.summary_line(rep),
                          "report": [rep["seconds"], rep["tokens_per_second"], rep["average_wait"], rep["slowest_wait"], rep["slots_in_use"]]})
    words = [{"args": [n, s, e, k], "want": [sim.commas(n), sim.fmt_seconds(s), sim.now_text(s, e), sim.time_at(k, e)]}
             for n, s, e, k in ((7, 0.0, 14.1, 0), (1234, 9.995, 103.9, 400), (999999, 10.0, 8.45, 1000), (0, 0.034, 0.5, 1), (12, 3.4349, 3.4349, 777))]
    steps = [sim.batch_step([], 100).seconds, sim.batch_step([1000]).seconds, sim.batch_step([1000] * 64).seconds]
    result = run_node(tmp_path, PARITY, {"cases": cases, "words": words, "steps": steps}, lambda d: [str(d / "fixture.json"), str(WIDGET)])
    assert result["bad"] == [] and result["checks"] >= len(cases) * (6 + len(TICKS) * 2 + 1)


def test_widget_page_shows_what_the_lesson_names(tmp_path):
    """Guards the part a reader sees: the opening state (8 requests a second, 40% in), every traffic stop, the time slider, the two settings, Play and Pause, the
    finished state under reduced motion, what is spoken on release, the theme message, and the words that keep a best case from reading as a measurement."""
    outside = re.sub(r"<details>.*?</details>", "", ROOT, flags=re.S)
    assert "A model at the chip&rsquo;s peak speed, not a measurement. The 48 requests are invented." in outside and sim.N_REQUESTS == 48
    assert "A new group starts only when every slot is empty." in outside and "A waiting request takes any empty slot at once." 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", "running out of memory for the caches", "Real servers are slower than these times", "Llama 3.1 8B with 16-bit weights on one A100"):
        assert said in re.sub(r"\s+", " ", leaves_out), said
    assert re.findall(r'<li><i data-state="(\w+)"></i><span>(.*?)</span></li>', ROOT) == [("use", "Reply under way"), ("empty", "Empty slot"), ("wait", "Empty while people wait")]
    assert 'class="bw-sr" id="bw-live" role="status" aria-live="polite" aria-atomic="true"' in ROOT      # what a screen reader hears
    assert ".bw-root button { font: inherit; font-weight: 600; font-size: 0.9rem; line-height: 1.2; min-height: 44px; min-width: 72px; padding: 2px 10px; white-space: nowrap;" in HTML
    assert ".bw-root [hidden] { display: none; }" in HTML                                              # Play really disappears under reduced motion
    views = {}
    mixes = list(sim.MIXES)
    for rate in sim.RATES:
        views[f"{rate}|mixed|8|400"] = view(rate, "mixed", 8, 400)
    for tick in TICKS + [100]:
        views[f"8|mixed|8|{tick}"] = view(8, "mixed", 8, tick)
    for slots in sim.SLOT_CHOICES:
        views[f"8|mixed|{slots}|400"] = view(8, "mixed", slots, 400)
    for mix in mixes:
        views[f"8|{mix}|8|400"] = view(8, mix, 8, 400)
    opening = views["8|mixed|8|400"]
    assert opening["lanes"][0]["state"] == "In use 2 of 8. 22 waiting. 15 of 48 done." and opening["lanes"][1]["state"] == "In use 7 of 8. 0 waiting. 32 of 48 done."
    assert opening["lanes"][0]["more"] == "+6" and opening["lanes"][1]["more"] == "" and opening["now"] == "5.63 s of 14.1 s"
    assert {c["state"] for c in opening["lanes"][0]["cells"]} == {"use", "wait"} and {c["state"] for c in opening["lanes"][1]["cells"]} == {"use", "empty"}
    said = {"traffic": {str(rate): said_after_a_run_setting(rate, "mixed", 8) for rate in sim.RATES},
            "time": {str(tick): said_at(tick) for tick in TICKS + [100]},
            "slots": {str(n): said_after_a_run_setting(8, "mixed", n) for n in sim.SLOT_CHOICES},
            "mix": {mix: said_after_a_run_setting(8, mix, 8) for mix in mixes},
            "finished": "Finished. " + said_at(1000)[said_at(1000).index("Fixed groups:"):]}
    ids = sorted(set(re.findall(r'id="(bw-[\w-]+)"', HTML)))
    fixture = {"ids": ids, "views": views, "said": said, "rates": sim.RATES, "slotChoices": sim.SLOT_CHOICES, "mixes": mixes, "ticks": TICKS,
               "mixLabels": [sim.mix_names()[m] for m in mixes]}
    result = run_node(tmp_path, PAGE, fixture, lambda d: [str(WIDGET), str(d / "fixture.json")])
    assert result["bad"] == [] and result["checks"] > 200