lnp_ml/lnp_ml/modeling/nested_cv_optuna.py

"""
嵌套交叉验证 + Optuna 超参调优

外层 5-fold StratifiedKFold（20% test / 80% train）
内层 3-fold StratifiedKFold（在 80% 上做 Optuna 超参搜索）

使用方法:
    python -m lnp_ml.modeling.nested_cv_optuna
    
或通过 Makefile:
    make nested_cv_tune DEVICE=cuda
"""

import json
import math
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
import torch
from torch.utils.data import DataLoader, Subset
from sklearn.model_selection import StratifiedKFold
from loguru import logger
import typer

try:
    import optuna
    from optuna.samplers import TPESampler
except ImportError:
    optuna = None
    TPESampler = None

from lnp_ml.config import MODELS_DIR, INTERIM_DATA_DIR
from lnp_ml.dataset import (
    LNPDataset,
    collate_fn,
    process_dataframe,
    TARGET_CLASSIFICATION_PDI,
    TARGET_CLASSIFICATION_EE,
    TARGET_TOXIC,
)
from tqdm import tqdm

from lnp_ml.modeling.models import LNPModel, LNPModelWithoutMPNN
from lnp_ml.modeling.layers.llm_prompt import DEFAULT_MOLT5_PATH
from lnp_ml.utils.seed import set_global_seed
from lnp_ml.modeling.encoders.rdkit_encoder import CachedRDKitEncoder
from lnp_ml.modeling.trainer_balanced import (
    ClassWeights,
    LossWeightsBalanced,
    compute_class_weights_from_loader,
    train_with_early_stopping,
    train_fixed_epochs,
    validate_balanced,
)

# MPNN ensemble 默认路径
DEFAULT_MPNN_ENSEMBLE_DIR = MODELS_DIR / "mpnn" / "all_amine_split_for_LiON"

app = typer.Typer()


# ============ CompositeStrata 复合分层标签 ============

def build_composite_strata(
    df: pd.DataFrame,
    min_stratum_count: int = 5,
) -> Tuple[np.ndarray, Dict]:
    """
    构建复合分层标签（toxic × PDI × EE）。
    
    Args:
        df: 处理后的 DataFrame
        min_stratum_count: 每个 stratum 最少样本数，低于此值合并为 RARE
    
    Returns:
        (strata_array, strata_info) 
        - strata_array: 每个样本的 stratum 编码（整数）
        - strata_info: 统计信息
    """
    n = len(df)
    strata_labels = []
    
    for i in range(n):
        # Toxic stratum
        if TARGET_TOXIC in df.columns:
            toxic_val = df[TARGET_TOXIC].iloc[i]
            if pd.notna(toxic_val) and toxic_val >= 0:
                toxic_str = str(int(toxic_val))
            else:
                toxic_str = "NA"
        else:
            toxic_str = "NA"
        
        # PDI stratum
        if all(col in df.columns for col in TARGET_CLASSIFICATION_PDI):
            pdi_vals = df[TARGET_CLASSIFICATION_PDI].iloc[i].values
            if pdi_vals.sum() > 0:
                pdi_str = str(int(np.argmax(pdi_vals)))
            else:
                pdi_str = "NA"
        else:
            pdi_str = "NA"
        
        # EE stratum
        if all(col in df.columns for col in TARGET_CLASSIFICATION_EE):
            ee_vals = df[TARGET_CLASSIFICATION_EE].iloc[i].values
            if ee_vals.sum() > 0:
                ee_str = str(int(np.argmax(ee_vals)))
            else:
                ee_str = "NA"
        else:
            ee_str = "NA"
        
        strata_labels.append(f"T{toxic_str}|P{pdi_str}|E{ee_str}")
    
    # 统计各 stratum 的样本数
    unique_strata, counts = np.unique(strata_labels, return_counts=True)
    strata_counts = dict(zip(unique_strata, counts))
    
    # 将稀疏 strata 合并为 RARE
    rare_strata = [s for s, c in strata_counts.items() if c < min_stratum_count]
    
    final_labels = []
    for label in strata_labels:
        if label in rare_strata:
            final_labels.append("RARE")
        else:
            final_labels.append(label)
    
    # 编码为整数
    unique_final, encoded = np.unique(final_labels, return_inverse=True)
    
    strata_info = {
        "original_strata_counts": strata_counts,
        "rare_strata": rare_strata,
        "final_strata": list(unique_final),
        "final_strata_counts": dict(zip(*np.unique(final_labels, return_counts=True))),
        "n_rare_merged": sum(strata_counts[s] for s in rare_strata) if rare_strata else 0,
    }
    
    logger.info(f"Built composite strata: {len(unique_final)} unique strata")
    logger.info(f"  Rare strata merged: {len(rare_strata)} types, {strata_info['n_rare_merged']} samples")
    
    return encoded.astype(np.int64), strata_info


# ============ RDKit 缓存预热 ============

def warmup_rdkit_cache(smiles_list: List[str], batch_size: int = 256) -> Dict:
    """预热 RDKit 特征缓存，返回可跨模型共享的缓存字典。"""
    encoder = CachedRDKitEncoder()
    unique_smiles = list(set(smiles_list))
    logger.info(f"Warming up RDKit cache for {len(unique_smiles)} unique SMILES...")
    for i in tqdm(range(0, len(unique_smiles), batch_size), desc="Cache warmup"):
        batch = unique_smiles[i:i + batch_size]
        encoder(batch)
    logger.success(f"Cache warmup complete. Cached {len(encoder._cache)} SMILES.")
    return encoder._cache


# ============ 模型创建 ============

def find_mpnn_ensemble_paths(base_dir: Path = DEFAULT_MPNN_ENSEMBLE_DIR) -> List[str]:
    """自动查找 MPNN ensemble 的 model.pt 文件。"""
    model_paths = sorted(base_dir.glob("cv_*/fold_*/model_*/model.pt"))
    if not model_paths:
        raise FileNotFoundError(f"No model.pt files found in {base_dir}")
    return [str(p) for p in model_paths]


