"""Guards for speculative_sim.py: the draft model's size, the chip it takes from lesson 23, the toy pair and the proof that its text is unchanged,
the sampling rule, the tokens-per-pass formula, the time model, and the numbers the lesson prints."""
from fractions import Fraction
import random

import pytest

import speculative_sim as sim

big = sim.big
DRAFT_MS = sim.draft_pass().seconds * 1000
RING = sim.WORDS * 3


def test_numbers_taken_from_lesson_23_are_the_ones_this_lesson_prints():
    """Guards a change to lesson 23's model or chip that would leave this lesson's prose stale: the weights, the 7.877 ms step, the ridge and the 153 / 154 flip."""
    assert big.PARAMS == 8_030_261_248 and big.A100_PEAK_FLOPS == 312e12 and big.A100_BANDWIDTH == 2039e9 and big.FP16_BYTES == 2
    assert big.decode_step(1).seconds * 1000 == pytest.approx(7.877, abs=5e-4)
    assert big.ridge() == pytest.approx(153.016, abs=1e-3) and big.first_arithmetic_limited_prompt() == 154


def test_draft_model_count_matches_the_listed_total_and_shares_the_big_models_vocabulary():
    """Guards a wrong config value: counting a separate output table (a second 263 million weights) would give 1.50 billion and a draft pass 21% too slow.
    The shared vocabulary is what lets the big model check the draft's words one for one."""
    assert sim.draft_parameter_count() == sim.DRAFT_PARAMS == sim.D_WEIGHT_FILES_PARAMS == 1_235_814_400
    separate = sim.DRAFT_PARAMS + sim.D_VOCAB * sim.D_HIDDEN                                  # a second table of the same size
    assert round(separate / 1e9, 2) == 1.50 and round((separate / sim.DRAFT_PARAMS - 1) * 100) == 21
    assert sim.D_VOCAB == big.VOCAB == 128_256
    assert sim.DRAFT_PARAMS * big.FP16_BYTES / 1e9 == pytest.approx(2.47, abs=5e-3)


def test_both_toy_models_have_an_answer_after_every_run_of_words_in_the_text():
    """Guards a gap in the tables: the text is a ring, so no context may be missing, and ties must break the same way every time."""
    for i in range(len(sim.WORDS)):
        context = RING[i:i + 2]
        assert sim.big_choice(context) == sim.big_ranked(context)[0] and sim.small_choice(context) == sim.small_ranked(context)[0]
    assert sim.big_choice(["xx", "kettle"]) == sim.small_choice(["kettle"]) == "."          # a pair never seen falls back to the last word alone
    assert sim.REPLY == sim.WORDS[2:2 + len(sim.REPLY)] and len(sim.REPLY) == 17


@pytest.mark.parametrize("k", [1, 2, 3, 4, 6, 8])
def test_speculative_text_equals_the_big_models_own_text(k):
    """The claim the lesson rests on, run: for every start in the text, keeping guesses up to the first disagreement and using the big model's word there
    gives exactly the word-by-word text of the big model. A draft that is always wrong changes only the speed, never the text."""
    for i in range(len(sim.WORDS)):
        prompt = RING[i:i + 2]
        plain = sim.greedy(sim.big_choice, prompt, 50)
        text, rounds = sim.speculative_greedy(prompt, 50, k)
        assert text == plain, (k, prompt)
        for r in rounds:
            assert len(r.guesses) == k and len(r.own) == k + 1 and 1 <= len(r.new) <= k + 1
            assert r.new[:r.kept] == r.guesses[:r.kept] and r.new[-1] == r.own[r.kept]       # kept guesses, then the big model's word
            assert (r.kept == k) == (len(r.new) == k + 1)                                    # all kept: the extra word came free
            if r.kept < k:
                assert r.guesses[r.kept] != r.own[r.kept]                                    # stopped at a real disagreement
        always_wrong = lambda context: next(w for w in sim.WORDS if w != sim.big_choice(context))
        junk, junk_rounds = sim.speculative_greedy(prompt, 50, k, draft=always_wrong)
        assert junk == plain and all(len(r.new) == 1 for r in junk_rounds)


def test_toy_pair_agreement_and_its_words_per_pass_against_the_formula():
    """Guards the lesson's note that real guesses are not independent: the toy pair agrees 68% of the time, the formula for that
    chance gives 2.69 words per pass at k = 4 and the toy measures 2.92. The gap is why the formula is called an assumption."""
    assert sim.agreement() == Fraction(28, 41)
    measured, formula = sim.toy_tokens_per_pass(4), sim.expected_tokens(float(sim.agreement()), 4)
    assert (round(measured, 2), round(formula, 2)) == (2.92, 2.69) and measured > 1.05 * formula


