"""UniMol 预计算表征的查表编码器"""
from pathlib import Path
from typing import Dict, List

import numpy as np
import torch
import torch.nn as nn


class UniMolEmbeddingEncoder(nn.Module):
    """从预计算 .npz 缓存按 SMILES 查 UniMol 表征，返回 {"unimol": [B, D]}。"""

    def __init__(self, cache_path: str) -> None:
        super().__init__()
        self.cache_path = str(cache_path)
        self._table: Dict[str, np.ndarray] = {}
        self._embed_dim: int = 0
        self._load()

    def _load(self) -> None:
        path = Path(self.cache_path)
        if not path.exists():
            raise FileNotFoundError(
                f"UniMol 缓存不存在: {path}。请先运行 scripts/precompute_unimol.py。"
            )
        data = np.load(path, allow_pickle=True)
        embeddings = np.asarray(data["embeddings"], dtype=np.float32)
        if embeddings.ndim != 2:
            raise ValueError(f"embeddings 应为 2D，实际 {embeddings.shape}")
        self._embed_dim = int(embeddings.shape[1])
        self._table = {str(s): embeddings[i] for i, s in enumerate(data["smiles"])}

    def forward(self, smiles_list: List[str]) -> Dict[str, torch.Tensor]:
        missing = [s for s in smiles_list if s not in self._table]
        if missing:
            raise KeyError(
                f"{len(missing)} 个 SMILES 不在 UniMol 缓存中，请重跑 precompute_unimol.py。"
                f"示例: {missing[:3]}"
            )
        mat = np.stack([self._table[s] for s in smiles_list])
        return {"unimol": torch.from_numpy(mat)}

    def clear_cache(self) -> None:
        """查表器无临时缓存，仅为接口对齐。"""

    @property
    def embed_dim(self) -> int:
        return self._embed_dim