"""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.",
    }