def test_sampling_rule_gives_the_big_models_odds_exactly():
    """Guards the rule's two halves: keep with chance min(1, p/q), redraw from max(0, p - q) scaled. Drop the redraw or use q and the odds change.
    The hand example is the one in the lesson; the loop is the same check on many invented pairs, in exact fractions."""
    p, q = sim.HAND_P, sim.HAND_Q
    assert {x: q[x] * sim.keep_chance(p[x], q[x]) for x in p} == {"A": Fraction(1, 5), "B": Fraction(3, 10), "C": Fraction(1, 5)}
    assert sim.keep_total(p, q) == Fraction(7, 10) and sim.residual(p, q) == {"A": 1, "B": 0, "C": 0}
    assert sim.output_distribution(p, q) == p
    rng = random.Random(7)
    for _ in range(30):
        size = rng.randint(2, 6)
        weights = lambda: [rng.randint(0, 9) + (i == 0) for i in range(size)]               # at least one nonzero
        wp, wq = weights(), weights()
        P = {i: Fraction(w, sum(wp)) for i, w in enumerate(wp)}
        Q = {i: Fraction(w, sum(wq)) for i, w in enumerate(wq)}
        assert sim.output_distribution(P, Q) == P
        assert sim.output_distribution(P, P) == P and sim.keep_total(P, P) == 1
    for rule in (sim.keep_total, sim.residual):                                              # odds for different tokens are a caller's mistake, not a KeyError
        with pytest.raises(AssertionError, match="same tokens"):
            rule({"A": Fraction(1)}, {"B": Fraction(1)})


@pytest.mark.parametrize("a, k", [(0.3, 1), (0.6, 2), (0.6, 4), (0.8, 4), (0.9, 8), (0.5, 12)])
def test_simulated_words_per_pass_match_the_formula(a, k):
    """Guards the formula against the process it describes: seeded rounds, each guess kept with chance a, within 1.5%. Also that the closed form is the sum
    1 + a + ... + a^k (the chance that the word at each place is reached) and ends at 1 and k + 1."""
    assert sim.simulate_tokens_per_pass(a, k, 100_000, seed=k) == pytest.approx(sim.expected_tokens(a, k), rel=0.015)
    assert sim.expected_tokens(a, k) == pytest.approx(sum(a**j for j in range(k + 1)), rel=1e-12)
    assert sim.expected_tokens(0, k) == 1 and sim.expected_tokens(1, k) == k + 1


def test_time_model_numbers_the_lesson_prints():
    """Pins the time and speed-up figures in lesson.md and the figures, so a change to the model cannot leave stale prose."""
    plain = big.decode_step(1).seconds
    assert plain * 1000 == pytest.approx(7.877, abs=5e-4) and DRAFT_MS == pytest.approx(1.212, abs=5e-4)
    assert sim.draft_pass().seconds / plain == pytest.approx(0.1539, abs=5e-5) == pytest.approx(sim.DRAFT_PARAMS / big.PARAMS, abs=5e-5)
    assert all(big.prefill(n).seconds == plain for n in range(1, 154)) and big.prefill(154).seconds > plain      # checking up to 153 inputs costs one plain pass
    assert sim.round_seconds(4) * 1000 == pytest.approx(12.725, abs=5e-4)
    assert [round(sim.expected_tokens(a, 4), 2) for a in (0.6, 0.8, 0.9)] == [2.31, 3.36, 4.10]
    assert [round(sim.speedup(a, 4), 2) for a in (0.3, 0.6, 0.8, 0.9)] == [0.88, 1.43, 2.08, 2.53]
    assert round(sim.round_seconds(4) * 1000 / sim.expected_tokens(0.8, 4), 2) == 3.79                          # ms per word at the opening setting
    assert [round(sim.break_even(k), 3) for k in (2, 4, 8)] == [0.247, 0.386, 0.554]
    assert [sim.best_k(a) for a in (0.5, 0.6, 0.8, 0.9, 0.95)] == [2, 2, 5, 8, 12]
    assert [round(sim.speedup(a, sim.best_k(a)), 2) for a in (0.6, 0.8, 0.9)] == [1.50, 2.08, 2.75]
    assert sim.speedup(0.386, 4) < 1 < sim.speedup(0.387, 4)


