"""Tests for the roofline / MFU / latency workbench.

    pytest test_lab.py -v                       # your lab.py
    LAB_MODULE=solution pytest test_lab.py -v   # the reference (must be green)
"""

import importlib
import math
import os

import pytest

lab = importlib.import_module(os.environ.get("LAB_MODULE", "lab"))


# ======================================================================================
# 1. The roofline
# ======================================================================================

def test_arithmetic_intensity_is_flops_per_byte():
    assert lab.arithmetic_intensity(1000.0, 100.0) == pytest.approx(10.0)


def test_arithmetic_intensity_rejects_zero_bytes():
    with pytest.raises(ValueError):
        lab.arithmetic_intensity(1000.0, 0.0)
    with pytest.raises(ValueError):
        lab.arithmetic_intensity(-1.0, 100.0)


def test_ridge_point_matches_the_spec_sheet():
    peak, _hbm, bw, _w = lab.HARDWARE["H100"]
    assert lab.ridge_point("H100") == pytest.approx(peak / bw)
    assert 250 < lab.ridge_point("H100") < 350        # ~296 FLOP/byte


def test_ridge_point_rejects_unknown_chip():
    with pytest.raises(ValueError):
        lab.ridge_point("Babbage Engine")


def test_roofline_is_bandwidth_limited_below_the_ridge():
    """BOUNDARY: below the ridge, throughput is bandwidth * intensity — the sloped roof."""
    _peak, _hbm, bw, _w = lab.HARDWARE["H100"]
    ai = lab.ridge_point("H100") / 2
    assert lab.roofline_throughput("H100", ai) == pytest.approx(bw * ai)


def test_roofline_is_capped_at_peak_above_the_ridge():
    """BOUNDARY: above the ridge you cannot exceed peak, no matter the intensity."""
    peak, _hbm, _bw, _w = lab.HARDWARE["H100"]
    assert lab.roofline_throughput("H100", 1e9) == pytest.approx(peak)


def test_roofline_is_continuous_at_the_ridge():
    """The two roof segments must meet exactly at the ridge point."""
    peak, _hbm, bw, _w = lab.HARDWARE["H100"]
    r = lab.ridge_point("H100")
    assert lab.roofline_throughput("H100", r) == pytest.approx(peak)
    assert lab.roofline_throughput("H100", r) == pytest.approx(bw * r)


def test_roofline_throughput_rejects_negative_intensity():
    with pytest.raises(ValueError):
        lab.roofline_throughput("H100", -1.0)


def test_big_matmul_is_compute_bound_and_decode_is_memory_bound():
    """The two ends of the roofline, and the whole reason it matters."""
    big = lab.roofline_report("H100", 2 * 8192 ** 3, 3 * 8192 * 8192 * 2)
    assert big["bound_by"] == "compute"
    assert big["lever"] == "reduce FLOPs"

    decode = lab.roofline_report("H100", 2 * 70e9, 70e9 * 2)
    assert decode["bound_by"] == "memory"
    assert decode["lever"] == "reduce bytes moved"
    assert decode["fraction_of_peak"] < 0.05


# ======================================================================================
# 2. MFU
# ======================================================================================

def test_mfu_of_a_perfectly_utilized_chip_is_one():
    """BOUNDARY: doing exactly peak FLOPs for exactly one second on one chip."""
    peak, _hbm, _bw, _w = lab.HARDWARE["H100"]
    assert lab.mfu(peak, 1.0, "H100", 1) == pytest.approx(1.0)


def test_mfu_halves_when_the_run_takes_twice_as_long():
    peak, _hbm, _bw, _w = lab.HARDWARE["H100"]
    assert lab.mfu(peak, 2.0, "H100", 1) == pytest.approx(0.5)


def test_mfu_halves_when_you_double_the_chips():
    """DIMENSIONAL ANALYSIS: same work, twice the hardware, half the utilization."""
    peak, _hbm, _bw, _w = lab.HARDWARE["H100"]
    assert lab.mfu(peak, 1.0, "H100", 2) == pytest.approx(0.5)


def test_mfu_rejects_nonpositive():
    with pytest.raises(ValueError):
        lab.mfu(1e15, 0.0, "H100", 1)
    with pytest.raises(ValueError):
        lab.mfu(1e15, 1.0, "H100", 0)