def _load_external_retrieval_pool(csv_path="data/external/all_data_LiON.csv"):
    """读外部 LiON 数据，筛 mRNA+Mouse，按唯一分子聚合 delivery，返回 (smiles_list, labels[N,1])。"""
    import pandas as pd
    import numpy as np
    df = pd.read_csv(csv_path, low_memory=False)
    sub = df[(df["Cargo_type"] == "mRNA") & (df["Model_type"] == "Mouse")].copy()
    sub["quantified_delivery"] = pd.to_numeric(sub["quantified_delivery"], errors="coerce")
    sub = sub.dropna(subset=["quantified_delivery", "smiles"])
    # 同一分子多器官记录取平均，聚合成唯一分子
    agg = sub.groupby("smiles")["quantified_delivery"].mean()
    smiles_list = agg.index.tolist()
    labels = agg.values.reshape(-1, 1).astype("float32")
    return smiles_list, labels


def create_model(
    d_model: int = 256,
    num_heads: int = 8,
    n_attn_layers: int = 4,
    fusion_strategy: str = "attention",
    head_hidden_dim: int = 128,
    dropout: float = 0.1,
    use_mpnn: bool = False,
    mpnn_device: str = "cpu",
    chemeleon_cache: Optional[str] = None,
    unimol_cache: Optional[str] = None,
    moleculestm_cache: Optional[str] = None,
    mole_cache: Optional[str] = None,
    set_transformer_block: str = "sab",
    # MoE
    use_moe: bool = False,
    moe_n_experts: int = 4,
    moe_top_k: int = 2,
    moe_expert_hidden_mult: int = 2,
    moe_jitter_noise: float = 0.0,
    # LLM（打包成 dict 以减少多进程参数传递）
    llm_kwargs: Optional[Dict] = None,
    # 检索增强
    use_retrieval: bool = False,
    retr_feature_dim: int = 0,
) -> Union[LNPModel, LNPModelWithoutMPNN]:
    """创建模型。llm_kwargs 含 use_llm / llm_model_path / llm_freeze / llm_use_lora / llm_lora_*。"""
    extra_kwargs = dict(
        set_transformer_block=set_transformer_block,
        use_moe=use_moe,
        moe_n_experts=moe_n_experts,
        moe_top_k=moe_top_k,
        moe_expert_hidden_mult=moe_expert_hidden_mult,
        moe_jitter_noise=moe_jitter_noise,
        use_retrieval=use_retrieval,
        retr_feature_dim=retr_feature_dim,
        chemeleon_cache_path=chemeleon_cache,
        unimol_cache_path=unimol_cache,
        moleculestm_cache_path=moleculestm_cache,
        mole_cache_path=mole_cache,
        **(llm_kwargs or {}),
    )

    if use_mpnn:
        ensemble_paths = find_mpnn_ensemble_paths()
        return LNPModel(
            d_model=d_model,
            num_heads=num_heads,
            n_attn_layers=n_attn_layers,
            fusion_strategy=fusion_strategy,
            head_hidden_dim=head_hidden_dim,
            dropout=dropout,
            mpnn_ensemble_paths=ensemble_paths,
            mpnn_device=mpnn_device,
            **extra_kwargs,
        )
    else:
        return LNPModelWithoutMPNN(
            d_model=d_model,
            num_heads=num_heads,
            n_attn_layers=n_attn_layers,
            fusion_strategy=fusion_strategy,
            head_hidden_dim=head_hidden_dim,
            dropout=dropout,
            **extra_kwargs,
        )


def _build_rag_pool(full_dataset, idx):
    """从 dataset + 索引构造 RAG 池：返回 (smiles, delivery[N], extra_labels)。"""
    import numpy as np
    idx = np.asarray(idx)
    smiles = [full_dataset.smiles[i] for i in idx]
    delivery = full_dataset.delivery[idx].reshape(-1)
    extra = {}
    if full_dataset.size is not None:
        extra["size"] = (full_dataset.size[idx], ~np.isnan(full_dataset.size[idx]))
    if full_dataset.pdi is not None:
        extra["pdi"] = (full_dataset.pdi[idx], full_dataset.pdi_valid[idx])
    if full_dataset.ee is not None:
        extra["ee"] = (full_dataset.ee[idx], full_dataset.ee_valid[idx])
    if full_dataset.toxic is not None:
        extra["toxic"] = (full_dataset.toxic[idx], full_dataset.toxic[idx] >= 0)
    if full_dataset.biodist is not None:
        extra["biodist"] = (full_dataset.biodist[idx], full_dataset.biodist_valid[idx])
    return smiles, delivery, extra


# ============ 评估指标 ============

