"""
LLM 编码器验证脚本（统一版）

对四种 LLM encoder 中的任意一种做四项检查：
    TEST 1: encoder 独立输出（shape / NaN / 梯度 / 缓存 / batch 一致性）
    TEST 2: 与 LNPModel 集成（任务头维度 / biodist 和为1 / 反向传播 / fusion token 数）
    TEST 3: 可复现性（同种子相同输出，不同种子不同输出）
    TEST 4: 过拟合 sanity（小数据上 loss 下降、R² > 0）

用法:
    python verify_llm_encoder.py \
        --model_path <llm_model_path> \
        --llm_type   <biot5|molt5|moleculestm|selfiested> \
        --lnp_repo_path <仓库根目录>

不同 encoder 的可训练主体属性名不同，本脚本会按 llm_type 自动选择：
    biot5       -> encoder._llm
    molt5       -> encoder._encoder
    moleculestm -> encoder._gnn
    selfiested  -> encoder._model
"""

import argparse
import sys
import random

import numpy as np
import torch
import torch.nn as nn


# llm_type -> (encoder 模块, 类名, 冻结主体属性名)
LLM_REGISTRY = {
    "biot5": ("lnp_ml.modeling.encoders.llm_encoder", "LLMEncoder", "_llm"),
    "molt5": ("lnp_ml.modeling.encoders.molt5_encoder", "MolT5Encoder", "_encoder"),
    "moleculestm": ("lnp_ml.modeling.encoders.moleculestm_encoder", "MoleculeSTMEncoder", "_gnn"),
    "selfiested": ("lnp_ml.modeling.encoders.selfiested_encoder", "SELFIESTEDEncoder", "_model"),
}


def set_seed(seed: int = 42) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    print(f"[Seed] All seeds fixed to {seed}")


def build_encoder(llm_type, model_path, d_model, device):
    """按 llm_type 构造 encoder，处理 model_path / checkpoint_path 参数名差异。"""
    import importlib
    module_name, class_name, _ = LLM_REGISTRY[llm_type]
    module = importlib.import_module(module_name)
    EncoderCls = getattr(module, class_name)

    if llm_type == "moleculestm":
        return EncoderCls(checkpoint_path=model_path, d_model=d_model, device=device)
    else:
        return EncoderCls(model_path=model_path, d_model=d_model, device=device)


def get_frozen_body(encoder, llm_type):
    """取出该 encoder 的冻结主体（用于检查参数是否冻结）。"""
    _, _, attr = LLM_REGISTRY[llm_type]
    return getattr(encoder, attr)


SAMPLE_SMILES = [
    "CCCCCCCCCCCCCCCC(=O)OCC(COC(=O)CCCCCCCCCCCCCCC)OC(=O)CCCCCCC",
    "CCN(CC)CCOC(=O)C1=CC=CC=C1",
    "CC(C)(C)OC(=O)N1CCC(CC1)C(=O)O",
    "O=C(O)c1ccccc1O",
]


def make_tabular(batch_size, device):
    return {
        "comp": torch.randn(batch_size, 5, device=device),
        "phys": torch.randn(batch_size, 12, device=device),
        "help": torch.randn(batch_size, 4, device=device),
        "exp": torch.randn(batch_size, 32, device=device),
    }


