"""Tests for the Mixture-of-Experts layer.

    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 random

import pytest

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

D_MODEL, D_FF, N_EXPERTS, TOP_K = 8, 16, 8, 2


def _tokens(n=64, d=D_MODEL, seed=1):
    rng = random.Random(seed)
    return [[rng.gauss(0, 1) for _ in range(d)] for _ in range(n)]


def _router(n_experts=N_EXPERTS, d=D_MODEL, seed=0, scale=0.5):
    rng = random.Random(seed)
    return [[rng.gauss(0, scale) for _ in range(d)] for _ in range(n_experts)]


def _experts(n=N_EXPERTS):
    return [lab.make_expert(D_MODEL, D_FF, seed=i) for i in range(n)]


# ======================================================================================
# 1. Numerical primitives
# ======================================================================================

def test_softmax_sums_to_one_and_is_positive():
    p = lab.softmax([1.0, 2.0, 3.0, -1.0])
    assert sum(p) == pytest.approx(1.0)
    assert all(x > 0 for x in p)


def test_softmax_survives_huge_logits():
    """NUMERICAL STABILITY: without max-subtraction, exp(1000) is inf and the whole
    layer produces NaN. Router logits really do get large during training."""
    p = lab.softmax([1000.0, 1001.0, 999.0])
    assert sum(p) == pytest.approx(1.0)
    assert all(math.isfinite(x) for x in p)
    assert p[1] > p[0] > p[2]


def test_softmax_is_shift_invariant():
    a = lab.softmax([1.0, 2.0, 3.0])
    b = lab.softmax([101.0, 102.0, 103.0])
    for x, y in zip(a, b):
        assert x == pytest.approx(y)


def test_softmax_rejects_empty():
    with pytest.raises(ValueError):
        lab.softmax([])


def test_logsumexp_matches_the_naive_form_when_safe():
    xs = [0.5, -1.0, 2.0]
    assert lab.logsumexp(xs) == pytest.approx(math.log(sum(math.exp(x) for x in xs)))


def test_logsumexp_survives_huge_logits():
    assert math.isfinite(lab.logsumexp([1000.0, 1001.0]))


def test_matvec_shape_check():
    with pytest.raises(ValueError):
        lab.matvec([[1.0, 2.0]], [1.0, 2.0, 3.0])


# ======================================================================================
# 2. The router
# ======================================================================================

def test_route_returns_exactly_top_k():
    r = _router()
    picked = lab.route_token(_tokens(1)[0], r, top_k=3)
    assert len(picked) == 3
    assert len({e for e, _ in picked}) == 3          # no duplicates


def test_renormalized_gates_sum_to_one():
    picked = lab.route_token(_tokens(1)[0], _router(), top_k=2, renormalize=True)
    assert sum(g for _, g in picked) == pytest.approx(1.0)


def test_raw_gates_do_not_sum_to_one():
    """The design choice: without renormalization the layer output is silently scaled
    down by whatever probability mass went to the experts you did NOT pick."""
    picked = lab.route_token(_tokens(1)[0], _router(), top_k=2, renormalize=False)
    assert sum(g for _, g in picked) < 1.0


def test_router_picks_the_highest_scoring_experts():
    r = _router()
    tok = _tokens(1)[0]
    probs = lab.softmax(lab.router_logits(tok, r))
    best_two = sorted(range(len(probs)), key=lambda i: (-probs[i], i))[:2]
    assert [e for e, _ in lab.route_token(tok, r, top_k=2)] == best_two


def test_gates_are_ordered_highest_first():
    picked = lab.route_token(_tokens(1)[0], _router(), top_k=4)
    gates = [g for _, g in picked]
    assert gates == sorted(gates, reverse=True)


def test_top_k_equal_to_n_experts_selects_everything():
    """BOUNDARY: k == E is a dense model wearing an MoE costume."""
    picked = lab.route_token(_tokens(1)[0], _router(), top_k=N_EXPERTS)
    assert sorted(e for e, _ in picked) == list(range(N_EXPERTS))
    assert sum(g for _, g in picked) == pytest.approx(1.0)


def test_top_k_of_one_is_switch_routing():
    """BOUNDARY: k == 1 (Switch Transformer). The single gate must be exactly 1.0."""
    picked = lab.route_token(_tokens(1)[0], _router(), top_k=1)
    assert len(picked) == 1
    assert picked[0][1] == pytest.approx(1.0)


@pytest.mark.parametrize("k", [0, -1, N_EXPERTS + 1])
def test_route_rejects_bad_top_k(k):
    with pytest.raises(ValueError):
        lab.route_token(_tokens(1)[0], _router(), top_k=k)


def test_routing_is_deterministic():
    """DETERMINISM: identical inputs must route identically, including ties."""
    r, toks = _router(), _tokens(16)
    assert lab.route_batch(toks, r, TOP_K) == lab.route_batch(toks, r, TOP_K)


def test_route_batch_shapes_line_up():
    toks = _tokens(10)
    a, g, p = lab.route_batch(toks, _router(), TOP_K)
    assert len(a) == len(g) == len(p) == 10
    assert all(len(row) == TOP_K for row in a)
    assert all(len(row) == N_EXPERTS for row in p)


# ======================================================================================
# 3. Auxiliary losses
# ======================================================================================

def test_balanced_routing_gives_aux_loss_of_exactly_one():
    """THE CALIBRATION POINT. Perfect balance -> E * E * (1/E) * (1/E) = 1.0.
    That is why 1.0 is the number on the dashboard, and larger is worse."""
    assign = [[i % N_EXPERTS] for i in range(N_EXPERTS * 8)]
    probs = [[1.0 / N_EXPERTS] * N_EXPERTS for _ in assign]
    assert lab.load_balance_loss(assign, probs, N_EXPERTS) == pytest.approx(1.0)


def test_collapsed_routing_is_penalized_far_more():
    assign = [[0] for _ in range(64)]
    probs = [[0.93] + [0.01] * (N_EXPERTS - 1) for _ in assign]
    collapsed = lab.load_balance_loss(assign, probs, N_EXPERTS)
    assert collapsed > 5.0


def test_aux_loss_is_monotone_in_imbalance():
    """More skew must mean more loss — otherwise the gradient points the wrong way."""
    probs = [[1.0 / N_EXPERTS] * N_EXPERTS for _ in range(64)]
    losses = []
    for skew in (0, 16, 32, 48, 64):
        assign = [[0] if i < skew else [i % N_EXPERTS] for i in range(64)]
        losses.append(lab.load_balance_loss(assign, probs, N_EXPERTS))
    assert losses == sorted(losses)


def test_aux_loss_validates():
    probs = [[0.25] * 4]
    with pytest.raises(ValueError):
        lab.load_balance_loss([], [], 4)
    with pytest.raises(ValueError):
        lab.load_balance_loss([[0]], probs, 0)
    with pytest.raises(ValueError):
        lab.load_balance_loss([[9]], probs, 4)          # expert id out of range
    with pytest.raises(ValueError):
        lab.load_balance_loss([[0], [1]], probs, 4)     # length mismatch


def test_z_loss_grows_with_logit_magnitude():
    """The whole point: keep router logits small so the softmax does not saturate."""
    base = [[1.0, 2.0, 0.5, -1.0]]
    small = lab.router_z_loss(base)
    big = lab.router_z_loss([[z * 10 for z in base[0]]])
    assert big > small > 0


def test_z_loss_rejects_empty():
    with pytest.raises(ValueError):
        lab.router_z_loss([])


def test_utilization_detects_dead_experts_and_collapse():
    u = lab.expert_utilization([[0] for _ in range(32)], N_EXPERTS)
    assert u["dead_experts"] == N_EXPERTS - 1
    assert u["max_over_mean"] == pytest.approx(N_EXPERTS)
    assert u["collapsed"] is True

    balanced = lab.expert_utilization([[i % N_EXPERTS] for i in range(64)], N_EXPERTS)
    assert balanced["dead_experts"] == 0
    assert balanced["max_over_mean"] == pytest.approx(1.0)
    assert balanced["collapsed"] is False


def test_utilization_fractions_sum_to_one():
    u = lab.expert_utilization([[i % N_EXPERTS, (i + 1) % N_EXPERTS]
                                for i in range(50)], N_EXPERTS)
    assert sum(u["fractions"]) == pytest.approx(1.0)


# ======================================================================================
# 4. Capacity, dropping, padding
# ======================================================================================

def test_capacity_formula():
    assert lab.expert_capacity(100, 4, 1, 1.25) == 31       # int(1.25*100*1/4)
    assert lab.expert_capacity(100, 4, 2, 1.0) == 50


def test_capacity_factor_below_one_is_rejected():
    """BOUNDARY: below 1.0 you drop tokens even under PERFECT balance. That is never
    what you want, so fail loudly rather than silently degrade."""
    with pytest.raises(ValueError):
        lab.expert_capacity(100, 4, 1, capacity_factor=0.9)


def test_perfectly_balanced_routing_drops_nothing_at_factor_one():
    """BOUNDARY: capacity_factor = 1.0 is exactly enough IF routing is perfect."""
    assign = [[i % 4] for i in range(100)]
    cap = lab.expert_capacity(100, 4, 1, 1.0)
    _, stats = lab.apply_capacity(assign, 4, cap)
    assert stats["dropped"] == 0


def test_imbalanced_routing_drops_the_overflow():
    assign = [[0]] * 40 + [[1]] * 30 + [[2]] * 20 + [[3]] * 10
    cap = lab.expert_capacity(100, 4, 1, 1.25)              # 31
    kept, stats = lab.apply_capacity(assign, 4, cap)
    assert stats["dropped"] == 9                           # expert 0 got 40, cap 31
    assert stats["drop_rate"] == pytest.approx(0.09)
    assert all(len(k) <= 1 for k in kept)


def test_higher_capacity_drops_less_and_wastes_more():
    """THE TRADE-OFF. This is what tuning capacity_factor actually buys and costs."""
    assign = [[0]] * 40 + [[1]] * 30 + [[2]] * 20 + [[3]] * 10
    prev_drop, prev_pad = None, None
    for cf in (1.0, 1.25, 2.0, 4.0):
        cap = lab.expert_capacity(100, 4, 1, cf)
        _, s = lab.apply_capacity(assign, 4, cap)
        if prev_drop is not None:
            assert s["dropped"] <= prev_drop
            assert s["padded_slots"] >= prev_pad
        prev_drop, prev_pad = s["dropped"], s["padded_slots"]


def test_zero_capacity_drops_everything():
    """BOUNDARY: cap of 0 means no token reaches any expert."""
    kept, s = lab.apply_capacity([[0], [1], [2]], 4, 0)
    assert s["dropped"] == 3
    assert all(k == [] for k in kept)


def test_apply_capacity_rejects_negative():
    with pytest.raises(ValueError):
        lab.apply_capacity([[0]], 4, -1)


# ======================================================================================
# 5. The forward pass
# ======================================================================================

def test_forward_preserves_shape():
    toks = _tokens(32)
    out, _ = lab.moe_forward(toks, _router(), _experts(), TOP_K)
    assert len(out) == len(toks)
    assert all(len(o) == D_MODEL for o in out)


def test_forward_is_deterministic():
    toks, r, ex = _tokens(32), _router(), _experts()
    a, da = lab.moe_forward(toks, r, ex, TOP_K)
    b, db = lab.moe_forward(toks, r, ex, TOP_K)
    assert a == b
    assert da["load_balance_loss"] == db["load_balance_loss"]


def test_fully_dropped_token_passes_through_on_the_residual():
    """THE SILENT FAILURE, made visible.

    A token whose every slot was dropped skips the FFN entirely. It is not an error and
    nothing logs it — the token just rides the residual stream unchanged. This is a real
    source of quality loss in production MoE models.
    """
    toks, r, ex = _tokens(32), _router(), _experts()
    out, diag = lab.moe_forward(toks, r, ex, TOP_K, capacity_factor=1.0)
    if diag["fully_dropped_tokens"] > 0:
        # At least one output must be byte-identical to its input.
        assert any(o == list(t) for o, t in zip(out, toks))


def test_shared_expert_prevents_full_dropping():
    """A shared expert runs for EVERY token, so no token can be fully dropped."""
    toks, r, ex = _tokens(32), _router(), _experts()
    shared = lab.make_expert(D_MODEL, D_FF, seed=999)
    _, diag = lab.moe_forward(toks, r, ex, TOP_K, capacity_factor=1.0,
                              shared_expert=shared)
    assert diag["fully_dropped_tokens"] == 0


def test_forward_rejects_empty_input():
    with pytest.raises(ValueError):
        lab.moe_forward([], _router(), _experts(), TOP_K)
    with pytest.raises(ValueError):
        lab.moe_forward(_tokens(4), _router(), [], TOP_K)


def test_total_loss_adds_the_auxiliary_terms():
    _, diag = lab.moe_forward(_tokens(32), _router(), _experts(), TOP_K)
    base = lab.total_training_loss(2.0, diag, aux_weight=0.0, z_weight=0.0)
    assert base == pytest.approx(2.0)
    with_aux = lab.total_training_loss(2.0, diag, aux_weight=0.01, z_weight=0.0)
    assert with_aux > base


def test_total_loss_rejects_negative_weights():
    _, diag = lab.moe_forward(_tokens(8), _router(), _experts(), TOP_K)
    with pytest.raises(ValueError):
        lab.total_training_loss(2.0, diag, aux_weight=-1.0)


# ======================================================================================
# 6. Parameter and communication accounting
# ======================================================================================

def test_total_exceeds_active_and_ratio_is_right():
    c = lab.moe_parameter_counts(D_MODEL, D_FF, n_experts=64, top_k=2)
    assert c["total"] > c["active"]
    assert c["sparsity_ratio"] == pytest.approx(c["total"] / c["active"])


def test_dense_equivalence_when_all_experts_active():
    """BOUNDARY: top_k == n_experts means sparsity 1.0 — you pay for everything."""
    c = lab.moe_parameter_counts(D_MODEL, D_FF, n_experts=8, top_k=8)
    assert c["total"] == c["active"]
    assert c["sparsity_ratio"] == pytest.approx(1.0)


def test_flops_track_top_k_not_expert_count():
    """THE MoE INVARIANT: adding experts costs MEMORY, not COMPUTE."""
    a = lab.moe_parameter_counts(D_MODEL, D_FF, n_experts=8, top_k=2)
    b = lab.moe_parameter_counts(D_MODEL, D_FF, n_experts=64, top_k=2)
    assert b["total"] > a["total"]                      # far more memory
    # Active differs only by the router, which grows with n_experts.
    assert b["active"] - a["active"] == D_MODEL * (64 - 8)


def test_parameter_counts_reject_bad_top_k():
    with pytest.raises(ValueError):
        lab.moe_parameter_counts(D_MODEL, D_FF, n_experts=4, top_k=5)


def test_comm_scales_with_layers_tokens_and_top_k():
    base = lab.expert_parallel_comm_bytes(8192, 8192, 60, 2)
    assert lab.expert_parallel_comm_bytes(8192, 8192, 120, 2) == 2 * base
    assert lab.expert_parallel_comm_bytes(16384, 8192, 60, 2) == 2 * base
    assert lab.expert_parallel_comm_bytes(8192, 8192, 60, 4) == 2 * base


def test_comm_time_has_a_transfer_and_a_latency_term():
    b = lab.expert_parallel_comm_bytes(8192, 8192, 60, 2)
    fast = lab.comm_seconds(b, 1e12, 120, latency_s_each=0.0)
    with_latency = lab.comm_seconds(b, 1e12, 120, latency_s_each=1e-3)
    assert with_latency > fast
    assert with_latency - fast == pytest.approx(120 * 1e-3)


def test_comm_rejects_zero_bandwidth():
    with pytest.raises(ValueError):
        lab.comm_seconds(1e9, 0, 10)


# ======================================================================================
# 7. Router collapse
# ======================================================================================

def test_router_collapses_without_an_auxiliary_loss():
    """THE MONEY TEST. Rich-get-richer is the DEFAULT behaviour of an MoE router, not
    an exotic failure. With no auxiliary loss it runs away from a tiny random asymmetry
    and you end up paying for E experts to run a dense model."""
    h = lab.simulate_collapse(N_EXPERTS, 1024, n_steps=600, aux_weight=0.0)
    assert h[-1] > 3.0


def test_auxiliary_loss_prevents_collapse():
    h = lab.simulate_collapse(N_EXPERTS, 1024, n_steps=600, aux_weight=0.01)
    assert h[-1] < 1.5


def test_collapse_is_monotone_in_the_aux_weight():
    """More balancing pressure must mean less imbalance."""
    finals = [lab.simulate_collapse(N_EXPERTS, 1024, 600, aux)[-1]
              for aux in (0.0, 0.002, 0.005, 0.01)]
    assert finals == sorted(finals, reverse=True)


def test_collapse_simulation_is_deterministic():
    a = lab.simulate_collapse(N_EXPERTS, 1024, 100, 0.0, seed=3)
    b = lab.simulate_collapse(N_EXPERTS, 1024, 100, 0.0, seed=3)
    assert a == b


def test_collapse_simulation_validates():
    with pytest.raises(ValueError):
        lab.simulate_collapse(1, 1024, 100, 0.01)
    with pytest.raises(ValueError):
        lab.simulate_collapse(8, 1024, 100, -0.01)