def evaluate_on_test(
    model: torch.nn.Module,
    test_loader: DataLoader,
    device: torch.device,
) -> Dict:
    """在测试集上评估模型"""
    from scipy.special import rel_entr
    from sklearn.metrics import (
        mean_squared_error,
        mean_absolute_error,
        r2_score,
        accuracy_score,
        precision_score,
        recall_score,
        f1_score,
    )
    
    model.eval()
    
    preds = {
        "size": [], "delivery": [], "pdi": [], "ee": [], "toxic": [], "biodist": []
    }
    targets = {
        "size": [], "delivery": [], "pdi": [], "ee": [], "toxic": [], "biodist": []
    }
    
    with torch.no_grad():
        for batch in test_loader:
            smiles = batch["smiles"]
            tabular = {k: v.to(device) for k, v in batch["tabular"].items()}
            tgts = batch["targets"]
            masks = batch["mask"]
            
            outputs = model(smiles, tabular)
            
            # 收集预测和真实值
            for task in ["size", "delivery"]:
                if task in masks and masks[task].any():
                    m = masks[task]
                    key = task if task == "size" else "delivery"
                    preds[task].extend(outputs[key].squeeze(-1)[m].cpu().numpy().tolist())
                    targets[task].extend(tgts[key][m].cpu().numpy().tolist())
            
            for task in ["pdi", "ee", "toxic"]:
                if task in masks and masks[task].any():
                    m = masks[task]
                    preds[task].extend(outputs[task][m].argmax(dim=-1).cpu().numpy().tolist())
                    targets[task].extend(tgts[task][m].cpu().numpy().tolist())
            
            if "biodist" in masks and masks["biodist"].any():
                m = masks["biodist"]
                preds["biodist"].extend(outputs["biodist"][m].cpu().numpy().tolist())
                targets["biodist"].extend(tgts["biodist"][m].cpu().numpy().tolist())
    
    # 计算指标
    results = {}
    
    # 回归任务
    for task in ["size", "delivery"]:
        if preds[task]:
            p = np.array(preds[task])
            t = np.array(targets[task])
            results[task] = {
                "n_samples": len(p),
                "mse": float(mean_squared_error(t, p)),
                "rmse": float(np.sqrt(mean_squared_error(t, p))),
                "mae": float(mean_absolute_error(t, p)),
                "r2": float(r2_score(t, p)),
            }
    
    # 分类任务
    for task in ["pdi", "ee", "toxic"]:
        if preds[task]:
            p = np.array(preds[task])
            t = np.array(targets[task])
            results[task] = {
                "n_samples": len(p),
                "accuracy": float(accuracy_score(t, p)),
                "precision": float(precision_score(t, p, average="macro", zero_division=0)),
                "recall": float(recall_score(t, p, average="macro", zero_division=0)),
                "f1": float(f1_score(t, p, average="macro", zero_division=0)),
            }
    
    # 分布任务
    if preds["biodist"]:
        p = np.array(preds["biodist"])
        t = np.array(targets["biodist"])
        
        def kl_divergence(p_arr, q_arr, eps=1e-10):
            p_arr = np.clip(p_arr, eps, 1.0)
            q_arr = np.clip(q_arr, eps, 1.0)
            return float(np.sum(rel_entr(p_arr, q_arr), axis=-1).mean())
        
        def js_divergence(p_arr, q_arr, eps=1e-10):
            p_arr = np.clip(p_arr, eps, 1.0)
            q_arr = np.clip(q_arr, eps, 1.0)
            m = 0.5 * (p_arr + q_arr)
            return float(0.5 * (np.sum(rel_entr(p_arr, m), axis=-1) + np.sum(rel_entr(q_arr, m), axis=-1)).mean())
        
        results["biodist"] = {
            "n_samples": len(p),
            "kl_divergence": kl_divergence(t, p),
            "js_divergence": js_divergence(t, p),
        }
    
    return results


# ============ 预训练权重加载 ============

def load_pretrain_weights_to_model(
    model: Union[LNPModel, LNPModelWithoutMPNN],
    pretrain_state_dict: Dict,
    d_model: int,
    pretrain_config: Dict,
    load_delivery_head: bool = True,
) -> bool:
    """
    加载预训练权重到模型。
    
    Returns:
        是否成功加载
    """
    if pretrain_config.get("d_model") != d_model:
        logger.warning(
            f"d_model mismatch: pretrain={pretrain_config.get('d_model')}, "
            f"current={d_model}. Skipping pretrain loading."
        )
        return False
    
    model.load_pretrain_weights(
        pretrain_state_dict=pretrain_state_dict,
        load_delivery_head=load_delivery_head,
        strict=False,
    )
    return True


