"""Deterministic tests for spec_decode_sim.py — closed-form edge cases,
monotonicity properties, hand-computed queueing/KV values, and the
colocated-vs-disaggregated crossover directions.

Run:  cd <this dir> && python -m pytest test_sim.py -q
"""
import math

import pytest

from spec_decode_sim import (
    BW_GRID,
    GB,
    LINKS,
    MODELS,
    SLO,
    NodeSpec,
    SpecDecodeConfig,
    Workload,
    best_disagg_split,
    colocated_metrics,
    crossover_bandwidth,
    disagg_metrics,
    erlang_c,
    expected_tokens,
    kv_transfer_s,
    mmc_wait_mean_s,
    mmc_wait_prob_le_s,
    optimal_k,
    speedup,
)

# -------------------------------------------------------------------------
# Part 1 — speculative decoding closed forms
# -------------------------------------------------------------------------


class TestExpectedTokens:
    def test_alpha_zero_always_one_token(self):
        # Every draft rejected -> the verification pass still emits exactly
        # one token from the target distribution.
        for k in range(0, 9):
            assert expected_tokens(0.0, k) == pytest.approx(1.0)

    def test_alpha_one_gives_k_plus_one(self):
        for k in range(0, 9):
            assert expected_tokens(1.0, k) == pytest.approx(k + 1)

    def test_matches_direct_geometric_sum(self):
        # Closed form must equal sum_{i=0..k} alpha^i.
        for alpha in (0.25, 0.5, 0.8, 0.99):
            for k in (1, 3, 8):
                direct = sum(alpha**i for i in range(k + 1))
                assert expected_tokens(alpha, k) == pytest.approx(direct)

    def test_monotone_in_alpha(self):
        alphas = [i / 20 for i in range(21)]
        vals = [expected_tokens(a, 4) for a in alphas]
        assert all(b >= a for a, b in zip(vals, vals[1:]))

    def test_monotone_in_k(self):
        vals = [expected_tokens(0.7, k) for k in range(0, 10)]
        assert all(b > a for a, b in zip(vals, vals[1:]))


class TestSpeedup:
    def test_spec_decode_can_hurt(self):
        # Low acceptance + expensive draft -> spec decode is a SLOWDOWN.
        cfg = SpecDecodeConfig(alpha=0.3, draft_cost=0.5, verify_overhead=0.05)
        assert speedup(cfg, 4) < 1.0

    def test_good_regime_speeds_up(self):
        cfg = SpecDecodeConfig(alpha=0.8, draft_cost=0.05,
                               verify_overhead=0.05)
        assert speedup(cfg, 4) > 1.5

    def test_optimal_k_monotone_in_alpha(self):
        ks = []
        for alpha in (0.6, 0.7, 0.8, 0.9):
            cfg = SpecDecodeConfig(alpha=alpha, draft_cost=0.10,
                                   verify_overhead=0.05)
            ks.append(optimal_k(cfg, k_max=8)[0])
        assert ks == sorted(ks)          # non-decreasing
        assert ks[-1] > ks[0]            # and actually grows

    def test_speedup_decays_monotonically_with_utilization(self):
        utils = [0.0, 0.2, 0.4, 0.6, 0.8, 0.9, 0.95]
        vals = [
            speedup(SpecDecodeConfig(alpha=0.8, draft_cost=0.10,
                                     verify_overhead=0.05, utilization=u), 4)
            for u in utils
        ]
        assert all(b < a for a, b in zip(vals, vals[1:]))

    def test_speedup_collapses_below_one_at_saturation(self):
        cfg = SpecDecodeConfig(alpha=0.8, draft_cost=0.10,
                               verify_overhead=0.05, utilization=0.99)
        _, best = optimal_k(cfg, k_max=8)
        assert best < 1.0                # even the best k hurts


# -------------------------------------------------------------------------
# Part 2 — Erlang-C / M/M/c queueing
# -------------------------------------------------------------------------


class TestErlangC:
    def test_low_load_no_waiting(self):
        assert erlang_c(2, 0.01) < 1e-2
        assert mmc_wait_mean_s(2, lam=0.01, mu=1.0) < 1e-2
        assert mmc_wait_prob_le_s(2, lam=0.01, mu=1.0, t=0.5) > 0.99

    def test_saturation_wait_explodes(self):
        # rho -> 1: mean wait grows without bound; at rho >= 1 it is inf.
        assert mmc_wait_mean_s(2, lam=1.999, mu=1.0) > 100.0
        assert math.isinf(mmc_wait_mean_s(2, lam=2.0, mu=1.0))

    def test_mm2_hand_computed_value(self):
        # M/M/2, lambda=1.5, mu=1 -> a=1.5, rho=0.75.
        # Erlang-C = (a^2/2!) / ((1-rho)(1+a) + a^2/2!) = 1.125/1.75 = 9/14.
        # E[W_q]  = C/(c*mu - lambda) = (9/14)/0.5 = 9/7 ~ 1.2857 s.
        assert erlang_c(2, 1.5) == pytest.approx(9 / 14, rel=1e-9)
        assert mmc_wait_mean_s(2, 1.5, 1.0) == pytest.approx(9 / 7, rel=1e-9)


# -------------------------------------------------------------------------
# Part 2 — KV transfer arithmetic
# -------------------------------------------------------------------------


