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

Four parts, each short:

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

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

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

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


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


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

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


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


DRAFT_PARAMS = draft_parameter_count()


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

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


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


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


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


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


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


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


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


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


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


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


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


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


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

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


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


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


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


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


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

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


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


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

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


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


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


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


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


MAX_K = 32


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


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


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


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

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


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


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


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

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

    return draw


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


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


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