def test_hfu_is_always_at_least_mfu():
    """HFU counts activation recomputation as useful work. It can never be lower.
    ALWAYS ask which number you are being shown — it is ~33% for free."""
    args = (6 * 70e9 * 1e9, 3600.0, "H100", 1000)
    assert lab.hfu(*args) >= lab.mfu(*args)
    assert lab.hfu(*args) == pytest.approx(lab.mfu(*args) * 8 / 6)


def test_hfu_equals_mfu_with_no_recomputation():
    """BOUNDARY: recompute_factor of 1.0 means no checkpointing, so the two agree."""
    args = (1e18, 10.0, "H100", 4)
    assert lab.hfu(*args, recompute_factor=1.0) == pytest.approx(lab.mfu(*args))


def test_hfu_rejects_impossible_recompute_factor():
    with pytest.raises(ValueError):
        lab.hfu(1e18, 10.0, "H100", 4, recompute_factor=0.5)


def test_mfu_budget_fractions_sum_to_one():
    b = lab.mfu_budget(100, 45, 60, 50, 25)
    total = sum(b[k] for k in ("matmul_fraction", "vector_fraction", "memory_fraction",
                               "comms_fraction", "optimizer_fraction"))
    assert total == pytest.approx(1.0)
    assert b["mfu"] == pytest.approx(b["matmul_fraction"])


def test_mfu_budget_reproduces_a_realistic_35_percent():
    """35% MFU is an ACCOUNTING IDENTITY, not a failure grade."""
    b = lab.mfu_budget(100, 45, 60, 50, 25)
    assert b["mfu"] == pytest.approx(100 / 280)
    assert 0.30 < b["mfu"] < 0.40


def test_mfu_budget_identifies_the_biggest_lever():
    """The breakdown is an agenda: it tells you what to fix next."""
    assert lab.mfu_budget(100, 5, 5, 200, 5)["biggest_lever"] == "comms"
    assert lab.mfu_budget(100, 5, 200, 5, 5)["biggest_lever"] == "memory"
    assert lab.mfu_budget(100, 200, 5, 5, 5)["biggest_lever"] == "vector"


def test_perfect_mfu_when_only_matmul_runs():
    """BOUNDARY: a pure matmul loop with no memory reads is the only way to hit 100% —
    and that is not a neural network."""
    assert lab.mfu_budget(100, 0, 0, 0, 0)["mfu"] == pytest.approx(1.0)


def test_mfu_budget_validates():
    with pytest.raises(ValueError):
        lab.mfu_budget(-1, 0, 0, 0, 0)
    with pytest.raises(ValueError):
        lab.mfu_budget(0, 0, 0, 0, 0)


# ======================================================================================
# 3. Prefill vs decode
# ======================================================================================

def test_prefill_scales_linearly_in_tokens_and_params():
    base = lab.prefill_seconds(8192, 70e9, "TPU v5e", 1)
    assert lab.prefill_seconds(16384, 70e9, "TPU v5e", 1) == pytest.approx(2 * base)
    assert lab.prefill_seconds(8192, 140e9, "TPU v5e", 1) == pytest.approx(2 * base)


def test_prefill_scales_inversely_with_chips():
    base = lab.prefill_seconds(8192, 70e9, "TPU v5e", 1)
    assert lab.prefill_seconds(8192, 70e9, "TPU v5e", 16) == pytest.approx(base / 16)


def test_prefill_reproduces_the_talks_number():
    """One v5e chip, 8k incremental prefill of a 70B model, compute-bound: ~5.8 s.
    The talk quotes ~5.7 s for the same setup."""
    s = lab.prefill_seconds(8192, 70e9, "TPU v5e", 1, mfu_frac=1.0)
    assert 5.0 < s < 6.5


def test_a_4x4_station_brings_prefill_under_the_half_second_limit():
    """The talk's conclusion: a 4x4 = 16-chip v5e prefill station."""
    assert lab.prefill_seconds(8192, 70e9, "TPU v5e", 16, mfu_frac=1.0) < 0.5


