"""ResidualConcatFusion 模块自检脚本
"""

from __future__ import annotations

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

import torch

from lnp_ml.modeling.layers.fusion import ResidualConcatFusion


def test_full_tokens_shape() -> None:
    fusion = ResidualConcatFusion(d_model=256)
    chem = torch.randn(2, 4, 256)
    tab = torch.randn(2, 4, 256)
    f_moe = torch.randn(2, 256)
    f_llm = torch.randn(2, 256)
    out = fusion(chem, tab, f_moe, f_llm)
    assert out.shape == (2, 256), f"out shape {out.shape} != (2, 256)"
    assert fusion.fusion_dim == 256


def test_without_extras() -> None:
    """关闭 MoE 和 LLM 时只拼接 chem + tab。"""
    fusion = ResidualConcatFusion(d_model=64)
    chem = torch.randn(3, 4, 64)
    tab = torch.randn(3, 4, 64)
    assert fusion(chem, tab).shape == (3, 64)


def test_nchem3_with_moe() -> None:
    fusion = ResidualConcatFusion(d_model=64)
    chem = torch.randn(2, 3, 64)
    tab = torch.randn(2, 4, 64)
    f_moe = torch.randn(2, 64)
    assert fusion(chem, tab, f_moe=f_moe).shape == (2, 64)


def test_attn_weights_shape() -> None:
    """注意力只池化真实 token：序列长度 = 4 + 4 = 8；f_moe/f_llm 走门控残差，不参与 softmax。"""
    fusion = ResidualConcatFusion(d_model=64)
    chem = torch.randn(2, 4, 64)
    tab = torch.randn(2, 4, 64)
    f_moe = torch.randn(2, 64)
    f_llm = torch.randn(2, 64)
    out, weights = fusion(chem, tab, f_moe, f_llm, return_attn_weights=True)
    assert out.shape == (2, 64)
    assert weights.shape == (2, 8), f"weights shape {weights.shape} != (2, 8)"


def test_concat_strategy_rejected() -> None:
    try:
        ResidualConcatFusion(d_model=64, strategy="concat")
    except ValueError:
        return
    raise AssertionError("concat 策略应被拒绝")


def test_grad_flows() -> None:
    fusion = ResidualConcatFusion(d_model=32)
    chem = torch.randn(2, 4, 32, requires_grad=True)
    tab = torch.randn(2, 4, 32, requires_grad=True)
    fusion(chem, tab).sum().backward()
    assert chem.grad is not None and torch.isfinite(chem.grad).all()


# ──────────────────────────────────────────────────────────────────────
# 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())