def test_speed_up_falls_as_users_are_added_and_ends_below_one():
    """Guards the claim that spare room is what pays for guessing: with several users the check covers users * (k + 1) inputs, so it stays
    free only up to the ridge. At k = 4 that is 31 users; past 154 inputs per word the speed-up is flat at E / (k * 0.154 + k + 1)."""
    assert [sim.users_when_checking_stops_being_free(k) for k in (2, 4, 6, 8)] == [52, 31, 22, 18]      # 6: 22 * 7 = 154 exactly, a tie counts as arithmetic
    assert [round(sim.speedup(0.8, 4, users=u), 2) for u in (1, 30, 32, 64, 128)] == [2.08, 2.08, 2.02, 1.24, 0.70]
    assert sim.users_when_it_stops_paying(0.8, 4) == 85
    limit = sim.expected_tokens(0.8, 4) / (4 * sim.DRAFT_PARAMS / big.PARAMS + 5)
    assert sim.speedup(0.8, 4, users=1000) == pytest.approx(limit, rel=1e-9) and round(limit, 2) == 0.60
    speeds = [sim.speedup(0.8, 4, users=u) for u in range(1, 400)]
    assert all(later <= earlier + 1e-12 for earlier, later in zip(speeds, speeds[1:]))


def test_plan_run_covers_the_reply_once_and_follows_the_chance():
    """Guards the widget's run: every word written once, rounds that stop at the end, none kept at 0%, all kept at 100%."""
    n = len(sim.REPLY)
    for percent in (0, 20, 50, 80, 100):
        for k in (1, 4, 8):
            for seed in range(1, 6):
                rounds = sim.plan_run(percent, k, seed)
                assert sum(r[4] for r in rounds) == n and [r[0] for r in rounds] == [sum(x[4] for x in rounds[:i]) for i in range(len(rounds))]
                for start, kept, rejected, own, words in rounds:
                    assert 0 <= kept <= min(k, words) and 1 <= words <= k + 1 and (own == 1 or words == kept)
    assert all(r[1] == 0 for r in sim.plan_run(0, 4, 1)) and len(sim.plan_run(0, 4, 1)) == n
    assert len(sim.plan_run(100, 4, 1)) == -(-n // 5) and all(r[2] == 0 for r in sim.plan_run(100, 4, 1))


def test_opening_run_is_typical_and_shows_each_kind_of_round():
    """Guards the first screen: the opening setting (80%, k = 4, seed 60) must show a rejection, a round with every guess kept and its free word, and a
    speed-up near the average (2.1 here), so the reader does not meet a lucky or unlucky run first. It must also move: a seed whose rejection draws all
    sat above 0.95 left the lanes frozen from 80% to 95%, so the chance slider looked dead. Here the speed-up changes at least once on that stretch."""
    view = sim.view(sim.OPENING["percent"], sim.OPENING["k"], DRAFT_MS, sim.OPENING["seed"])
    rounds = view["rounds"]
    speed = lambda v: v["plainMs"] * len(sim.REPLY) / (len(v["rounds"]) * v["roundMs"])
    assert speed(view) == pytest.approx(sim.speedup(0.8, 4), rel=0.05)
    assert len({round(speed(sim.view(pct, 4, DRAFT_MS, sim.OPENING["seed"])), 6) for pct in (80, 85, 90, 95)}) > 1
    assert any(r[2] for r in rounds) and any(r[1] == sim.OPENING["k"] and r[3] for r in rounds) and any(r[1] == 0 for r in rounds)
    assert view["plain"] == "17 big passes, 1 token each, 7.88 ms per token."
    assert view["spec"] == "This reply: 5 big passes, 3.4 tokens each, 3.74 ms per token, 2.1 times as fast."
    assert view["average"] == "Averaged over many replies at this setting: 3.36 tokens per big pass, 2.1 times as fast." and view["verdict"] == ""


def test_view_says_so_when_guessing_loses_and_rounds_ties_up():
    """Guards the low-speed message and the rounding: half up as JavaScript's toFixed does (2.25 is 2.3, not Python's 2.2). The verdict follows the averaged
    speed-up printed below it, whatever the cause: with 95% of guesses kept and the draft pass as slow as the big one it is still slower, and it must not
    blame the guesses. With every guess kept the same slow draft is only 0.2% behind, which prints as 1.0 times, so the verdict says about the same."""
    assert sim.fmt(2.25, 1) == "2.3" and sim.fmt(0.125, 2) == "0.13" and sim.fmt(7.8766, 2) == "7.88" and sim.fmt(12.0, 1) == "12.0"
    low = sim.view(20, 4, DRAFT_MS, 1)
    assert low["verdict"].startswith("On average, slower than plain decoding") and low["spec"].endswith("times as fast.")
    slow_draft = sim.view(95, 4, sim.DRAFT_STOPS_MS[-1], 1)
    assert slow_draft["verdict"] == "On average, slower than plain decoding: the draft passes cost more than the kept guesses save." and "0.9 times as fast" in slow_draft["average"]
    assert sim.view(100, 4, sim.DRAFT_STOPS_MS[-1], 1)["verdict"] == "On average, about the same speed as plain decoding."
    assert sim.view(100, 8, DRAFT_MS, 1)["verdict"] == "" and sim.DRAFT_STOPS_MS[sim.DEFAULT_DRAFT_STOP] == DRAFT_MS