"""LNP 多任务预测模型"""

import torch
import torch.nn as nn
from typing import Dict, List, Optional, Literal

from lnp_ml.modeling.encoders import CachedRDKitEncoder, CachedMPNNEncoder
from lnp_ml.modeling.layers import (
    TokenProjector,
    SetTransformer,
    ResidualConcatFusion,
    MoEBlock,
    LLMPromptEncoder,
)
from lnp_ml.modeling.layers.llm_prompt import DEFAULT_MOLT5_PATH
from lnp_ml.modeling.heads import MultiTaskHead


PoolingStrategy = Literal["attention", "avg", "max"]


# Token 维度配置
DEFAULT_INPUT_DIMS = {
    # Channel A: 化学特征
    "mpnn": 600,      # D-MPNN embedding
    "morgan": 1024,   # Morgan fingerprint
    "maccs": 167,     # MACCS keys
    "desc": 217,      # RDKit descriptors
    # Channel B: 配方/实验条件
    "comp": 5,        # 配方比例
    "phys": 12,       # 物理参数 one-hot
    "help": 4,        # Helper lipid one-hot
    "exp": 32,        # 实验条件 one-hot
}

# 化学 / 配方 token 的键顺序
CHEM_KEYS_WITH_MPNN = ["mpnn", "morgan", "maccs", "desc"]
CHEM_KEYS_NO_MPNN = ["morgan", "maccs", "desc"]
TAB_KEYS = ["comp", "phys", "help", "exp"]

# backbone 权重前缀（用于预训练加载与导出）
BACKBONE_PREFIXES = (
    "token_projector.",
    "set_transformer.",
    "fusion.",
    "moe.",
    "llm_prompt.",
)
# 冻结的 MolT5 encoder 权重前缀，不纳入 backbone（由本地权重加载，不进 checkpoint）
LLM_FROZEN_PREFIX = "llm_prompt.encoder."