class TestKVTransfer:
    def test_llama70b_kv_bytes_per_token(self):
        # 2 * 80 layers * 8 kv_heads * 128 head_dim * 2 B = 327,680 B/token.
        assert MODELS["llama-70b-gqa8"].kv_bytes_per_token() == 327680

    def test_transfer_time_linear_in_prompt(self):
        m = MODELS["llama-70b-gqa8"]
        t1 = kv_transfer_s(m, 1000, LINKS["nvlink"])
        t2 = kv_transfer_s(m, 2000, LINKS["nvlink"])
        assert t2 == pytest.approx(2 * t1, rel=1e-12)

    def test_llama70b_1k_tokens_hand_calc(self):
        m = MODELS["llama-70b-gqa8"]
        # ~0.33 GB of KV for a 1k-token prompt in FP16:
        assert m.kv_bytes(1000) / GB == pytest.approx(0.32768, rel=1e-9)
        # NVLink-class 400 GB/s -> ~0.8 ms.
        t_nvlink = kv_transfer_s(m, 1000, 400 * GB)
        assert t_nvlink * 1e3 == pytest.approx(0.8192, rel=1e-9)
        assert t_nvlink * 1e3 == pytest.approx(0.8, rel=0.15)
        # 10G TCP 1.25 GB/s -> ~a quarter second.
        t_tcp = kv_transfer_s(m, 1000, 1.25 * GB)
        assert t_tcp * 1e3 == pytest.approx(262.144, rel=1e-9)
        assert t_tcp * 1e3 == pytest.approx(256, rel=0.15)


# -------------------------------------------------------------------------
# Part 2 — colocated vs disaggregated crossover directions
# -------------------------------------------------------------------------

NODE = NodeSpec()  # prefill 8000 tok/s, decode 4000 tok/s, base TPOT 25 ms


class TestColocated:
    def test_tpot_inflation_monotone_in_prefill_load(self):
        slo = SLO(ttft_ms=5000, tpot_ms=1000)
        tpots = []
        for lam in (1.0, 4.0, 8.0, 12.0):
            m = colocated_metrics(Workload(lam, 2048, 256), NODE, 8, slo)
            tpots.append(m.tpot_ms)
        assert all(b > a for a, b in zip(tpots, tpots[1:]))

    def test_overload_gives_zero_goodput(self):
        # lambda*O = 100*1024 tok/s >> 8 nodes * 4000 tok/s decode capacity.
        slo = SLO(ttft_ms=5000, tpot_ms=1000)
        m = colocated_metrics(Workload(100.0, 512, 1024), NODE, 8, slo)
        assert not m.stable
        assert m.goodput_rps == 0.0
        assert math.isinf(m.ttft_mean_ms)


class TestCrossover:
    def test_decode_heavy_fast_link_tight_tpot_disagg_wins(self):
        # Agent-like: long outputs, tight TPOT SLO, NVLink-class transfer.
        # Colocated prefill interference (rho_p=0.16 -> TPOT 27 ms) blows the
        # 26 ms SLO; disagg isolates decode at 25 ms and wins outright.
        w = Workload(10.0, 1024, 2048)
        slo = SLO(ttft_ms=1000, tpot_ms=26)
        colo = colocated_metrics(w, NODE, 8, slo)
        n_p, dis = best_disagg_split(w, NODE, 8, MODELS["llama-70b-gqa8"],
                                     LINKS["nvlink"], slo)
        assert colo.goodput_rps == 0.0            # fails the TPOT SLO
        assert dis.goodput_rps > 9.0              # ~all of lambda=10
        assert dis.goodput_rps >= colo.goodput_rps
        assert dis.tpot_ms <= slo.tpot_ms < colo.tpot_ms

    def test_tiny_scale_slow_link_colocated_wins(self):
        # 2 nodes, 4k-token prompts over 10G TCP: the ~1.07 s KV transfer
        # alone eats the 1.5 s TTFT budget; pooling both nodes colocated
        # keeps TTFT ~0.52 s and the loose TPOT SLO absorbs interference.
        w = Workload(0.5, 4096, 256)
        slo = SLO(ttft_ms=1500, tpot_ms=40)
        colo = colocated_metrics(w, NODE, 2, slo)
        _, dis = best_disagg_split(w, NODE, 2, MODELS["llama-70b-gqa8"],
                                   LINKS["tcp"], slo)
        assert colo.goodput_rps > 0.45            # ~all of lambda=0.5
        assert dis.goodput_rps == 0.0             # transfer blows TTFT
        assert colo.goodput_rps > dis.goodput_rps

    def test_crossover_bandwidth_found_for_tight_tpot_workload(self):
        # Same agent-like workload: colocated is SLO-infeasible, so disagg
        # should win from some (low) bandwidth upward, and losing the link
        # entirely below the KV budget flips it back.
        w = Workload(10.0, 1024, 2048)
        slo = SLO(ttft_ms=1000, tpot_ms=26)
        bw = crossover_bandwidth(w, NODE, 8, MODELS["llama-70b-gqa8"], slo)
        assert bw is not None
        assert bw in BW_GRID
        # At the crossover bandwidth disagg must genuinely beat colocated.
        colo = colocated_metrics(w, NODE, 8, slo)
        _, dis = best_disagg_split(w, NODE, 8, MODELS["llama-70b-gqa8"],
                                   bw, slo)
        assert dis.goodput_rps > colo.goodput_rps

    def test_disagg_infeasible_when_decode_pool_too_small(self):
        # Force a split whose decode pool cannot absorb lambda*O.
        w = Workload(10.0, 1024, 2048)          # needs 20480 decode tok/s
        slo = SLO(ttft_ms=1000, tpot_ms=26)
        m = disagg_metrics(w, NODE, n_prefill=7, n_decode=1,
                           model=MODELS["llama-70b-gqa8"],
                           link_bandwidth=LINKS["nvlink"], slo=slo)
        assert not m.stable
        assert m.goodput_rps == 0.0