#!/usr/bin/env python3
"""
SMI-TED 推理正确性验证 — 不启用 FlagGems。
使用 MOSES 测试集 (1000 个分子), encode→decode 后计算 Morgan 指纹 Tanimoto 相似度。

用法:
  python scripts/inference/validate_accuracy_native.py
"""

import sys
import os
import time
import warnings
warnings.filterwarnings("ignore")

_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
_PROJECT_ROOT = os.path.dirname(os.path.dirname(_SCRIPT_DIR))
sys.path.insert(0, _PROJECT_ROOT)

import torch
import pandas as pd
import numpy as np
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.DataStructs import TanimotoSimilarity

from models.smi_ted.smi_ted_light.load import load_smi_ted, normalize_smiles


def main():
    model_dir = os.path.join(_PROJECT_ROOT, "models", "smi_ted", "smi_ted_light")
    data_path = os.path.join(_PROJECT_ROOT, "models", "smi_ted", "notebooks", "data", "moses_test.csv")

    if not os.path.isfile(data_path):
        print(f"[ERROR] 数据集不存在: {data_path}")
        sys.exit(1)

    device = "cuda" if torch.cuda.is_available() else "cpu"

    # ── 加载数据 ──
    df = pd.read_csv(data_path, nrows=1000)
    df["norm_smiles"] = df["SMILES"].apply(normalize_smiles)
    df = df.dropna()
    smiles_list = df["norm_smiles"].tolist()
    print(f"数据集: {len(smiles_list)} 个分子 (MOSES test)")

    # ── 加载模型 ──
    print("加载模型...")
    model = load_smi_ted(folder=model_dir, ckpt_filename="smi-ted-Light_40.pt")
    model.eval()
    if device == "cuda":
        model = model.cuda()

    # ── Encode + Decode ──
    print("推理中...")
    t0 = time.time()

    with torch.no_grad():
        embeddings = model.encode(smiles_list, batch_size=128, return_torch=True)
        reconstructed = model.decode(embeddings)

    elapsed = time.time() - t0
    print(f"  耗时: {elapsed:.1f}s ({len(smiles_list)/elapsed:.0f} mol/s)")

    # ── 计算 Tanimoto 相似度 ──
    similarities = []
    failed = 0
    for orig, rec in zip(smiles_list, reconstructed):
        mol_orig = Chem.MolFromSmiles(orig)
        mol_rec = Chem.MolFromSmiles(rec)
        if mol_orig is None or mol_rec is None:
            failed += 1
            continue
        fp_orig = AllChem.GetMorganFingerprintAsBitVect(mol_orig, 2)
        fp_rec = AllChem.GetMorganFingerprintAsBitVect(mol_rec, 2)
        similarities.append(TanimotoSimilarity(fp_orig, fp_rec))

    # ── 输出结果 ──
    mean_sim = np.mean(similarities)
    min_sim = np.min(similarities)
    perfect = sum(1 for s in similarities if s >= 0.999)

    print(f"\n{'=' * 50}")
    print(f"验证结果 (FlagGems: 未启用)")
    print(f"{'=' * 50}")
    print(f"  总样本数:    {len(smiles_list)}")
    print(f"  有效样本:    {len(similarities)}")
    print(f"  RDKit 失败:  {failed}")
    print(f"  均值 Tanimoto: {mean_sim:.4f}")
    print(f"  最小 Tanimoto: {min_sim:.4f}")
    print(f"  完美重建 (≥0.999): {perfect}/{len(similarities)} ({perfect/len(similarities)*100:.1f}%)")

    if mean_sim >= 0.99:
        print(f"\n  ✓ 推理正确性验证通过 (均值 ≥ 0.99)")
    else:
        print(f"\n  ✗ 推理正确性验证未通过!")
    print(f"{'=' * 50}")


if __name__ == "__main__":
    main()