"""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]