"""Tests on the serve-spec recording of Lab 12, script 10: Qwen3-8B checked by Qwen3-0.6B drafting,
8 prompts per task, greedy, on one laptop (Apple M4 Max). Run: pytest test_spec_practice.py"""
from spec_practice import alpha_from_counts, one_round, speedup, tokens_per_pass

RATIO_16 = 0.1027  # a draft step costs about a tenth of a big step, at full size (315 vs 32 tokens/s)
RATIO_4 = 0.1789  # both models 4-bit, run 1


def test_one_round_the_video_shows():
    # "Thank you for" -> the draft guessed 4 tokens; the big model kept 2, then wrote " and"
    guesses = [" reaching", " out", ".", " I"]
    target = [" reaching", " out", " and", " for", " your"]
    assert one_round(guesses, target) == [" reaching", " out", " and"]  # 3 tokens from one pass
    assert one_round(["a", "b"], ["a", "b", "c"]) == ["a", "b", "c"]  # all right: k + 1
    assert one_round(["x", "b"], ["a", "b", "c"]) == ["a"]  # first wrong: the rest is thrown away too


def test_tokens_per_pass():
    assert tokens_per_pass(0, 4) == 1
    assert tokens_per_pass(1, 4) == 5
    assert round(tokens_per_pass(0.6, 2), 2) == 1.96
    # stories: all six guesses right only about 1 time in 20, so k = 6 adds little over k = 4
    assert round(0.6**6, 3) == 0.047
    assert tokens_per_pass(0.6, 6) - tokens_per_pass(0.6, 4) < 0.2


def test_alpha_from_recorded_counts():
    assert round(alpha_from_counts(202, 303, 2), 2) == 1.00  # ticket JSON, k = 2: 202 of 303 tokens from the draft
    assert round(alpha_from_counts(526, 1024, 2), 2) == 0.64  # support replies
    assert round(alpha_from_counts(504, 1024, 2), 2) == 0.60  # short stories
    assert round(alpha_from_counts(569, 1024, 6), 2) == 0.56  # stories at 4 bits, k = 6: guesses barely worse
    assert alpha_from_counts(0, 100, 2) < 0.001


def test_formula_against_measured():
    # predicted at k = 2, full size; measured (median of 3 runs) was 1.37x, 1.46x, 1.46x
    assert round(speedup(alpha_from_counts(202, 303, 2), 2, RATIO_16), 2) == 2.49
    assert round(speedup(alpha_from_counts(526, 1024, 2), 2, RATIO_16), 2) == 1.71
    assert round(speedup(alpha_from_counts(504, 1024, 2), 2, RATIO_16), 2) == 1.63
    # 4-bit stories at k = 6: the formula still says faster (about 1.1x); measured 0.47x, half speed,
    # because at 4 bits checking 7 tokens costs 2.5x a single step, not 1x as the formula assumes
    assert round(speedup(alpha_from_counts(569, 1024, 6), 6, RATIO_4), 1) == 1.1
    assert speedup(0.6, 2, 0) > speedup(0.6, 2, RATIO_16)  # drafting isn't free