def test_1_standalone(llm_type, model_path, d_model, device):
    print("\n" + "=" * 60)
    print("TEST 1: encoder standalone")
    print("=" * 60)

    print("\n[1a] Testing output shape...")
    encoder = build_encoder(llm_type, model_path, d_model, device)
    out = encoder(SAMPLE_SMILES)
    assert out.shape == (len(SAMPLE_SMILES), d_model), f"bad shape {out.shape}"
    print(f"  Output shape: {out.shape}  OK")

    print("[1b] Checking for NaN/Inf...")
    assert not torch.isnan(out).any() and not torch.isinf(out).any()
    print(f"  No NaN/Inf  OK")
    print(f"  Output stats: mean={out.mean():.4f}, std={out.std():.4f}, "
          f"min={out.min():.4f}, max={out.max():.4f}")

    print("[1c] Checking gradients...")
    body = get_frozen_body(encoder, llm_type)
    n_trainable_body = sum(p.numel() for p in body.parameters() if p.requires_grad)
    assert n_trainable_body == 0, f"frozen body has {n_trainable_body} trainable params"
    print(f"  Frozen body ({0} trainable params)  OK")
    n_proj = sum(p.numel() for p in encoder.projection.parameters() if p.requires_grad)
    assert n_proj > 0
    print(f"  Projection trainable ({n_proj:,} params)  OK")
    loss = encoder(SAMPLE_SMILES).sum()
    loss.backward()
    grad_ok = any(p.grad is not None for p in encoder.projection.parameters())
    assert grad_ok
    print(f"  Gradient flows to projection  OK")

    print("[1e] Single vs batch consistency...")
    encoder.eval()
    with torch.no_grad():
        batch_out = encoder(SAMPLE_SMILES)
        single_out = torch.cat([encoder([s]) for s in SAMPLE_SMILES], dim=0)
    assert torch.allclose(batch_out, single_out, atol=1e-4)
    print(f"  Batch == single  OK")
    print("\nTEST 1 PASSED")


def test_2_integration(llm_type, model_path, d_model, device):
    print("\n" + "=" * 60)
    print("TEST 2: Integration with LNPModel")
    print("=" * 60)

    from lnp_ml.modeling.models import LNPModelWithoutMPNN
    model = LNPModelWithoutMPNN(
        d_model=d_model, use_llm=True, llm_type=llm_type,
        llm_model_path=model_path, llm_device=device,
    ).to(device)

    print("\n[2a] Testing full forward pass shapes...")
    tab = make_tabular(len(SAMPLE_SMILES), device)
    outputs = model(SAMPLE_SMILES, tab)
    expected = {"size": 1, "pdi": 4, "ee": 3, "delivery": 1, "biodist": 7, "toxic": 2}
    for k, dim in expected.items():
        assert outputs[k].shape == (len(SAMPLE_SMILES), dim), f"{k}: {outputs[k].shape}"
        print(f"  outputs['{k}']: {tuple(outputs[k].shape)}  OK")

    print("[2b] Checking biodist sums to 1...")
    sums = outputs["biodist"].sum(dim=-1)
    assert torch.allclose(sums, torch.ones_like(sums), atol=1e-4)
    print(f"  biodist sums: {[round(s.item(),4) for s in sums]}  OK")

    print("[2c] Testing backward pass...")
    loss = sum(o.sum() for o in outputs.values())
    loss.backward()
    proj_grad = any(p.grad is not None for p in model.llm_encoder.projection.parameters())
    assert proj_grad
    print(f"  llm_encoder.projection gets gradients  OK")
    body = get_frozen_body(model.llm_encoder, llm_type)
    body_grad = any(p.grad is not None and p.grad.abs().sum() > 0 for p in body.parameters())
    assert not body_grad
    print(f"  LLM body has no gradients (frozen)  OK")

    print("[2d] Checking fusion token count...")
    # 无 MPNN 时主干 7 个 token + 1 个 V5 = 8
    assert model.n_fusion_tokens == 8, f"n_fusion_tokens={model.n_fusion_tokens}"
    print(f"  fusion.n_tokens={model.n_fusion_tokens}  OK")
    print("\nTEST 2 PASSED")