# ============ 内层 Optuna 调参 ============
def run_inner_optuna(
    full_dataset: LNPDataset,
    inner_train_indices: np.ndarray,
    strata: np.ndarray,
    device: torch.device,
    n_trials: int = 20,
    epochs_per_trial: int = 30,
    patience: int = 10,
    batch_size: int = 32,
    n_inner_folds: int = 3,
    use_mpnn: bool = False,
    chemeleon_cache: Optional[str] = None,
    unimol_cache: Optional[str] = None,
    moleculestm_cache: Optional[str] = None,
    mole_cache: Optional[str] = None,
    seed: int = 42,
    study_path: Optional[Path] = None,
    pretrain_state_dict: Optional[Dict] = None,
    pretrain_config: Optional[Dict] = None,
    load_delivery_head: bool = True,
    rdkit_cache: Optional[Dict] = None,
    use_moe: bool = False,
    moe_n_experts: int = 4,
    moe_top_k: int = 2,
    moe_expert_hidden_mult: int = 2,
    moe_jitter_noise: float = 0.0,
    set_transformer_block: str = "sab",
    llm_kwargs: Optional[Dict] = None,
) -> Tuple[Dict, int, optuna.Study]:
    """
    在内层数据上运行 Optuna 超参搜索。
    
    Args:
        full_dataset: 完整数据集
        inner_train_indices: 内层训练数据的索引（相对于 full_dataset）
        strata: 每个样本的分层标签
        device: 设备
        n_trials: Optuna 试验数
        epochs_per_trial: 每个试验的最大 epoch
        patience: 早停耐心值
        batch_size: 批次大小
        n_inner_folds: 内层折数
        use_mpnn: 是否使用 MPNN
        seed: 随机种子
        study_path: 可选的 study 持久化路径
        pretrain_state_dict: 预训练权重
        pretrain_config: 预训练配置
        load_delivery_head: 是否加载 delivery head 权重
    
    Returns:
        (best_params, epoch_mean, study)
    """
    if optuna is None:
        raise ImportError("Optuna not installed. Run: pip install optuna")
    
    inner_strata = strata[inner_train_indices]
    
    # 固定架构参数（与预训练一致，确保权重完整加载）
    _cfg = pretrain_config or {}
    fixed_d_model = _cfg.get("d_model", 256)
    fixed_num_heads = _cfg.get("num_heads", 8)
    fixed_n_attn_layers = _cfg.get("n_attn_layers", 4)
    fixed_fusion_strategy = _cfg.get("fusion_strategy", "attention")
    fixed_head_hidden_dim = _cfg.get("head_hidden_dim", 128)
    logger.info(
        f"Fixed architecture params: d_model={fixed_d_model}, num_heads={fixed_num_heads}, "
        f"n_attn_layers={fixed_n_attn_layers}, fusion={fixed_fusion_strategy}, "
        f"head_hidden_dim={fixed_head_hidden_dim}"
    )
    
    def objective(trial: optuna.Trial) -> float:
        d_model = fixed_d_model
        num_heads = fixed_num_heads
        n_attn_layers = fixed_n_attn_layers
        fusion_strategy = fixed_fusion_strategy
        head_hidden_dim = fixed_head_hidden_dim
        
        # 搜索训练超参数
        dropout = trial.suggest_float("dropout", 0.1, 0.5)
        lr = trial.suggest_float("lr", 1e-5, 1e-3, log=True)
        weight_decay = trial.suggest_float("weight_decay", 1e-5, 1e-1, log=True)
        backbone_lr_ratio = trial.suggest_float("backbone_lr_ratio", 0.01, 1.0, log=True)
        
        # MoE / LLM
        use_llm = bool(llm_kwargs.get("use_llm", False))
        moe_ne_t = trial.suggest_categorical("moe_n_experts", [2, 4]) if use_moe else moe_n_experts
        moe_tk_t = trial.suggest_int("moe_top_k", 1, 2) if use_moe else moe_top_k
        moe_hm_t = trial.suggest_categorical("moe_expert_hidden_mult", [1, 2]) if use_moe else moe_expert_hidden_mult
        llm_kwargs_t = dict(llm_kwargs)
        if use_llm:
            llm_kwargs_t["llm_use_lora"] = True
            llm_kwargs_t["llm_lora_r"] = trial.suggest_categorical("llm_lora_r", [8, 16, 32])

        # 内层 3-fold CV
        inner_cv = StratifiedKFold(
            n_splits=n_inner_folds, shuffle=True, random_state=seed
        )
        
        fold_val_losses = []
        fold_best_epochs = []
        
        for inner_fold, (inner_train_idx, inner_val_idx) in enumerate(
            inner_cv.split(inner_train_indices, inner_strata)
        ):
            # 获取实际的数据集索引
            actual_train_idx = inner_train_indices[inner_train_idx]
            actual_val_idx = inner_train_indices[inner_val_idx]
            
            # 创建 DataLoader
            train_subset = Subset(full_dataset, actual_train_idx.tolist())
            val_subset = Subset(full_dataset, actual_val_idx.tolist())
            
            train_loader = DataLoader(
                train_subset, batch_size=batch_size, shuffle=True,
                collate_fn=collate_fn, drop_last=True,
            )
            val_loader = DataLoader(
                val_subset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn
            )
            
            # 计算类权重
            class_weights = compute_class_weights_from_loader(train_loader)
            
            # 创建模型
            model = create_model(
                d_model=d_model,
                num_heads=num_heads,
                n_attn_layers=n_attn_layers,
                fusion_strategy=fusion_strategy,
                head_hidden_dim=head_hidden_dim,
                dropout=dropout,
                use_mpnn=use_mpnn,
                mpnn_device=device.type,
                chemeleon_cache=chemeleon_cache,
                unimol_cache=unimol_cache,
                moleculestm_cache=moleculestm_cache,
                mole_cache=mole_cache,
                use_moe=use_moe,
                moe_n_experts=moe_ne_t,
                moe_top_k=moe_tk_t,
                moe_expert_hidden_mult=moe_hm_t,
                moe_jitter_noise=moe_jitter_noise,
                set_transformer_block=set_transformer_block,
                llm_kwargs=llm_kwargs_t,
            )
            if rdkit_cache is not None:
                model.rdkit_encoder._cache = rdkit_cache

            # 内层 RAG 池：只用内层训练集（防 inner-val 泄漏）
            if llm_kwargs_t.get("use_rag", False) and getattr(model, "llm_prompt", None) is not None:
                _s, _d, _ex = _build_rag_pool(full_dataset, actual_train_idx)
                model.llm_prompt.set_retrieval_pool(
                    _s, _d, pool_id=f"inner{inner_fold}", extra_labels=_ex)
            
            # 加载预训练权重
            if pretrain_state_dict is not None and pretrain_config is not None:
                load_pretrain_weights_to_model(
                    model, pretrain_state_dict, d_model, pretrain_config, load_delivery_head
                )
            
            # 训练（带早停）
            result = train_with_early_stopping(
                model=model,
                train_loader=train_loader,
                val_loader=val_loader,
                device=device,
                lr=lr,
                weight_decay=weight_decay,
                epochs=epochs_per_trial,
                patience=patience,
                class_weights=class_weights,
                backbone_lr_ratio=backbone_lr_ratio,
            )
            
            fold_val_losses.append(result["best_val_loss"])
            fold_best_epochs.append(result["best_epoch"])
        
        # 记录 epoch_mean 到 trial
        epoch_mean = int(round(np.mean(fold_best_epochs)))
        trial.set_user_attr("epoch_mean", epoch_mean)
        trial.set_user_attr("fold_best_epochs", fold_best_epochs)
        
        return np.mean(fold_val_losses)
    
    # 创建 study
    storage = None
    if study_path is not None:
        storage = f"sqlite:///{study_path}"
    
    study = optuna.create_study(
        direction="minimize",
        sampler=TPESampler(seed=seed),
        storage=storage,
        study_name="inner_optuna",
        load_if_exists=True,
    )
    
    study.optimize(objective, n_trials=n_trials, show_progress_bar=True)
    
    best_params = dict(study.best_trial.params)
    best_params.update({
        "d_model": fixed_d_model,
        "num_heads": fixed_num_heads,
        "n_attn_layers": fixed_n_attn_layers,
        "fusion_strategy": fixed_fusion_strategy,
        "head_hidden_dim": fixed_head_hidden_dim,
        "set_transformer_block": set_transformer_block,
    })
    epoch_mean = study.best_trial.user_attrs.get("epoch_mean", epochs_per_trial)
    
    logger.info(f"Best trial: {study.best_trial.number}")
    logger.info(f"Best val_loss: {study.best_trial.value:.4f}")
    logger.info(f"Best params: {best_params}")
    logger.info(f"Epoch mean: {epoch_mean}")
    
    return best_params, epoch_mean, study