def test_prefill_validates():
    with pytest.raises(ValueError):
        lab.prefill_seconds(8192, 70e9, "TPU v5e", 1, mfu_frac=1.5)
    with pytest.raises(ValueError):
        lab.prefill_seconds(0, 70e9, "TPU v5e", 1)


def test_decode_is_driven_by_bytes_not_flops():
    """Decode re-reads every weight per token — halving the bytes per parameter
    halves the time, which is the entire case for quantization."""
    bf16 = lab.decode_seconds(128, 70e9, "TPU v5e", 1, bytes_per_param=2)
    int8 = lab.decode_seconds(128, 70e9, "TPU v5e", 1, bytes_per_param=1)
    assert int8 == pytest.approx(bf16 / 2)


def test_decode_dominates_prefill_at_batch_one():
    """The correction to the naive reading of the napkin: at batch 1, generating 128
    tokens costs MORE than prefilling 8192 of them."""
    p = lab.prefill_seconds(8192, 70e9, "TPU v5e", 16)
    d = lab.decode_seconds(128, 70e9, "TPU v5e", 16)
    assert d > p


def test_decode_validates():
    with pytest.raises(ValueError):
        lab.decode_seconds(128, 70e9, "TPU v5e", 1, bandwidth_efficiency=0)
    with pytest.raises(ValueError):
        lab.decode_seconds(128, 70e9, "TPU v5e", 0)


def test_interactive_latency_sums_its_parts():
    r = lab.interactive_latency(8192, 128, 70e9, "TPU v5e", 16, scaffolding_s=0.25)
    assert r["total_s"] == pytest.approx(r["prefill_s"] + r["decode_s"] + 0.25)
    assert 0.0 < r["prefill_fraction"] < 1.0


def test_interactive_latency_rejects_negative_scaffolding():
    with pytest.raises(ValueError):
        lab.interactive_latency(8192, 128, 70e9, "TPU v5e", 16, scaffolding_s=-1)