class LNPModel(nn.Module):
    """
    LNP 药物递送性能预测模型。

    架构流程:
        1. Encoders: SMILES -> 化学特征; tabular -> 配方/实验特征
        2. TokenProjector: 统一到 d_model
        3. SetTransformer: 对化学 token 集合做置换等变编码 -> chem'
        4. MoE (可选): router 看 tab，expert 吃 chem' -> F_moe
        5. LLM (可选): chem'/tab 注入 MolT5 prompt -> F_llm
        6. ResidualConcatFusion: 拼接 chem'/tab/F_moe/F_llm -> attention pooling
        7. MultiTaskHead: 多任务预测
    """

    def __init__(
        self,
        # 模型维度
        d_model: int = 256,
        # Set Transformer
        num_heads: int = 8,
        n_attn_layers: int = 4,
        set_transformer_block: str = "sab",
        # Fusion
        fusion_strategy: PoolingStrategy = "attention",
        # Head
        head_hidden_dim: int = 128,
        # Dropout
        dropout: float = 0.1,
        # MPNN encoder
        mpnn_checkpoint: Optional[str] = None,
        mpnn_ensemble_paths: Optional[List[str]] = None,
        mpnn_device: str = "cpu",
        # 输入维度配置
        input_dims: Optional[Dict[str, int]] = None,
        # ============ 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 相关 ============
        use_llm: bool = False,
        llm_model_path: str = DEFAULT_MOLT5_PATH,
        llm_freeze: bool = True,
        llm_use_lora: bool = False,
        llm_lora_r: int = 8,
        llm_lora_alpha: int = 16,
        llm_lora_dropout: float = 0.05,
        use_rag: bool = False,
        rag_top_k: int = 4,
        use_retrieval: bool = False,
        retr_feature_dim: int = 0,
    ) -> None:
        super().__init__()

        self.input_dims = input_dims or DEFAULT_INPUT_DIMS
        self.d_model = d_model
        self.use_mpnn = mpnn_checkpoint is not None or mpnn_ensemble_paths is not None

        # ============ Encoders ============
        self.rdkit_encoder = CachedRDKitEncoder()
        if self.use_mpnn:
            self.mpnn_encoder = CachedMPNNEncoder(
                checkpoint_path=mpnn_checkpoint,
                ensemble_paths=mpnn_ensemble_paths,
                device=mpnn_device,
            )
        else:
            self.mpnn_encoder = None

        # ============ Token Projector ============
        proj_input_dims = {k: v for k, v in self.input_dims.items()}
        if not self.use_mpnn:
            proj_input_dims.pop("mpnn", None)
        self.token_projector = TokenProjector(
            input_dims=proj_input_dims,
            d_model=d_model,
            dropout=dropout,
        )

        # token 顺序与化学侧 token 数
        self.chem_keys = CHEM_KEYS_WITH_MPNN if self.use_mpnn else CHEM_KEYS_NO_MPNN
        self.tab_keys = TAB_KEYS
        self.token_order = self.chem_keys + self.tab_keys
        self.split_idx = len(self.chem_keys)

        # ============ Set Transformer ============
        self.set_transformer = SetTransformer(
            d_model=d_model,
            num_heads=num_heads,
            n_layers=n_attn_layers,
            dropout=dropout,
            block=set_transformer_block,
        )

        # ============ MoE Block (可选) ============
        self.use_moe = use_moe
        self._last_moe_extras: Optional[Dict[str, torch.Tensor]] = None
        if use_moe:
            self.moe = MoEBlock(
                d_model=d_model,
                n_chem_tokens=self.split_idx,
                n_experts=moe_n_experts,
                top_k=moe_top_k,
                expert_hidden_mult=moe_expert_hidden_mult,
                dropout=dropout,
                jitter_noise=moe_jitter_noise,
            )
        else:
            self.moe = None

        # ============ LLM Prompt (可选) ============
        self.use_llm = use_llm
        if use_llm:
            self.llm_prompt = LLMPromptEncoder(
                d_model=d_model,
                n_chem_tokens=self.split_idx,
                n_cond_tokens=len(self.tab_keys),
                model_name_or_path=llm_model_path,
                freeze=llm_freeze,
                use_lora=llm_use_lora,
                lora_r=llm_lora_r,
                lora_alpha=llm_lora_alpha,
                lora_dropout=llm_lora_dropout,
                use_rag=use_rag,
                rag_top_k=rag_top_k,
            )
        else:
            self.llm_prompt = None

        # ============ 检索增强 (可选) ============
        self.use_retrieval = use_retrieval
        self.retriever = None  # 由 nested_cv 每折用训练集构建后赋值
        if use_retrieval:
            self.retr_proj = nn.Linear(retr_feature_dim, d_model)
        else:
            self.retr_proj = None

        # ============ Residual Concat + Fusion ============
        self.fusion = ResidualConcatFusion(d_model=d_model, strategy=fusion_strategy)

        # ============ Multi-Task Head ============
        self.head = MultiTaskHead(
            in_dim=self.fusion.fusion_dim,
            hidden_dim=head_hidden_dim,
            dropout=dropout,
        )

    def _encode_and_project(
        self,
        smiles: List[str],
        tabular: Dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """
        编码 SMILES 和 tabular，返回 stacked tokens。

        Returns:
            stacked: [B, n_tokens, d_model]，顺序为 chem 在前、tab 在后
        """
        device = tabular["comp"].device

        rdkit_features = self.rdkit_encoder(smiles)

        all_features: Dict[str, torch.Tensor] = {}
        if self.use_mpnn:
            mpnn_features = self.mpnn_encoder(smiles)
            all_features["mpnn"] = mpnn_features["mpnn"].to(device)
        all_features["morgan"] = rdkit_features["morgan"].to(device)
        all_features["maccs"] = rdkit_features["maccs"].to(device)
        all_features["desc"] = rdkit_features["desc"].to(device)
        all_features["comp"] = tabular["comp"]
        all_features["phys"] = tabular["phys"]
        all_features["help"] = tabular["help"]
        all_features["exp"] = tabular["exp"]

        projected = self.token_projector(all_features)
        stacked = torch.stack([projected[k] for k in self.token_order], dim=1)
        return stacked

    def _backbone_from_stacked(
        self, stacked: torch.Tensor, smiles: Optional[List[str]] = None
    ) -> torch.Tensor:
        chem = stacked[:, : self.split_idx, :]
        tab = stacked[:, self.split_idx :, :]

        chem = self.set_transformer(chem)  # chem'

        f_moe = None
        if self.moe is not None:
            f_moe, extras = self.moe(chem, tab)
            self._last_moe_extras = extras
        else:
            self._last_moe_extras = None

        f_llm = None
        if self.llm_prompt is not None and smiles is not None:
            f_llm = self.llm_prompt(smiles)

        f_retr = None
        if self.retr_proj is not None and self.retriever is not None and smiles is not None:
            _feats = self.retriever.query_batch(smiles, exclude_self="auto")
            _feats_t = torch.as_tensor(_feats, dtype=chem.dtype, device=chem.device)
            f_retr = self.retr_proj(_feats_t)

        return self.fusion(chem, tab, f_moe=f_moe, f_llm=f_llm, f_retr=f_retr)

    def forward_from_projected(
        self,
        stacked: torch.Tensor,
        task: Optional[str] = None,
    ) -> torch.Tensor:
        """
        从已投影的 stacked tokens 开始 forward，用于 Captum 归因。

        Args:
            stacked: [B, n_tokens, d_model]
            task: 单任务名；None 时返回 delivery head 输出。

        Returns:
            对应任务的预测输出。
        """
        fused = self._backbone_from_stacked(stacked)

        if task is None:
            task = "delivery"

        task_heads = {
            "size": self.head.size_head,
            "pdi": self.head.pdi_head,
            "ee": self.head.ee_head,
            "delivery": self.head.delivery_head,
            "biodist": self.head.biodist_head,
            "toxic": self.head.toxic_head,
        }
        return task_heads[task](fused)

    def forward_replacing_token(
        self,
        raw_feature: torch.Tensor,
        feature_key: str,
        base_projected: torch.Tensor,
        task: Optional[str] = None,
    ) -> torch.Tensor:
        """
        用原始特征替换 base_projected 中指定 token 的投影，然后 forward。

        用于对单个 token 内部特征做 Captum 归因（如 desc 的 210 维）。
        """
        projected = self.token_projector.projectors[feature_key](raw_feature)
        gate = torch.sigmoid(self.token_projector.weights[feature_key])
        projected = projected * gate

        token_order = list(self.token_projector.keys)
        token_idx = token_order.index(feature_key)

        stacked = base_projected.clone()
        stacked[:, token_idx, :] = projected

        return self.forward_from_projected(stacked, task=task)

    def forward_backbone(
        self,
        smiles: List[str],
        tabular: Dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """Backbone forward：编码 -> 投影 -> set transformer -> (MoE/LLM) -> 融合。"""
        stacked = self._encode_and_project(smiles, tabular)
        return self._backbone_from_stacked(stacked, smiles=smiles)

    def forward_delivery(
        self,
        smiles: List[str],
        tabular: Dict[str, torch.Tensor],
    ) -> torch.Tensor:
        """仅预测 delivery（用于 pretrain）。返回 [B, 1]。"""
        fused = self.forward_backbone(smiles, tabular)
        return self.head.delivery_head(fused)

    def forward(
        self,
        smiles: List[str],
        tabular: Dict[str, torch.Tensor],
    ) -> Dict[str, torch.Tensor]:
        """
        完整的多任务 forward。

        Returns:
            Dict[str, Tensor]: size [B,1], pdi [B,4], ee [B,3],
            delivery [B,1], biodist [B,7], toxic [B,2]
        """
        fused = self.forward_backbone(smiles, tabular)
        return self.head(fused)

    def clear_cache(self) -> None:
        """清空所有 encoder 的缓存"""
        self.rdkit_encoder.clear_cache()
        if self.mpnn_encoder is not None:
            self.mpnn_encoder.clear_cache()
        if self.llm_prompt is not None and hasattr(self.llm_prompt, "clear_cache"):
            self.llm_prompt.clear_cache()

    def get_last_moe_extras(self) -> Optional[Dict[str, torch.Tensor]]:
        """返回最近一次 forward 中 MoE 模块的副产物（aux loss、gates、probs）。"""
        return self._last_moe_extras

    def get_backbone_state_dict(self) -> Dict[str, torch.Tensor]:
        """
        获取 backbone 部分的 state_dict（不含任务头，且排除冻结的 MolT5 encoder）。
        """
        return {
            k: v for k, v in self.state_dict().items()
            if k.startswith(BACKBONE_PREFIXES) and not k.startswith(LLM_FROZEN_PREFIX)
        }

    def get_delivery_head_state_dict(self) -> Dict[str, torch.Tensor]:
        """获取 delivery head 的 state_dict"""
        return {
            k: v for k, v in self.state_dict().items()
            if k.startswith("head.delivery_head.")
        }

    def load_pretrain_weights(
        self,
        pretrain_state_dict: Dict[str, torch.Tensor],
        load_delivery_head: bool = True,
        strict: bool = False,
    ) -> None:
        """
        从预训练 checkpoint 加载 backbone 和（可选）delivery head 权重。

        冻结的 MolT5 encoder 权重不在加载范围内。
        """
        keys_to_load = []
        for name in pretrain_state_dict.keys():
            if name.startswith(BACKBONE_PREFIXES) and not name.startswith(LLM_FROZEN_PREFIX):
                keys_to_load.append(name)
            elif load_delivery_head and name.startswith("head.delivery_head."):
                keys_to_load.append(name)

        filtered_state_dict = {
            k: v for k, v in pretrain_state_dict.items() if k in keys_to_load
        }

        unexpected = []
        model_state = self.state_dict()
        for k, v in filtered_state_dict.items():
            if k in model_state and model_state[k].shape == v.shape:
                model_state[k] = v
            else:
                unexpected.append(k)


        self.load_state_dict(model_state, strict=False)

        if strict and unexpected:
            raise RuntimeError(f"Unexpected keys: {unexpected}")


class LNPModelWithoutMPNN(LNPModel):
    """不使用 MPNN 的简化版本（化学 token 为 3 个）"""

    def __init__(
        self,
        d_model: int = 256,
        num_heads: int = 8,
        n_attn_layers: int = 4,
        set_transformer_block: str = "sab",
        fusion_strategy: PoolingStrategy = "attention",
        head_hidden_dim: int = 128,
        dropout: float = 0.1,
        input_dims: Optional[Dict[str, int]] = None,
        # ============ 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 相关 ============
        use_llm: bool = False,
        llm_model_path: str = DEFAULT_MOLT5_PATH,
        llm_freeze: bool = True,
        llm_use_lora: bool = False,
        llm_lora_r: int = 8,
        llm_lora_alpha: int = 16,
        llm_lora_dropout: float = 0.05,
        use_rag: bool = False,
        rag_top_k: int = 4,
        use_retrieval: bool = False,
        retr_feature_dim: int = 0,
    ) -> None:
        dims = input_dims or DEFAULT_INPUT_DIMS.copy()
        dims.pop("mpnn", None)

        super().__init__(
            d_model=d_model,
            num_heads=num_heads,
            n_attn_layers=n_attn_layers,
            set_transformer_block=set_transformer_block,
            fusion_strategy=fusion_strategy,
            head_hidden_dim=head_hidden_dim,
            dropout=dropout,
            mpnn_checkpoint=None,
            mpnn_ensemble_paths=None,
            input_dims=dims,
            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_llm=use_llm,
            use_rag=use_rag,
            rag_top_k=rag_top_k,
            use_retrieval=use_retrieval,
            retr_feature_dim=retr_feature_dim,
            llm_model_path=llm_model_path,
            llm_freeze=llm_freeze,
            llm_use_lora=llm_use_lora,
            llm_lora_r=llm_lora_r,
            llm_lora_alpha=llm_lora_alpha,
            llm_lora_dropout=llm_lora_dropout,
        )