# ============ 单 fold 执行（可跨进程调用） ============
def _run_single_outer_fold(
    outer_fold: int,
    outer_train_idx: np.ndarray,
    outer_test_idx: np.ndarray,
    df: pd.DataFrame,
    strata: np.ndarray,
    fold_dir: Path,
    n_trials: int,
    epochs_per_trial: int,
    inner_patience: int,
    batch_size: int,
    n_inner_folds: int,
    use_mpnn: bool,
    chemeleon_cache: Optional[str],
    unimol_cache: Optional[str],
    moleculestm_cache: Optional[str],
    mole_cache: Optional[str],
    seed: int,
    pretrain_state_dict: Optional[Dict],
    pretrain_config: Optional[Dict],
    load_delivery_head: bool,
    device_str: str,
    use_retrieval: bool = False,
    retrieval_source: str = "internal",
    use_moe: bool = False,
    moe_n_experts: int = 4,
    moe_top_k: int = 2,
    moe_expert_hidden_mult: int = 2,
    moe_jitter_noise: float = 0.0,
    set_transformer_block: str = "sab",
    llm_kwargs: Optional[Dict] = None,
    precomputed_best_params: Optional[Dict] = None,
    precomputed_epoch_mean: Optional[int] = None,
) -> Dict:
    """
    执行单个外层 fold 的完整流程（内层调参 + 外层训练 + 评估）。

    所有参数均为可序列化类型，以支持 spawn 多进程。
    """
    device = torch.device(device_str)
    set_global_seed(seed + outer_fold)
    fold_dir = Path(fold_dir)
    fold_dir.mkdir(parents=True, exist_ok=True)

    # ===== 断点续跑 =====
    _f_metrics = fold_dir / "test_metrics.json"
    _f_bp = fold_dir / "best_params.json"
    _f_em = fold_dir / "epoch_mean.json"
    # (1) 整折已完成 → 读回结果直接跳过
    if _f_metrics.exists() and _f_bp.exists() and _f_em.exists():
        logger.info(f"[RESUME] Outer fold {outer_fold} 已完成，跳过。")
        with open(_f_metrics) as f:
            _tm = json.load(f)
        with open(_f_bp) as f:
            _bp = json.load(f)
        with open(_f_em) as f:
            _em = json.load(f)["epoch_mean"]
        return {"fold": outer_fold, "best_params": _bp,
                "epoch_mean": _em, "test_metrics": _tm}
    # (2) 内层已完成、外层中断 → 复用超参，跳过内层 Optuna
    if precomputed_best_params is None and _f_bp.exists() and _f_em.exists():
        with open(_f_bp) as f:
            precomputed_best_params = json.load(f)
        with open(_f_em) as f:
            precomputed_epoch_mean = json.load(f)["epoch_mean"]
        logger.info(f"[RESUME] fold {outer_fold} 复用已存 best_params，跳过内层 Optuna。")

    full_dataset = LNPDataset(df)
    # === 断点续跑:已完成的 fold 直接跳过(读回磁盘结果)===
    _tm = fold_dir / "test_metrics.json"
    _bp = fold_dir / "best_params.json"
    _em = fold_dir / "epoch_mean.json"
    if precomputed_best_params is None and _tm.exists() and _bp.exists() and _em.exists():
        logger.success(f"[SKIP] Outer fold {outer_fold} already done, loading cached results.")
        with open(_tm) as _f: _tmd = json.load(_f)
        with open(_bp) as _f: _bpd = json.load(_f)
        with open(_em) as _f: _emd = json.load(_f)
        return {
            "fold": outer_fold,
            "best_params": _bpd,
            "epoch_mean": int(_emd.get("epoch_mean", _emd) if isinstance(_emd, dict) else _emd),
            "test_metrics": _tmd,
        }

    logger.info(f"\n{'='*60}")
    logger.info(f"OUTER FOLD {outer_fold}")
    logger.info(f"{'='*60}")
    logger.info(f"Train: {len(outer_train_idx)}, Test: {len(outer_test_idx)}")

    # 预热 RDKit 缓存（在整个 fold 内共享）
    rdkit_cache = warmup_rdkit_cache(full_dataset.smiles)

    # 保存 split indices
    with open(fold_dir / "splits.json", "w") as f:
        json.dump({
            "outer_train_idx": outer_train_idx.tolist(),
            "outer_test_idx": outer_test_idx.tolist(),
        }, f)

    # 内层 Optuna 调参（方式 B：已有超参则跳过，直接复用）
    if precomputed_best_params is not None and precomputed_epoch_mean is not None:
        logger.info("Reusing precomputed best_params (skip inner Optuna).")
        best_params = dict(precomputed_best_params)
        epoch_mean = int(precomputed_epoch_mean)
    else:
        logger.info(f"\nRunning inner Optuna with {n_trials} trials...")
        study_path = fold_dir / "optuna_study.sqlite3"
        best_params, epoch_mean, study = run_inner_optuna(
            full_dataset=full_dataset,
            inner_train_indices=outer_train_idx,
            strata=strata,
            device=device,
            n_trials=n_trials,
            epochs_per_trial=epochs_per_trial,
            patience=inner_patience,
            batch_size=batch_size,
            n_inner_folds=n_inner_folds,
            use_mpnn=use_mpnn,
            chemeleon_cache=chemeleon_cache,
            unimol_cache=unimol_cache,
            moleculestm_cache=moleculestm_cache,
            mole_cache=mole_cache,
            seed=seed + outer_fold,
            study_path=study_path,
            pretrain_state_dict=pretrain_state_dict,
            pretrain_config=pretrain_config,
            load_delivery_head=load_delivery_head,
            rdkit_cache=rdkit_cache,
            use_moe=use_moe,
            moe_n_experts=moe_n_experts,
            moe_top_k=moe_top_k,
            moe_expert_hidden_mult=moe_expert_hidden_mult,
            moe_jitter_noise=moe_jitter_noise,
            set_transformer_block=set_transformer_block,
            llm_kwargs=llm_kwargs,
        )

    # 保存最佳参数
    with open(fold_dir / "best_params.json", "w") as f:
        json.dump(best_params, f, indent=2)
    with open(fold_dir / "epoch_mean.json", "w") as f:
        json.dump({"epoch_mean": epoch_mean}, f)

    # 外层训练（使用最优超参，固定 epoch 数，不 early-stop）
    logger.info(f"\nTraining outer fold with best params, epochs={epoch_mean}...")

    # 从 outer_train 再切 10% 作为"监控用 val"（只画曲线、看 gap，不参与选模型/早停）
    _rng = np.random.RandomState(seed + outer_fold)
    _perm = _rng.permutation(len(outer_train_idx))
    _n_val = max(1, int(0.1 * len(outer_train_idx)))
    fit_idx = outer_train_idx[_perm[_n_val:]]
    monitor_idx = outer_train_idx[_perm[:_n_val]]

    train_subset = Subset(full_dataset, fit_idx.tolist())
    monitor_subset = Subset(full_dataset, monitor_idx.tolist())
    test_subset = Subset(full_dataset, outer_test_idx.tolist())

    train_loader = DataLoader(
        train_subset, batch_size=batch_size, shuffle=True,
        collate_fn=collate_fn, drop_last=True,
    )
    monitor_loader = DataLoader(
        monitor_subset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn
    )
    test_loader = DataLoader(
        test_subset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn
    )

    class_weights = compute_class_weights_from_loader(train_loader)

    arch = {
        "moe_n_experts": best_params.get("moe_n_experts", moe_n_experts),
        "moe_top_k": best_params.get("moe_top_k", moe_top_k),
        "moe_expert_hidden_mult": best_params.get("moe_expert_hidden_mult", moe_expert_hidden_mult),
        "moe_jitter_noise": moe_jitter_noise,
        "set_transformer_block": best_params.get("set_transformer_block", "sab"),
    }
    llm_kwargs_resolved = {
        **(llm_kwargs or {}),
        **({"llm_use_lora": True, "llm_lora_r": best_params["llm_lora_r"]}
           if ((llm_kwargs or {}).get("use_llm", False) and "llm_lora_r" in best_params) else {}),
    }

    model = create_model(
        d_model=best_params["d_model"],
        num_heads=best_params["num_heads"],
        n_attn_layers=best_params["n_attn_layers"],
        fusion_strategy=best_params["fusion_strategy"],
        head_hidden_dim=best_params["head_hidden_dim"],
        dropout=best_params["dropout"],
        use_mpnn=use_mpnn,
        mpnn_device=device.type,
        chemeleon_cache=chemeleon_cache,
        unimol_cache=unimol_cache,
        moleculestm_cache=moleculestm_cache,
        mole_cache=mole_cache,
        use_moe=use_moe,
        llm_kwargs=llm_kwargs_resolved,
        use_retrieval=use_retrieval,
        retr_feature_dim=(3 if use_retrieval else 0),
        **arch,
    )
    model.rdkit_encoder._cache = rdkit_cache
    if use_retrieval:
        from lnp_ml.modeling.retrieval import MorganRetriever
        if retrieval_source == "external":
            _tr_smiles, _tr_labels = _load_external_retrieval_pool()
            logger.info(f"[retrieval] 外部检索池(mRNA+Mouse)={len(_tr_smiles)} 分子")
        else:
            _tr_smiles = [full_dataset.smiles[i] for i in outer_train_idx]
            _tr_labels = full_dataset.delivery[outer_train_idx].reshape(-1, 1)
            logger.info(f"[retrieval] 内部检索池={len(_tr_smiles)} 训练分子（不含测试集）")
        model.retriever = MorganRetriever(_tr_smiles, _tr_labels, k=5)

    # ===== RAG 检索池设置（防泄漏：只用训练集分子+标签）=====
    if llm_kwargs and llm_kwargs.get("use_rag", False) and getattr(model, "llm_prompt", None) is not None:
        _s, _d, _ex = _build_rag_pool(full_dataset, outer_train_idx)
        model.llm_prompt.set_retrieval_pool(
            _s, _d, pool_id=f"fold{outer_fold}", extra_labels=_ex)
        logger.info(f"[RAG] 检索池={len(_s)} 训练分子（多任务标签，不含测试集，防泄漏）")

    if pretrain_state_dict is not None and pretrain_config is not None:
        loaded = load_pretrain_weights_to_model(
            model, pretrain_state_dict, best_params["d_model"],
            pretrain_config, load_delivery_head
        )
        if loaded:
            logger.info(f"Loaded pretrain weights for outer fold {outer_fold}")

    train_result = train_fixed_epochs(
        model=model,
        train_loader=train_loader,
        val_loader=monitor_loader,
        device=device,
        lr=best_params["lr"],
        weight_decay=best_params["weight_decay"],
        epochs=epoch_mean,
        class_weights=class_weights,
        use_cosine_annealing=True,
        backbone_lr_ratio=best_params.get("backbone_lr_ratio", 1.0),
        freeze_backbone_epochs=3,
    )

    model.load_state_dict(train_result["final_state"], strict=False)
    model = model.to(device)

    config = {
        "d_model": best_params["d_model"],
        "num_heads": best_params["num_heads"],
        "n_attn_layers": best_params["n_attn_layers"],
        "fusion_strategy": best_params["fusion_strategy"],
        "head_hidden_dim": best_params["head_hidden_dim"],
        "dropout": best_params["dropout"],
        "use_mpnn": use_mpnn,
        "use_chemeleon": chemeleon_cache is not None,
        "chemeleon_cache": chemeleon_cache,
        "use_unimol": unimol_cache is not None,
        "unimol_cache": unimol_cache,
        "use_moleculestm": moleculestm_cache is not None,
        "moleculestm_cache": moleculestm_cache,
        "use_mole": mole_cache is not None,
        "mole_cache": mole_cache,
        "use_moe": use_moe,
        "use_retrieval": use_retrieval,
        "retr_feature_dim": (3 if use_retrieval else 0),
        **arch,
        **llm_kwargs_resolved,
    }

    _full_state = train_result["final_state"]
    _slim_state = {
        k: v for k, v in _full_state.items()
        if (not k.startswith("llm_prompt.encoder")) or ("lora_" in k)
    }
    _n_lora = sum(1 for k in _slim_state if "lora_" in k)
    logger.info(f"保存 checkpoint: {len(_slim_state)} 个张量，其中 LoRA 适配器 {_n_lora} 个")
    torch.save({
        "model_state_dict": _slim_state,
        "config": config,
        "epoch_mean": epoch_mean,
        "best_params": best_params,
    }, fold_dir / "model.pt")

    with open(fold_dir / "history.json", "w") as f:
        json.dump(train_result["history"], f, indent=2)

    # 在测试集上评估
    logger.info("Evaluating on outer test set...")
    test_metrics = evaluate_on_test(model, test_loader, device)

    with open(fold_dir / "test_metrics.json", "w") as f:
        json.dump(test_metrics, f, indent=2)

    logger.info(f"\nOuter Fold {outer_fold} Test Results:")
    for task, metrics in test_metrics.items():
        if "rmse" in metrics:
            logger.info(f"  {task}: RMSE={metrics['rmse']:.4f}, R²={metrics['r2']:.4f}")
        elif "accuracy" in metrics:
            logger.info(f"  {task}: Acc={metrics['accuracy']:.4f}, F1={metrics['f1']:.4f}")
        elif "kl_divergence" in metrics:
            logger.info(f"  {task}: KL={metrics['kl_divergence']:.4f}, JS={metrics['js_divergence']:.4f}")

    return {
        "fold": outer_fold,
        "best_params": best_params,
        "epoch_mean": epoch_mean,
        "test_metrics": test_metrics,
    }


