"""LNPModel 整图自检脚本

base/+moe 路径只需 rdkit + torch；+llm 路径需 MolT5。
"""

from __future__ import annotations

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

import torch

from lnp_ml.featurization.smiles import RDKitFeaturizer
from lnp_ml.modeling.models import LNPModelWithoutMPNN

MOLT5_PATH = os.environ.get("MOLT5_PATH", "models/molt5-base")

# 期望的任务输出维度
EXPECTED = {"size": 1, "pdi": 2, "ee": 3, "delivery": 1, "biodist": 7, "toxic": 2}


def _input_dims() -> Dict[str, int]:
    """根据实际 RDKit 输出确定无 MPNN 时的输入维度。"""
    feat = RDKitFeaturizer().transform(["CCO"])
    return {
        "morgan": feat["morgan"].shape[1],
        "maccs": feat["maccs"].shape[1],
        "desc": feat["desc"].shape[1],
        "comp": 5, "phys": 12, "help": 4, "exp": 32,
    }


def _batch(B: int = 2) -> Tuple[List[str], Dict[str, torch.Tensor]]:
    smiles = ["CCO", "CCN"][:B]
    tabular = {
        "comp": torch.randn(B, 5),
        "phys": torch.randn(B, 12),
        "help": torch.randn(B, 4),
        "exp": torch.randn(B, 32),
    }
    return smiles, tabular


def _check_outputs(outputs: Dict[str, torch.Tensor], B: int) -> None:
    for task, dim in EXPECTED.items():
        assert outputs[task].shape == (B, dim), f"{task}: {outputs[task].shape} != ({B}, {dim})"


def test_forward_base() -> None:
    model = LNPModelWithoutMPNN(d_model=64, input_dims=_input_dims()).eval()
    smiles, tabular = _batch()
    _check_outputs(model(smiles, tabular), B=2)


def test_forward_moe() -> None:
    model = LNPModelWithoutMPNN(d_model=64, input_dims=_input_dims(), use_moe=True).eval()
    smiles, tabular = _batch()
    _check_outputs(model(smiles, tabular), B=2)
    assert model.get_last_moe_extras() is not None, "use_moe=True 应产生 MoE extras"


def test_forward_from_projected() -> None:
    model = LNPModelWithoutMPNN(d_model=64, input_dims=_input_dims()).eval()
    smiles, tabular = _batch()
    stacked = model._encode_and_project(smiles, tabular)
    out = model.forward_from_projected(stacked, task="biodist")
    assert out.shape == (2, 7), f"biodist {out.shape} != (2, 7)"


def test_split_idx_no_mpnn() -> None:
    model = LNPModelWithoutMPNN(d_model=64, input_dims=_input_dims())
    assert model.split_idx == 3, "无 MPNN 时化学 token 应为 3"


def test_forward_llm() -> None:
    """需要 MolT5；缺失则跳过。"""
    try:
        model = LNPModelWithoutMPNN(
            d_model=64, input_dims=_input_dims(),
            use_moe=True, use_llm=True, llm_model_path=MOLT5_PATH,
        ).eval()
    except Exception as e:
        print(f"[SKIP] test_forward_llm（MolT5 不可用）：{e}")
        return
    smiles, tabular = _batch()
    _check_outputs(model(smiles, tabular), B=2)
    # 冻结的 MolT5 encoder 不应进入 backbone state dict
    bb = model.get_backbone_state_dict()
    assert not any(k.startswith("llm_prompt.encoder.") for k in bb), "backbone 不应含冻结 encoder 权重"


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