"""LLM 分子特征分支：用 MolT5 直接编码 SMILES 文本。"""

import os
from typing import Dict, List, Optional

import torch
import torch.nn as nn


# 权重默认路径，可用环境变量 MOLT5_PATH 覆盖
DEFAULT_MOLT5_PATH = os.environ.get("MOLT5_PATH", "models/molt5-base")


class LLMPromptEncoder(nn.Module):
    """用 MolT5 编码 SMILES 文本，输出分子特征 F_llm [B, d_model]。

    - 冻结 encoder 时：对每个 SMILES 缓存其句向量，避免重复前向，降方差、提速。
    - use_lora=True 时：encoder 可训练（LoRA），不缓存。
    """

    def __init__(
        self,
        d_model: int,
        n_chem_tokens: int = 4,
        n_cond_tokens: int = 4,
        model_name_or_path: str = DEFAULT_MOLT5_PATH,
        freeze: bool = True,
        use_lora: bool = False,
        lora_r: int = 8,
        lora_alpha: int = 16,
        lora_dropout: float = 0.05,
        max_length: int = 128,
    ) -> None:
        super().__init__()
        from transformers import AutoTokenizer, T5EncoderModel

        self.use_lora = use_lora
        self.max_length = max_length
        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
        self.encoder = T5EncoderModel.from_pretrained(model_name_or_path)
        self.hidden_size = self.encoder.config.d_model

        if use_lora:
            self._apply_lora(lora_r, lora_alpha, lora_dropout)
        elif freeze:
            for p in self.encoder.parameters():
                p.requires_grad = False
        self._frozen = freeze and not use_lora

        self.proj_down = nn.Sequential(
            nn.Linear(self.hidden_size, d_model), nn.LayerNorm(d_model)
        )
        # 冻结特征缓存：smiles -> [H]（CPU）
        self._cache: Dict[str, torch.Tensor] = {}

    def _apply_lora(self, r: int, alpha: int, dropout: float) -> None:
        from peft import LoraConfig, get_peft_model

        for p in self.encoder.parameters():
            p.requires_grad = False
        cfg = LoraConfig(
            r=r, lora_alpha=alpha, lora_dropout=dropout,
            target_modules=["q", "k", "v", "o"], bias="none",
        )
        self.encoder = get_peft_model(self.encoder, cfg)

    def _mean_pool(self, last_hidden: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        m = mask.unsqueeze(-1).float()                       # [B, L, 1]
        return (last_hidden * m).sum(1) / m.sum(1).clamp(min=1e-6)

    @torch.no_grad()
    def _encode_frozen(self, smiles: List[str], device: torch.device) -> torch.Tensor:
        missing = [s for s in smiles if s not in self._cache]
        if missing:
            uniq = list(dict.fromkeys(missing))
            for i in range(0, len(uniq), 256):
                chunk = uniq[i:i + 256]
                enc = self.tokenizer(
                    chunk, padding=True, truncation=True,
                    max_length=self.max_length, return_tensors="pt",
                ).to(device)
                out = self.encoder(**enc).last_hidden_state
                pooled = self._mean_pool(out, enc["attention_mask"])
                for s, v in zip(chunk, pooled):
                    self._cache[s] = v.cpu()
        return torch.stack([self._cache[s] for s in smiles]).to(device)

    def _encode_trainable(self, smiles: List[str], device: torch.device) -> torch.Tensor:
        enc = self.tokenizer(
            list(smiles), padding=True, truncation=True,
            max_length=self.max_length, return_tensors="pt",
        ).to(device)
        out = self.encoder(**enc).last_hidden_state
        return self._mean_pool(out, enc["attention_mask"])

    def forward(self, smiles: List[str], tab: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Args: smiles [B] SMILES 字符串列表。Returns: [B, d_model]。"""
        device = self.proj_down[0].weight.device
        feat = self._encode_frozen(smiles, device) if self._frozen \
            else self._encode_trainable(smiles, device)
        return self.proj_down(feat)

    def clear_cache(self) -> None:
        self._cache.clear()