"""Set Transformer 模块自检脚本

约定:
    - 每个 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.set_transformer import SetTransformer


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

def test_shape_preserved() -> None:
    st = SetTransformer(d_model=256, num_heads=8, n_layers=4)
    x = torch.randn(2, 4, 256)
    out = st(x)
    assert out.shape == (2, 4, 256), f"out shape {out.shape} != (2, 4, 256)"


def test_variable_set_size() -> None:
    """n_chem=3（无 MPNN）和 4（有 MPNN）都要支持。"""
    st = SetTransformer(d_model=64, num_heads=8, n_layers=2)
    for n in (3, 4):
        x = torch.randn(5, n, 64)
        assert st(x).shape == (5, n, 64), f"n={n} 形状不符"


def test_permutation_equivariance() -> None:
    """无位置编码，输出应随输入 token 置换而同步置换。"""
    torch.manual_seed(0)
    st = SetTransformer(d_model=64, num_heads=8, n_layers=3).eval()
    x = torch.randn(2, 4, 64)
    perm = torch.tensor([2, 0, 3, 1])
    out = st(x)
    out_perm = st(x[:, perm, :])
    assert torch.allclose(out[:, perm, :], out_perm, atol=1e-4), "未满足置换等变性"


def test_grad_flows() -> None:
    st = SetTransformer(d_model=32, num_heads=4, n_layers=2)
    x = torch.randn(3, 4, 32, requires_grad=True)
    st(x).sum().backward()
    assert x.grad is not None and torch.isfinite(x.grad).all(), "梯度异常"


def test_isab_shape() -> None:
    st = SetTransformer(d_model=64, num_heads=8, n_layers=2, block="isab", num_inducing=8)
    x = torch.randn(2, 4, 64)
    assert st(x).shape == (2, 4, 64), "ISAB 形状不符"


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