"""LLM 分子特征分支：SMILES 文本 + chem'/tab soft token 的混合 prompt 编码。"""

import os
from typing import Dict, List, Optional

import torch
import torch.nn as nn


DEFAULT_MOLT5_PATH = os.environ.get("MOLT5_PATH", "models/molt5-base")

# 检索池支持的额外多任务（除 delivery 外）
_EXTRA_TASKS = ["size", "pdi", "ee", "toxic", "biodist"]


class LLMPromptEncoder(nn.Module):
    """用 LLM 编码分子，输出 F_llm [B, d_model]。

    三种模式：
      - use_soft_prompt=True：原始 SMILES 文本(+RAG 邻居) + chem'/tab soft token，
        inputs_embeds 注入，全程带梯度（配合 LoRA/QLoRA 真微调），不缓存特征。
      - use_rag=True 且非 soft：旧的冻结+缓存 RAG 路径。
      - 其他：冻结缓存 / 可训练 mean-pool。
    """

    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,
        use_qlora: bool = False,
        lora_r: int = 8,
        lora_alpha: int = 16,
        lora_dropout: float = 0.05,
        max_length: int = 256,
        use_rag: bool = False,
        rag_top_k: int = 4,
        use_soft_prompt: bool = False,
    ) -> None:
        super().__init__()
        from transformers import AutoTokenizer, T5EncoderModel, AutoModel

        self.use_lora = use_lora or use_qlora
        self.use_qlora = use_qlora
        self.max_length = max_length
        self.use_rag = use_rag
        self.rag_top_k = rag_top_k
        self.use_soft_prompt = use_soft_prompt

        _name_l = model_name_or_path.lower()
        _is_t5 = "t5" in _name_l
        _is_qwen = "qwen" in _name_l
        self._is_qwen = _is_qwen
        _is_biot5 = "biot5" in _name_l
        self._is_biot5 = _is_biot5

        self.tokenizer = AutoTokenizer.from_pretrained(
            model_name_or_path,
            trust_remote_code=_is_qwen,
            use_fast=not _is_biot5, 
        )
        if _is_qwen and self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        if _is_t5:
            self.encoder = T5EncoderModel.from_pretrained(model_name_or_path)
            self.hidden_size = self.encoder.config.d_model
        elif _is_qwen and use_qlora:
            from transformers import BitsAndBytesConfig
            bnb = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
                bnb_4bit_compute_dtype=torch.bfloat16,
            )
            self.encoder = AutoModel.from_pretrained(
                model_name_or_path, trust_remote_code=True,
                quantization_config=bnb, torch_dtype=torch.bfloat16)
            self.hidden_size = self.encoder.config.hidden_size
        elif _is_qwen:
            # 微调(use_lora)时用 4-bit 量化(QLoRA)省显存；纯冻结时用 fp16
            if use_lora:
                from transformers import BitsAndBytesConfig
                _bnb = BitsAndBytesConfig(
                    load_in_4bit=True, bnb_4bit_quant_type="nf4",
                    bnb_4bit_compute_dtype=torch.float16,
                    bnb_4bit_use_double_quant=True)
                self.encoder = AutoModel.from_pretrained(
                    model_name_or_path, trust_remote_code=True,
                    quantization_config=_bnb, device_map={"": 0})
            else:
                self.encoder = AutoModel.from_pretrained(
                    model_name_or_path, trust_remote_code=True, torch_dtype=torch.float16)
            self.hidden_size = self.encoder.config.hidden_size
        else:
            self.encoder = AutoModel.from_pretrained(model_name_or_path)
            self.hidden_size = self.encoder.config.hidden_size

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

        # soft token 投影：chem'(d_model) / tab(d_model) -> hidden_size
        if use_soft_prompt:
            self.chem_to_llm = nn.Linear(d_model, self.hidden_size)
            self.tab_to_llm = nn.Linear(d_model, self.hidden_size)
        else:
            self.chem_to_llm = None
            self.tab_to_llm = None

        self.proj_down = nn.Sequential(
            nn.Linear(self.hidden_size, d_model), nn.LayerNorm(d_model)
        )

        # 缓存：冻结句向量缓存 / 文本 prompt 缓存
        self._cache: Dict[str, torch.Tensor] = {}
        self._prompt_cache: Dict[str, str] = {}

        # RAG 检索池
        self._rag_pool_smiles: List[str] = []
        self._rag_pool_labels = None          # delivery [N]
        self._rag_pool_extra = None           # dict: task -> (values, valid)
        self._rag_pool_fps = None
        self._rag_pool_id: str = "none"

    def _apply_lora(self, r, alpha, dropout, prepare_kbit=False):
        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

        if prepare_kbit:
            self.encoder = prepare_model_for_kbit_training(self.encoder)
        else:
            for p in self.encoder.parameters():
                p.requires_grad = False

        # 4-bit 量化模型(QLoRA)需先 prepare，才能正确接收梯度
        if getattr(self.encoder, "is_loaded_in_4bit", False) or getattr(self.encoder, "is_loaded_in_8bit", False):
            from peft import prepare_model_for_kbit_training
            self.encoder = prepare_model_for_kbit_training(self.encoder)
        for p in self.encoder.parameters():
            p.requires_grad = False
        # 不同架构注意力层命名不同：
        # T5: q/k/v/o；Roberta/ChemBERTa: query/key/value；Qwen/Llama: q_proj/k_proj/v_proj/o_proj
        _enc_name = type(self.encoder).__name__.lower()
        if getattr(self, "_is_qwen", False) or "qwen" in _enc_name or "llama" in _enc_name:
            _targets = ["q_proj", "k_proj", "v_proj", "o_proj"]
        elif "t5" in _enc_name or hasattr(self.encoder.config, "d_model"):
            _targets = ["q", "k", "v", "o"]
        else:
            _targets = ["query", "key", "value"]
        cfg = LoraConfig(r=r, lora_alpha=alpha, lora_dropout=dropout,
                         target_modules=_targets, bias="none")
        self.encoder = get_peft_model(self.encoder, cfg)

    # ---------- 检索池 ----------
    def set_retrieval_pool(self, smiles_list, labels, pool_id="train", extra_labels=None):
        """设置 RAG 检索池（防泄漏：只传训练集分子与标签）。
        labels: delivery [N]; extra_labels: dict task -> (values, valid_bool)。"""
        import numpy as np
        from lnp_ml.modeling.retrieval import _smiles_to_fp
        self._rag_pool_smiles = list(smiles_list)
        self._rag_pool_labels = np.asarray(labels, dtype=np.float32).reshape(-1)
        self._rag_pool_extra = extra_labels
        self._rag_pool_fps = [_smiles_to_fp(s) for s in self._rag_pool_smiles]
        if pool_id != self._rag_pool_id:
            self._cache = {k: v for k, v in self._cache.items() if not k.startswith("RAG::")}
            self._prompt_cache.clear()
            self._rag_pool_id = pool_id

    def _retrieve_topk(self, query_smiles: str):
        """检索 Top-K 邻居，返回 [dict(smiles, sim, delivery, extra)]；排除查询分子自己。"""
        from rdkit import DataStructs
        from lnp_ml.modeling.retrieval import _smiles_to_fp
        qfp = _smiles_to_fp(query_smiles)
        if qfp is None or self._rag_pool_fps is None:
            return []
        sims = []
        for j, (smi, fp) in enumerate(zip(self._rag_pool_smiles, self._rag_pool_fps)):
            if fp is None or smi == query_smiles:
                continue
            sims.append((j, DataStructs.TanimotoSimilarity(qfp, fp)))
        sims.sort(key=lambda t: t[1], reverse=True)
        out = []
        for j, sim in sims[: self.rag_top_k]:
            extra = {}
            if self._rag_pool_extra is not None:
                for task in _EXTRA_TASKS:
                    if task in self._rag_pool_extra:
                        vals, valid = self._rag_pool_extra[task]
                        extra[task] = vals[j] if bool(valid[j]) else None
                    else:
                        extra[task] = None
            out.append({
                "smiles": self._rag_pool_smiles[j],
                "sim": sim,
                "delivery": float(self._rag_pool_labels[j]),
                "extra": extra,
            })
        return out

    @staticmethod
    def _fmt(v, kind="float"):
        if v is None:
            return "unknown"
        if kind == "int":
            return str(int(v))
        if kind == "vec":
            return "[" + ", ".join(f"{x:.3f}" for x in v) + "]"
        return f"{float(v):.3f}"

    def _fmt_mol(self, smiles: str) -> str:
        """按 backbone 期望格式化分子。
        BioT5：SMILES -> SELFIES，用 <bom>...<eom> 紧贴包裹（官方格式，token 间无空格）；
        其他 backbone：原样返回 SMILES。"""
        if not getattr(self, "_is_biot5", False):
            return smiles
        try:
            import selfies as sf
            sfs = sf.encoder(smiles)          # CCO -> [C][C][O]
        except Exception:
            return smiles                      # 转换失败退回 SMILES，避免整批中断
        return f"<bom>{sfs}<eom>"
        
    def _build_rag_prompt(self, target_smiles: str, neighbors) -> str:
        """构造 RAG prompt：原始 SMILES + 邻居多任务结果（numeric）。"""
        blocks = []
        for rank, nb in enumerate(neighbors, 1):
            ex = nb["extra"]
            blocks.append(
                f"Retrieved sample {rank}:\n"
                f"Molecule: {self._fmt_mol(nb['smiles'])}\n"
                f"Similarity score: {nb['sim']:.3f}\n"
                f"delivery_log: {self._fmt(nb['delivery'])}\n"
                f"size_z: {self._fmt(ex.get('size'))}\n"
                f"pdi_class: {self._fmt(ex.get('pdi'), 'int')}\n"
                f"ee_class: {self._fmt(ex.get('ee'), 'int')}\n"
                f"toxic: {self._fmt(ex.get('toxic'), 'int')}\n"
                f"biodist: {self._fmt(ex.get('biodist'), 'vec')}"
            )
        retrieved_block = "\n\n".join(blocks) if blocks else "(no retrieved samples)"
        return (
            "Task: Encode the target LNP molecule into a retrieval-aware representation "
            "for downstream multi-task property prediction. Do not output predictions.\n\n"
            f"[Target Molecule]\nMolecule: {self._fmt_mol(target_smiles)}\n\n"
            "[Retrieved Similar LNP Samples]\n"
            "Retrieved from the training set by fingerprint similarity, with their known "
            "multi-task outcomes (delivery_log, size_z, pdi_class, ee_class, toxic, biodist; "
            "'unknown' means the measurement is missing):\n\n"
            f"{retrieved_block}\n\n"
            "[Encoding Instructions]\n"
            "Capture the target structure and the retrieval evidence (structural similarity "
            "and consistency of retrieved outcomes) into your internal representation."
        )

    def _get_prompt(self, s: str) -> str:
        if not self.use_rag:
            return self._fmt_mol(s)
        key = f"{self._rag_pool_id}::{s}"
        if key not in self._prompt_cache:
            self._prompt_cache[key] = self._build_rag_prompt(s, self._retrieve_topk(s))
        return self._prompt_cache[key]
    def _rag_encode_batch(self, prompts, device):
        """编码一批 RAG prompt，取最后有效 token。grad 由调用方上下文决定。"""
        outs = []
        for i in range(0, len(prompts), 4):
            bt = prompts[i:i+4]
            enc = self.tokenizer(
                bt, padding=True, truncation=True,
                max_length=self.max_length, return_tensors="pt",
            ).to(device)
            out = self.encoder(**enc).last_hidden_state          # [B,L,H]
            lengths = enc["attention_mask"].sum(1) - 1
            b = torch.arange(out.size(0), device=device)
            outs.append(out[b, lengths.long(), :].float())       # [B,H]
        return torch.cat(outs, 0)

    def _encode_rag(self, smiles, device):
        """RAG 编码。冻结时：no_grad + 缓存（快）。微调时：算梯度 + 不缓存。"""
        if self._frozen:
            # 冻结路径：缓存复用，no_grad
            keys = [f"RAG::{self._rag_pool_id}::{s}" for s in smiles]
            to_compute = [s for s, k in zip(smiles, keys) if k not in self._cache]
            if to_compute:
                uniq = list(dict.fromkeys(to_compute))
                prompts = [self._build_rag_prompt(s, self._retrieve_topk(s)) for s in uniq]
                with torch.no_grad():
                    feat = self._rag_encode_batch(prompts, device)
                for s, v in zip(uniq, feat):
                    self._cache[f"RAG::{self._rag_pool_id}::{s}"] = v.float().cpu()
            return torch.stack([self._cache[k] for k in keys]).to(device)
        else:
            # 微调路径：每次重新编码，保留计算图（算梯度），不缓存
            prompts = [self._build_rag_prompt(s, self._retrieve_topk(s)) for s in smiles]
            return self._rag_encode_batch(prompts, device)

    # ---------- soft-prompt 编码（带梯度，不缓存特征）----------
    def _encode_softrag(self, smiles, chem, tab, device) -> torch.Tensor:
        prompts = [self._get_prompt(s) for s in smiles]
        enc = self.tokenizer(
            prompts, padding=True, truncation=True,
            max_length=self.max_length, return_tensors="pt",
        ).to(device)

        embed_layer = self.encoder.get_input_embeddings()
        text_embeds = embed_layer(enc["input_ids"])     # [B, L, H]
        text_mask = enc["attention_mask"]               # [B, L]

        soft_list = []
        if self.chem_to_llm is not None and chem is not None:
            soft_list.append(self.chem_to_llm(chem))    # [B, n_chem, H]
        if self.tab_to_llm is not None and tab is not None:
            soft_list.append(self.tab_to_llm(tab))      # [B, n_tab, H]

        if soft_list:
            soft = torch.cat(soft_list, dim=1).to(text_embeds.dtype)
            inputs_embeds = torch.cat([soft, text_embeds], dim=1)
            soft_mask = torch.ones(soft.size(0), soft.size(1),
                                   device=device, dtype=text_mask.dtype)
            attn_mask = torch.cat([soft_mask, text_mask], dim=1)
        else:
            inputs_embeds = text_embeds
            attn_mask = text_mask

        out = self.encoder(inputs_embeds=inputs_embeds, attention_mask=attn_mask).last_hidden_state
        lengths = attn_mask.sum(1) - 1
        b = torch.arange(out.size(0), device=device)
        feat = out[b, lengths.long(), :]                # [B, H] 最后有效 token
        return self.proj_down(feat.float())

    # ---------- 旧路径（保留，向后兼容）----------
    def _mean_pool(self, last_hidden, mask):
        m = mask.unsqueeze(-1).float()
        return (last_hidden * m).sum(1) / m.sum(1).clamp(min=1e-6)

    @torch.no_grad()
    def _encode_frozen(self, smiles, device):
        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]
                mols = [self._fmt_mol(s) for s in chunk]
                enc = self.tokenizer(mols, 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, device):
        mols = [self._fmt_mol(s) for s in smiles]
        enc = self.tokenizer(mols, 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], chem: Optional[torch.Tensor] = None,
                tab: Optional[torch.Tensor] = None) -> torch.Tensor:
        device = self.proj_down[0].weight.device
        if self.use_soft_prompt:
            return self._encode_softrag(smiles, chem, tab, device)
        if self.use_rag:
            return self._encode_softrag(smiles, None, None, 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()
        self._prompt_cache.clear()