def test_more_chips_eventually_meet_the_budget():
    n = lab.chips_for_latency_budget(1.0, 8192, 128, 70e9, "TPU v5e")
    assert n is not None
    r = lab.interactive_latency(8192, 128, 70e9, "TPU v5e", n)
    assert r["total_s"] <= 1.0
    # ...and one step smaller must NOT meet it (it is the smallest such power of two).
    if n > 1:
        assert lab.interactive_latency(8192, 128, 70e9, "TPU v5e",
                                       n // 2)["total_s"] > 1.0


def test_a_smaller_model_needs_fewer_chips():
    """THE CONCLUSION OF THIS PHASE. Halving N halves both prefill and decode, so the
    chip count drops. This is the economic case for Flash-class models."""
    big = lab.chips_for_latency_budget(1.0, 8192, 128, 70e9, "TPU v5e")
    small = lab.chips_for_latency_budget(1.0, 8192, 128, 8e9, "TPU v5e")
    assert small < big


def test_budget_below_scaffolding_is_impossible():
    """BOUNDARY: you cannot beat a budget smaller than your fixed overhead."""
    with pytest.raises(ValueError):
        lab.chips_for_latency_budget(0.2, 8192, 128, 70e9, "TPU v5e",
                                     scaffolding_s=0.25)


def test_impossible_budget_returns_none():
    """No amount of hardware fixes a budget that is out of reach — the answer is a
    smaller model, and the function must say so rather than looping forever."""
    assert lab.chips_for_latency_budget(0.2501, 8192, 128, 70e9, "TPU v5e",
                                        scaffolding_s=0.25, max_chips=8) is None


def test_weights_fit_chips_rounds_up():
    """A 70B model in bf16 is 140 GB; a v5e has 16 GB -> 9 chips just to hold it."""
    assert lab.weights_fit_chips(70e9, "TPU v5e") == 9
    assert lab.weights_fit_chips(8e9, "TPU v5e") == 1


# ======================================================================================
# 4. Shape co-design
# ======================================================================================

def test_tile_efficiency_is_one_on_tile_multiples():
    for d in (128, 4096, 8192, 11008):
        assert lab.tile_efficiency(d) == pytest.approx(1.0)


def test_tile_efficiency_penalizes_off_tile_dimensions():
    assert lab.tile_efficiency(4097) < 1.0
    assert lab.tile_efficiency(4097) == pytest.approx(4097 / 4224)


def test_tile_efficiency_worst_case_is_just_over_a_boundary():
    """BOUNDARY: one element past a tile boundary wastes almost a whole tile."""
    assert lab.tile_efficiency(129, tile=128) == pytest.approx(129 / 256)


def test_tile_efficiency_validates():
    with pytest.raises(ValueError):
        lab.tile_efficiency(0)
    with pytest.raises(ValueError):
        lab.tile_efficiency(128, tile=0)


def test_kv_cache_is_linear_in_every_dimension():
    base = lab.kv_cache_bytes(80, 8, 128, 8192, 1)
    assert lab.kv_cache_bytes(160, 8, 128, 8192, 1) == 2 * base
    assert lab.kv_cache_bytes(80, 16, 128, 8192, 1) == 2 * base
    assert lab.kv_cache_bytes(80, 8, 128, 8192, 2) == 2 * base


def test_gqa_ratio_equals_the_group_size():
    """THE inference co-design lever, and it is exactly the group factor."""
    g = lab.gqa_saving(80, 64, 8, 128, 8192, 32)
    assert g["group_size"] == 8
    assert g["ratio"] == pytest.approx(8.0)
    assert g["bytes_saved"] == g["mha_bytes"] - g["gqa_bytes"]


def test_mqa_is_the_extreme_case():
    """BOUNDARY: one KV head shared by all queries."""
    g = lab.gqa_saving(80, 64, 1, 128, 8192, 32)
    assert g["ratio"] == pytest.approx(64.0)


def test_mha_is_the_identity_case():
    """BOUNDARY: n_kv == n_query is plain MHA — no saving at all."""
    g = lab.gqa_saving(80, 64, 64, 128, 8192, 32)
    assert g["ratio"] == pytest.approx(1.0)
    assert g["bytes_saved"] == 0


def test_gqa_rejects_invalid_head_configuration():
    with pytest.raises(ValueError):
        lab.gqa_saving(80, 8, 64, 128, 8192, 32)          # more kv than query heads
    with pytest.raises(ValueError):
        lab.gqa_saving(80, 64, 7, 128, 8192, 32)          # does not divide evenly


def test_batching_raises_arithmetic_intensity():
    """Batching is the ONLY lever that moves decode rightward on the roofline, because
    it amortizes the weight read across the batch."""
    ais = [lab.decode_batch_intensity(70e9, 80, 8, 128, 8192, b)
           for b in (1, 8, 64, 256, 1024)]
    assert ais == sorted(ais)


def test_decode_stays_memory_bound_even_at_huge_batch():
    """The fact that explains GQA, quantization and speculative decoding: decode never
    reaches the ridge at realistic batch sizes."""
    ridge = lab.ridge_point("H100")
    assert lab.decode_batch_intensity(70e9, 80, 8, 128, 8192, 1024) < ridge


def test_depth_vs_width_reports_serial_cost():
    """Depth is SERIAL — it costs decode latency directly. Width is parallel."""
    deep = lab.depth_vs_width(70e9, d_model=4096, n_layers=160)
    wide = lab.depth_vs_width(70e9, d_model=8192, n_layers=40)
    assert deep["serial_steps_per_token"] > wide["serial_steps_per_token"]
    assert wide["aspect_ratio"] > deep["aspect_ratio"]


def test_depth_vs_width_validates():
    with pytest.raises(ValueError):
        lab.depth_vs_width(70e9, d_model=4096, n_layers=0)


# ======================================================================================
# 5. Determinism
# ======================================================================================

def test_everything_is_deterministic():
    """No clocks, no randomness — same inputs, byte-identical output."""
    a = lab.roofline_report("H100", 2 * 70e9, 70e9 * 2)
    b = lab.roofline_report("H100", 2 * 70e9, 70e9 * 2)
    assert repr(a) == repr(b)
    c = lab.interactive_latency(8192, 128, 70e9, "TPU v5e", 16)
    d = lab.interactive_latency(8192, 128, 70e9, "TPU v5e", 16)
    assert repr(c) == repr(d)