"""MoE 模块自检脚本

约定:
    - 每个 test_xxx 函数返回 None；失败时直接抛 AssertionError 或 RuntimeError。
    - main() 收集所有 test_xxx 函数依次执行，统计 pass/fail 并以非零状态退出。
"""

from __future__ import annotations

import sys
import traceback
from typing import Callable, List, Tuple

import torch

from lnp_ml.modeling.layers.moe import MoEAttentionPool, MoEBlock, MoERouter


# ──────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────

def _expect_value_error(fn: Callable[[], object], match: str = "") -> None:
    """断言调用 fn 会抛 ValueError，且消息包含 match（若给定）。"""
    try:
        fn()
    except ValueError as e:
        if match and match not in str(e):
            raise AssertionError(
                f"ValueError 抛出了，但消息中未包含 {match!r}: {e}"
            )
        return
    raise AssertionError(f"期望抛出 ValueError，但 {fn} 正常返回")


# ──────────────────────────────────────────────────────────────────────
# Tests
# ──────────────────────────────────────────────────────────────────────

def test_moe_block_shape() -> None:
    B, T_chem, T_tab, d, K = 2, 4, 4, 256, 4
    block = MoEBlock(d_model=d, n_chem_tokens=T_chem, n_experts=K, top_k=2)
    chem = torch.randn(B, T_chem, d)
    tab = torch.randn(B, T_tab, d)

    F_moe, extras = block(chem, tab)

    assert F_moe.shape == (B, d), f"F_moe shape {F_moe.shape} != ({B}, {d})"
    assert extras["lb_loss"].dim() == 0
    assert extras["lb_loss"].requires_grad
    assert extras["gates"].shape == (B, K)
    assert extras["probs"].shape == (B, K)


def test_moe_block_nchem3() -> None:
    B, T_chem, T_tab, d, K = 2, 3, 4, 256, 4
    block = MoEBlock(d_model=d, n_chem_tokens=T_chem, n_experts=K, top_k=2)
    chem = torch.randn(B, T_chem, d)
    tab = torch.randn(B, T_tab, d)

    F_moe, extras = block(chem, tab)

    assert F_moe.shape == (B, d), f"F_moe shape {F_moe.shape} != ({B}, {d})"
    assert block.experts[0].net[0].in_features == T_chem * d, "expert 输入维未随 n_chem 自适应"
    assert extras["gates"].shape == (B, K)


def test_moe_gates_topk_normalization() -> None:
    block = MoEBlock(d_model=64, n_chem_tokens=4, n_experts=4, top_k=2)
    chem = torch.randn(8, 4, 64)
    tab = torch.randn(8, 4, 64)
    _, extras = block(chem, tab)
    gates = extras["gates"]

    assert torch.allclose(gates.sum(dim=-1), torch.ones(8), atol=1e-5), \
        f"gates 每行未归一为 1: {gates.sum(dim=-1)}"
    nonzero_per_row = (gates > 0).sum(dim=-1).max().item()
    assert nonzero_per_row <= 2, f"top_k=2 但某行有 {nonzero_per_row} 个非零"


def test_moe_load_balancing_loss_uniform() -> None:
    """随机初始化 + 大 batch，lb_loss 应接近 Switch 风格的 uniform 极限。

    注意:
        Switch 风格 lb_loss = K * Σ f_i * p_i。
        - Σ f_i = top_k （top-k mask 的列和约束）
        - Σ p_i = 1
        - 完全 uniform 极限 ⇒ lb_loss = top_k
        随机初始化下 f 与 p 因 top-k 选择天然正相关，
        实际值会略高于 top_k，但远低于 K（K 是塌缩极限）。
    """
    torch.manual_seed(42)
    block = MoEBlock(d_model=64, n_chem_tokens=4, n_experts=4, top_k=2)
    chem = torch.randn(64, 4, 64)
    tab = torch.randn(64, 4, 64)
    _, extras = block(chem, tab)
    lb = extras["lb_loss"].item()

    expected = block.top_k
    upper = (block.top_k + block.n_experts) / 2  # uniform 与塌缩的中点
    assert expected * 0.8 <= lb <= upper, (
        f"lb_loss={lb:.3f} 偏离理想 uniform 值 {expected} 过远"
        f"（健康区间 ~[{expected*0.8:.2f}, {upper:.2f}]）"
    )

def test_moe_invalid_top_k() -> None:
    _expect_value_error(
        lambda: MoEBlock(d_model=64, n_chem_tokens=4, n_experts=4, top_k=5)
    )


def test_moe_chem_token_mismatch() -> None:
    block = MoEBlock(d_model=64, n_chem_tokens=4, n_experts=4, top_k=2)
    chem = torch.randn(2, 3, 64)
    tab = torch.randn(2, 4, 64)
    _expect_value_error(lambda: block(chem, tab), match="chem token")


def test_moe_router_eval_no_jitter() -> None:
    """评估态下 jitter 不生效：相同输入应得相同 gates。"""
    router = MoERouter(d_model=32, n_experts=4, top_k=2, jitter_noise=0.5)
    router.eval()
    q = torch.randn(4, 32)
    g1, _, _ = router(q)
    g2, _, _ = router(q)
    assert torch.allclose(g1, g2), "eval 模式下相同输入产生了不同 gates"


def test_moe_attention_pool_shape() -> None:
    pool = MoEAttentionPool(d_model=128)
    x = torch.randn(3, 5, 128)
    out = pool(x)
    assert out.shape == (3, 128), f"pool out shape {out.shape} != (3, 128)"


# ──────────────────────────────────────────────────────────────────────
# Runner
# ──────────────────────────────────────────────────────────────────────

def _collect_tests() -> List[Tuple[str, Callable[[], None]]]:
    return [
        (name, fn) for name, fn in globals().items()
        if name.startswith("test_") and callable(fn)
    ]


def main() -> int:
    tests = _collect_tests()
    pass_count = 0
    fail: List[Tuple[str, str]] = []

    for name, fn in tests:
        try:
            fn()
        except Exception:
            fail.append((name, traceback.format_exc()))
            print(f"[FAIL] {name}")
        else:
            pass_count += 1
            print(f"[ OK ] {name}")

    print(f"\n{'='*60}")
    print(f"Passed: {pass_count}/{len(tests)}")
    if fail:
        print(f"\n{'='*60}\nFailures:")
        for name, tb in fail:
            print(f"\n--- {name} ---\n{tb}")
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())