# ============ 主流程 ============
@app.command()
def main(
    input_path: Path = INTERIM_DATA_DIR / "internal.csv",
    output_dir: Path = MODELS_DIR / "nested_cv",
    resume_dir: Optional[Path] = None,
    # CV 参数
    n_outer_folds: int = 5,
    n_inner_folds: int = 3,
    min_stratum_count: int = 5,
    seed: int = 42,
    # Optuna 参数
    n_trials: int = 20,
    epochs_per_trial: int = 30,
    inner_patience: int = 10,
    # 训练参数
    batch_size: int = 32,
    # 预训练权重
    init_from_pretrain: Optional[Path] = None,
    load_delivery_head: bool = False,
    # MPNN
    use_mpnn: bool = False,
    # CheMeleon
    use_chemeleon: bool = False,
    chemeleon_cache: str = "data/processed/chemeleon_embeddings.npz",
    use_unimol: bool = False,
    unimol_cache: str = "data/processed/unimol_embeddings.npz",
    use_moleculestm: bool = False,
    moleculestm_cache: str = "data/processed/moleculestm_embeddings.npz",
    use_mole: bool = False,
    mole_cache: str = "data/processed/mole_embeddings.npz",
    n_repeats: int = 1,
    repeat_seed_step: int = 1000,
    # MoE（消融开关）
    use_moe: bool = False,
    moe_n_experts: int = 4,
    moe_top_k: int = 2,
    moe_expert_hidden_mult: int = 2,
    moe_jitter_noise: float = 0.0,
    # Set Transformer
    set_transformer_block: str = "sab",
    # 回归旁路（消融开关）
    reg_bypass: str = "on",
    # LLM（消融开关）
    use_llm: bool = False,
    use_rag: bool = False,
    rag_top_k: int = 4,
    llm_max_length: int = 1536,
    use_retrieval: bool = False,
    retrieval_source: str = "internal",
    llm_model_path: str = DEFAULT_MOLT5_PATH,
    llm_freeze: bool = True,
    llm_use_lora: bool = False,
    llm_use_qlora: bool = False,
    use_soft_prompt: bool = False,
    llm_lora_r: int = 8,
    llm_lora_alpha: int = 16,
    llm_lora_dropout: float = 0.05,
    fix_hparams_json: Optional[Path] = None,
    fix_epoch_mean: int = 12,
    # 并行
    parallel: bool = False,
    # 设备
    device: str = "cuda" if torch.cuda.is_available() else "cpu",
):
    """
    嵌套交叉验证 + Optuna 超参调优。
    
    外层 5-fold（20% test / 80% train），内层 3-fold Optuna 调参。
    外层训练不使用 early-stopping，epoch 数使用内层 best trial 的 epoch_mean。
    
    使用 --init-from-pretrain 从预训练 checkpoint 初始化模型权重。
    使用 --parallel 同时运行所有外层 fold（需要足够 GPU 显存）。
    """
    if optuna is None:
        logger.error("Optuna not installed. Run: pip install optuna")
        raise typer.Exit(1)
    
    logger.info(f"Using device: {device}")
    device = torch.device(device)
    set_global_seed(seed)

    llm_kwargs = dict(
        reg_bypass=reg_bypass,
        use_rag=use_rag,
        rag_top_k=rag_top_k,
        llm_max_length=llm_max_length,
        use_llm=use_llm,
        llm_model_path=llm_model_path,
        llm_freeze=llm_freeze,
        llm_use_lora=llm_use_lora,
        llm_use_qlora=llm_use_qlora,
        use_soft_prompt=use_soft_prompt,
        llm_lora_r=llm_lora_r,
        llm_lora_alpha=llm_lora_alpha,
        llm_lora_dropout=llm_lora_dropout,
    )

    _fixed_bp = None
    _fixed_em = None
    if fix_hparams_json is not None:
        _fixed_bp = json.loads(Path(fix_hparams_json).read_text())
        _fixed_em = int(fix_epoch_mean)
        if use_moe:
            _fixed_bp["moe_n_experts"] = moe_n_experts
            _fixed_bp["moe_top_k"] = moe_top_k
        logger.info(f"[FIX] 固定超参={_fixed_bp}, epoch_mean={_fixed_em}, 跳过内层 Optuna")

    # 加载预训练权重
    pretrain_state_dict = None
    pretrain_config = None
    if init_from_pretrain is not None:
        if init_from_pretrain.exists():
            logger.info(f"Loading pretrain weights from {init_from_pretrain}")
            checkpoint = torch.load(init_from_pretrain, map_location="cpu", weights_only=False)
            pretrain_state_dict = checkpoint["model_state_dict"]
            pretrain_config = checkpoint.get("config", {})
            logger.success(f"Loaded pretrain checkpoint (d_model={pretrain_config.get('d_model')})")
        else:
            logger.warning(f"Pretrain checkpoint not found: {init_from_pretrain}, skipping")
    
    # 创建输出目录（带时间戳；--resume-dir 指定则复用，支持断点续跑）
    if resume_dir is not None:
        run_dir = Path(resume_dir)
        run_dir.mkdir(parents=True, exist_ok=True)
        logger.info(f"[RESUME] 复用已有运行目录: {run_dir}")
    else:
        run_name = datetime.now().strftime("%Y%m%d_%H%M%S")
        run_dir = output_dir / run_name
        run_dir.mkdir(parents=True, exist_ok=True)
    logger.info(f"Output directory: {run_dir}")
    
    # 加载数据
    logger.info(f"Loading data from {input_path}")
    df = pd.read_csv(input_path)
    logger.info(f"Loaded {len(df)} samples")
    
    # 处理数据
    logger.info("Processing dataframe...")
    df = process_dataframe(df)
    
    # 构建复合分层标签
    logger.info("Building composite strata...")
    strata, strata_info = build_composite_strata(df, min_stratum_count)
    
    # 保存 strata 信息
    with open(run_dir / "strata_info.json", "w") as f:
        json.dump(strata_info, f, indent=2, default=str)
    
    # 创建完整数据集（仅用于获取样本数做 split）
    n_samples = len(LNPDataset(df))

    # 外层 CV split
    outer_cv = StratifiedKFold(
        n_splits=n_outer_folds, shuffle=True, random_state=seed
    )

    device_str = str(device)
    fold_args = []
    for outer_fold, (outer_train_idx, outer_test_idx) in enumerate(
        outer_cv.split(np.arange(n_samples), strata)
    ):
        fold_args.append(dict(
            outer_fold=outer_fold,
            outer_train_idx=outer_train_idx,
            outer_test_idx=outer_test_idx,
            df=df,
            strata=strata,
            fold_dir=run_dir / f"outer_fold_{outer_fold}",
            n_trials=n_trials,
            epochs_per_trial=epochs_per_trial,
            inner_patience=inner_patience,
            batch_size=batch_size,
            n_inner_folds=n_inner_folds,
            use_mpnn=use_mpnn,
            chemeleon_cache=(chemeleon_cache if use_chemeleon else None),
            unimol_cache=(unimol_cache if use_unimol else None),
            moleculestm_cache=(moleculestm_cache if use_moleculestm else None),
            mole_cache=(mole_cache if use_mole else None),
            seed=seed,
            pretrain_state_dict=pretrain_state_dict,
            pretrain_config=pretrain_config,
            load_delivery_head=load_delivery_head,
            device_str=device_str,
            use_retrieval=use_retrieval,
            retrieval_source=retrieval_source,
            use_moe=use_moe,
            moe_n_experts=moe_n_experts,
            moe_top_k=moe_top_k,
            moe_expert_hidden_mult=moe_expert_hidden_mult,
            moe_jitter_noise=moe_jitter_noise,
            set_transformer_block=set_transformer_block,
            llm_kwargs=llm_kwargs,
            precomputed_best_params=_fixed_bp,
            precomputed_epoch_mean=_fixed_em,
        ))
    
    if parallel:
        import multiprocessing as mp
        from concurrent.futures import ProcessPoolExecutor
        
        ctx = mp.get_context("spawn")
        logger.info(f"Running {n_outer_folds} outer folds in PARALLEL (spawn)")
        with ProcessPoolExecutor(max_workers=n_outer_folds, mp_context=ctx) as executor:
            futures = [executor.submit(_run_single_outer_fold, **args) for args in fold_args]
            outer_results = [f.result() for f in futures]
        outer_results.sort(key=lambda r: r["fold"])
    else:
        logger.info(f"Running {n_outer_folds} outer folds SEQUENTIALLY")
        outer_results = []
        per_fold_params = {}
        for args in fold_args:
            result = _run_single_outer_fold(**args)
            outer_results.append(result)
            per_fold_params[result["fold"]] = (result["best_params"], result["epoch_mean"])

        for rep in range(1, n_repeats):
            logger.info(f"\n===== Repeat {rep}/{n_repeats - 1} (reuse hyperparams) =====")
            for args in fold_args:
                bp, em = per_fold_params[args["outer_fold"]]
                rep_args = dict(args)
                rep_args["seed"] = seed + rep * repeat_seed_step
                rep_args["fold_dir"] = run_dir / f"repeat_{rep}" / f"outer_fold_{args['outer_fold']}"
                rep_args["precomputed_best_params"] = bp
                rep_args["precomputed_epoch_mean"] = em
                outer_results.append(_run_single_outer_fold(**rep_args))
    
    # 汇总结果
    logger.info("\n" + "=" * 60)
    logger.info("NESTED CV COMPLETE")
    logger.info("=" * 60)
    
    # 计算汇总统计
    summary = {"fold_results": outer_results}
    
    # 对每个任务计算均值和标准差
    tasks_with_metrics = {}
    for result in outer_results:
        for task, metrics in result["test_metrics"].items():
            if task not in tasks_with_metrics:
                tasks_with_metrics[task] = {k: [] for k in metrics.keys() if k != "n_samples"}
            for k, v in metrics.items():
                if k != "n_samples":
                    tasks_with_metrics[task][k].append(v)
    
    summary["summary_stats"] = {}
    for task, metrics_dict in tasks_with_metrics.items():
        summary["summary_stats"][task] = {}
        for metric_name, values in metrics_dict.items():
            summary["summary_stats"][task][f"{metric_name}_mean"] = float(np.mean(values))
            summary["summary_stats"][task][f"{metric_name}_std"] = float(np.std(values))
    
    # 打印汇总
    logger.info("\n[Summary Statistics]")
    for task, stats in summary["summary_stats"].items():
        if "rmse_mean" in stats:
            logger.info(
                f"  {task}: RMSE={stats['rmse_mean']:.4f}±{stats['rmse_std']:.4f}, "
                f"R²={stats['r2_mean']:.4f}±{stats['r2_std']:.4f}"
            )
        elif "accuracy_mean" in stats:
            logger.info(
                f"  {task}: Acc={stats['accuracy_mean']:.4f}±{stats['accuracy_std']:.4f}, "
                f"F1={stats['f1_mean']:.4f}±{stats['f1_std']:.4f}"
            )
        elif "kl_divergence_mean" in stats:
            logger.info(
                f"  {task}: KL={stats['kl_divergence_mean']:.4f}±{stats['kl_divergence_std']:.4f}, "
                f"JS={stats['js_divergence_mean']:.4f}±{stats['js_divergence_std']:.4f}"
            )
    
    # 保存汇总
    with open(run_dir / "summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    
    logger.success(f"\nAll results saved to {run_dir}")


if __name__ == "__main__":
    app()