def test_3_reproducibility(llm_type, model_path, d_model, device):
    print("\n" + "=" * 60)
    print("TEST 3: Reproducibility")
    print("=" * 60)

    print("\n[3a] Running twice with same seed=42...")
    set_seed(42)
    enc1 = build_encoder(llm_type, model_path, d_model, device)
    enc1.eval()
    with torch.no_grad():
        o1 = enc1(SAMPLE_SMILES)
    set_seed(42)
    enc2 = build_encoder(llm_type, model_path, d_model, device)
    enc2.eval()
    with torch.no_grad():
        o2 = enc2(SAMPLE_SMILES)
    assert torch.allclose(o1, o2, atol=1e-5)
    print(f"  Same seed -> identical outputs  OK")

    print("[3b] Running with different seeds...")
    set_seed(123)
    enc3 = build_encoder(llm_type, model_path, d_model, device)
    enc3.eval()
    with torch.no_grad():
        o3 = enc3(SAMPLE_SMILES)
    assert not torch.allclose(o1, o3, atol=1e-5)
    print(f"  Different seeds -> different outputs  OK")
    print("\nTEST 3 PASSED")


def test_4_overfit(llm_type, model_path, d_model, device):
    print("\n" + "=" * 60)
    print("TEST 4: Overfit sanity (gradient + metric check)")
    print("=" * 60)
    set_seed(42)

    from lnp_ml.modeling.models import LNPModelWithoutMPNN
    model = LNPModelWithoutMPNN(
        d_model=d_model, use_llm=True, llm_type=llm_type,
        llm_model_path=model_path, llm_device=device,
    ).to(device)

    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"\n  Trainable params: {n_trainable:,}")

    smiles = SAMPLE_SMILES * 2  # 8 samples
    tab = make_tabular(len(smiles), device)
    target = torch.randn(len(smiles), 1, device=device)

    opt = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=1e-3)
    lossfn = nn.MSELoss()

    print(f"  Training 50 steps on {len(smiles)} samples...")
    first_loss = None
    model.train()
    for step in range(1, 51):
        opt.zero_grad()
        pred = model.forward_delivery(smiles, tab)
        loss = lossfn(pred, target)
        loss.backward()
        opt.step()
        if first_loss is None:
            first_loss = loss.item()
        if step % 10 == 0:
            print(f"    step {step:3d}: loss={loss.item():.4f}")

    last_loss = loss.item()
    assert last_loss < first_loss
    print(f"  Loss decreased: {first_loss:.4f} -> {last_loss:.4f}  OK")

    model.eval()
    with torch.no_grad():
        pred = model.forward_delivery(smiles, tab)
    ss_res = ((pred - target) ** 2).sum().item()
    ss_tot = ((target - target.mean()) ** 2).sum().item()
    r2 = 1 - ss_res / (ss_tot + 1e-9)
    rmse = (ss_res / len(smiles)) ** 0.5
    print(f"\n  After 50 steps on {len(smiles)} samples (overfit test):")
    print(f"    R2   = {r2:.4f}  (should be > 0 after overfitting)")
    print(f"    RMSE = {rmse:.4f}")
    assert r2 > 0
    print(f"  R2 > 0  OK")
    print("\nTEST 4 PASSED")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", required=True, help="LLM 权重路径或 HuggingFace ID")
    parser.add_argument("--llm_type", required=True,
                        choices=["biot5", "molt5", "moleculestm", "selfiested"])
    parser.add_argument("--lnp_repo_path", default=".", help="仓库根目录")
    parser.add_argument("--d_model", type=int, default=256)
    args = parser.parse_args()

    if args.lnp_repo_path not in sys.path:
        sys.path.insert(0, args.lnp_repo_path)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}")
    set_seed(42)

    test_1_standalone(args.llm_type, args.model_path, args.d_model, device)
    test_2_integration(args.llm_type, args.model_path, args.d_model, device)
    test_3_reproducibility(args.llm_type, args.model_path, args.d_model, device)
    test_4_overfit(args.llm_type, args.model_path, args.d_model, device)

    print("\n" + "=" * 60)
    print("ALL TESTS PASSED")
    print("=" * 60)


if __name__ == "__main__":
    main()