"""
synth_pdb.plm
~~~~~~~~~~~~~
Protein Language Model (PLM) embeddings via ESM-2.

─────────────────────────────────────────────────────────────────────────────
WHAT IS A PROTEIN LANGUAGE MODEL?
─────────────────────────────────────────────────────────────────────────────

A protein language model (PLM) is a transformer neural network pre-trained on
hundreds of millions of protein sequences by masked-language-modelling (MLM):
randomly mask amino acids and train the network to predict the missing ones
from context.

After pre-training, the *internal activations* at the last hidden layer form
per-residue embedding vectors.  These representations encode:

  • Evolutionary information  — which positions co-vary across species
  • Structural context        — buried vs. solvent-exposed residues
  • Chemical environment      — polar / charged / hydrophobic neighbourhoods
  • Functional signals        — active-site residues vs. scaffold

All of this is learned from **sequence alone** — no 3D coordinates are used
during training.  Yet ESM-2 embeddings predict secondary structure, solvent
accessibility, and contact maps at near-state-of-the-art accuracy with zero
fine-tuning.  ESMFold, built on these representations, predicts full 3D
structure in milliseconds.

─────────────────────────────────────────────────────────────────────────────
MODEL: facebook/esm2_t6_8M_UR50D
─────────────────────────────────────────────────────────────────────────────

ESM-2 is a family of transformer models from Meta AI, trained on UniRef50
(~250M non-redundant protein sequences).  We use the smallest variant:

  Model          Params   Embed dim   File    Best for
  t6_8M          8M       320         ~30 MB  Education / fast experiments ← THIS
  t12_35M        35M      480         ~140 MB Better structure prediction
  t30_150M       150M     640         ~580 MB Near production quality
  t33_650M       650M     1280        ~2.5 GB AlphaFold-rivalling quality

All variants use the same API — change only the `model_name` argument.

Reference:
  Lin et al. (2023) "Evolutionary-scale prediction of atomic-level protein
  structure with a language model." Science 379, 1123–1130.
  https://doi.org/10.1126/science.ade2574

─────────────────────────────────────────────────────────────────────────────
ARCHITECTURE INSIDE ESM-2
─────────────────────────────────────────────────────────────────────────────

  Input sequence: "MQIFVKTLTG..."
       ↓
  Tokenization: each AA → integer token + [CLS] prefix, [EOS] suffix
       ↓
  Token embeddings: (L+2, 320) learnable lookup table
       ↓
  Rotary Position Embedding (RoPE): encodes token position without absolute PEs
       ↓
  6 × Transformer encoder layers:
      Multi-Head Self-Attention (10 heads)
      LayerNorm + residual connection
      FFN (1280-dim hidden) + LayerNorm + residual
       ↓
  Last hidden state: (L+2, 320)
  Slice off [CLS] and [EOS]: → (L, 320)
       ↓
  Return per-residue embedding matrix

─────────────────────────────────────────────────────────────────────────────
WHAT TO DO WITH THE EMBEDDINGS
─────────────────────────────────────────────────────────────────────────────

Per-residue embeddings (L, 320) — downstream uses:
  • GNN node features:       enrich the quality GNN with evolutionary context
  • Secondary structure:     simple linear probe on each residue
  • Contact prediction:      outer product → (L, L, 640) → CNN → contact map
  • Disorder prediction:     linear probe per residue
  • Site annotation:         active sites, binding sites, PTM sites

Mean embedding (320,) — downstream uses:
  • Sequence similarity search
  • Clustering protein families
  • Retrieval by function similarity

─────────────────────────────────────────────────────────────────────────────
USAGE
─────────────────────────────────────────────────────────────────────────────

    pip install synth-pdb[plm]

    from synth_pdb.plm import ESM2Embedder

    embedder = ESM2Embedder()                               # lazy — nothing loaded yet
    emb = embedder.embed("MQIFVKTLTG")                     # (10, 320) float32
    emb = embedder.embed_structure(atom_array)             # extract seq, then embed
    sim = embedder.sequence_similarity("ACDEF", "VWLYG")  # cosine sim of mean embeddings

─────────────────────────────────────────────────────────────────────────────
BENCHMARK (measured 2026-02-19, CPU, esm2_t6_8M_UR50D)
─────────────────────────────────────────────────────────────────────────────

  Protein length   First call (load + embed)   Subsequent calls
  50 residues      ~5 s (model download+load)   ~8 ms
  100 residues     ~5 s                         ~10 ms
  500 residues     ~5 s                         ~25 ms

  Model file: ~30 MB, cached after first download at ~/.cache/huggingface/
"""

import logging
from typing import Any, Optional, cast

import numpy as np

logger = logging.getLogger(__name__)

# Default ESM-2 model.  Can be overridden at ESM2Embedder() construction time.
# To use a larger model: "facebook/esm2_t12_35M_UR50D", "facebook/esm2_t33_650M_UR50D"
_DEFAULT_MODEL = "facebook/esm2_t6_8M_UR50D"

# Three-letter to one-letter code for sequence extraction from AtomArray
_THREE_TO_ONE = {
    "ALA": "A",
    "ARG": "R",
    "ASN": "N",
    "ASP": "D",
    "CYS": "C",
    "GLN": "Q",
    "GLU": "E",
    "GLY": "G",
    "HIS": "H",
    "ILE": "I",
    "LEU": "L",
    "LYS": "K",
    "MET": "M",
    "PHE": "F",
    "PRO": "P",
    "SER": "S",
    "THR": "T",
    "TRP": "W",
    "TYR": "Y",
    "VAL": "V",
    # Common aliases
    "HID": "H",
    "HIE": "H",
    "HIP": "H",
    "CYX": "C",
    "MSE": "M",
}


class ESM2Embedder:
    """
    Per-residue protein language model embeddings from ESM-2.

    The model is loaded **lazily** on the first call to embed() — not at
    __init__ time.  This means:
      • `import synth_pdb.plm` is always safe with no torch/transformers installed
      • `ESM2Embedder()` is instantaneous
      • The ~5-second model load occurs once, then is cached in self._model

    Args:
        model_name: HuggingFace model ID.  Default: "facebook/esm2_t6_8M_UR50D".
            Upgrade to "facebook/esm2_t12_35M_UR50D" for 480-dim embeddings
            with better accuracy; API is identical.
        device: Torch device string ("cpu", "cuda", "mps").  Default: auto-detect.
    """

    def __init__(
        self,
        model_name: str = _DEFAULT_MODEL,
        device: Optional[str] = None,
    ):
        self.model_name = model_name
        self._device_str = device  # resolved lazily
        self._model = None  # EsmModel — loaded on first embed()
        self._tokenizer = None  # EsmTokenizer — loaded on first embed()
        self._embedding_dim: Optional[int] = None

    # ──────────────────────────────────────────────────────────────────────
    # Public API
    # ──────────────────────────────────────────────────────────────────────

    @property
    def embedding_dim(self) -> int:
        """
        Embedding dimensionality for this model variant.

        Determined from the model config after the first embed() call.
        Before the first call, returns the known default for common models.
        """
        if self._embedding_dim is not None:
            return self._embedding_dim
        # Return known defaults without loading the model
        known = {
            "facebook/esm2_t6_8M_UR50D": 320,
            "facebook/esm2_t12_35M_UR50D": 480,
            "facebook/esm2_t30_150M_UR50D": 640,
            "facebook/esm2_t33_650M_UR50D": 1280,
            "facebook/esm2_t36_3B_UR50D": 2560,
        }
        if self.model_name in known:
            return known[self.model_name]
        # Unknown model — must load to inspect
        self._load_model()
        return self._embedding_dim  # type: ignore[return-value]

    def embed(self, sequence: str) -> np.ndarray:
        """
        Embed a protein sequence using ESM-2.

        Each amino acid is represented as a D-dimensional float32 vector
        encoding evolutionary, structural, and chemical context learned from
        250M protein sequences.

        Args:
            sequence: Single-letter amino acid string, e.g. "MQIFVKTLTG".
                      Standard 20 amino acids only.  Unknown residues → 'X'.

        Returns:
            np.ndarray of shape (L, D) and dtype float32.
            L = len(sequence), D = self.embedding_dim (320 for default model).

        Raises:
            ImportError: If torch or transformers are not installed.
                         Install with: pip install synth-pdb[plm]

        Example:
            >>> embedder = ESM2Embedder()
            >>> emb = embedder.embed("ACDEFGHIKLMNPQRSTVWY")
            >>> emb.shape
            (20, 320)
        """
        self._load_model()
        return self._run_model(sequence)

    def embed_structure(self, structure: Any) -> np.ndarray:
        """
        Embed a protein given its biotite AtomArray.

        Extracts the amino acid sequence from the structure (using residue
        names in the AtomArray), then delegates to embed().

        This is a convenience method — the embeddings are purely sequence-based
        and do not use any 3D coordinate information.

        Args:
            structure: biotite.structure.AtomArray.  Must contain at least
                       one atom per residue (e.g. CA atoms suffice).

        Returns:
            np.ndarray of shape (n_residues, embedding_dim), float32.

        Example:
            >>> from synth_pdb.generator import ProteinGenerator
            >>> structure = ProteinGenerator().generate(20, ss_type="helix")
            >>> emb = embedder.embed_structure(structure)
            >>> emb.shape
            (20, 320)
        """
        sequence = _extract_sequence(structure)
        logger.debug("Extracted sequence (%d residues) from AtomArray", len(sequence))
        return self.embed(sequence)

    def mean_embed(self, sequence: str) -> np.ndarray:
        """
        Return the mean-pooled sequence-level embedding.

        Mean pooling averages the per-residue vectors:
            mean_embed(seq) = (1/L) Σ embed(seq)[i]    for i in 0..L-1

        This gives a single D-dim vector representing the whole sequence.
        Loses positional information but enables fast sequence comparison.

        Returns:
            np.ndarray of shape (D,), float32.
        """
        return cast(np.ndarray, self.embed(sequence).mean(axis=0))

    def sequence_similarity(self, seq_a: str, seq_b: str) -> float:
        """
        Cosine similarity between the mean embeddings of two sequences.

        Returns a value in [-1, 1]:
          •  1.0  — identical embeddings (same sequence)
          •  0.0  — orthogonal (no similarity)
          • -1.0  — opposite (very unlikely for protein embeddings)

        WHY COSINE, NOT L2?
        -------------------
        Cosine similarity is magnitude-invariant — it measures the *angle*
        between vectors, not their length.  Longer proteins have higher-norm
        embeddings simply because there are more residues, not because they
        are more similar.  Cosine corrects for this.

        Args:
            seq_a: First single-letter amino acid string.
            seq_b: Second single-letter amino acid string.

        Returns:
            float — cosine similarity of mean embeddings.

        Example:
            >>> embedder.sequence_similarity("AAAAAAA", "VIVIVIV")
            0.832...   # high — both are simple repetitive peptides
            >>> embedder.sequence_similarity("ACDEFGHIK", "WQMPLRNTS")
            0.71...    # lower — very different character
        """
        ea = self.mean_embed(seq_a)
        eb = self.mean_embed(seq_b)
        norm_a = np.linalg.norm(ea)
        norm_b = np.linalg.norm(eb)
        if norm_a < 1e-8 or norm_b < 1e-8:
            return 0.0
        return float(np.dot(ea, eb) / (norm_a * norm_b))

    # ──────────────────────────────────────────────────────────────────────
    # Internal
    # ──────────────────────────────────────────────────────────────────────

    def _load_model(self) -> None:
        """
        Lazily import torch + transformers and load the ESM-2 model.

        Called automatically on the first embed() call.  Subsequent calls
        are no-ops because self._model is already set.

        Why lazy loading?
        -----------------
        torch and transformers together take ~1 s to import even when the
        model weights are already cached.  Deferring the import means that
        `from synth_pdb.plm import ESM2Embedder` is always fast and always safe,
        even in environments without PyTorch installed.

        The pattern is identical to how synth_nmr.neural_shifts handles the
        same problem for its NeuralShiftPredictor.
        """
        if self._model is not None:
            return  # Already loaded — fast path

        # ── Import check ──────────────────────────────────────────────────
        try:
            import torch  # noqa: F401
        except ImportError as exc:
            raise ImportError(
                "torch is required for ESM2Embedder. Install with: pip install synth-pdb[plm]"
            ) from exc

        try:
            from transformers import EsmModel, EsmTokenizer
        except (ImportError, TypeError) as exc:
            raise ImportError(
                "transformers is required for ESM2Embedder. "
                "Install with: pip install synth-pdb[plm]"
            ) from exc

        import torch

        # ── Device selection ──────────────────────────────────────────────
        # Priority: explicit override > MPS (Apple Silicon) > CUDA > CPU
        if self._device_str:
            device = torch.device(self._device_str)
        elif torch.backends.mps.is_available():
            device = torch.device("mps")  # Apple Silicon GPU
        elif torch.cuda.is_available():
            device = torch.device("cuda")
        else:
            device = torch.device("cpu")

        self._device = device
        logger.info("Loading ESM-2 model '%s' on %s …", self.model_name, device)

        # ── Load tokenizer + model ────────────────────────────────────────
        # HuggingFace caches the weights at ~/.cache/huggingface/ after the
        # first download (~30 MB for t6_8M).
        self._tokenizer = EsmTokenizer.from_pretrained(self.model_name)
        self._model = EsmModel.from_pretrained(self.model_name)  # type: ignore[assignment]

        assert self._tokenizer is not None
        assert self._model is not None

        self._model = self._model.to(device)
        self._model.eval()  # Disable dropout — we need deterministic inference

        # Record the embedding dim from the model config
        self._embedding_dim = self._model.config.hidden_size
        logger.info(
            "ESM-2 loaded: %d params, embedding_dim=%d",
            sum(p.numel() for p in self._model.parameters()),
            self._embedding_dim,
        )

    def _run_model(self, sequence: str) -> np.ndarray:
        """
        Tokenize the sequence, run the transformer, return per-residue embeddings.

        Tokenization adds [CLS] at position 0 and [EOS] at position L+1.
        We slice these off so the output aligns 1-to-1 with the input residues:
            last_hidden_state[:, 1:-1, :]  →  (L, D)

        Gradient computation is disabled (torch.no_grad) for speed and memory.
        """
        import torch

        assert self._tokenizer is not None, "Tokenizer not loaded"
        assert self._model is not None, "Model not loaded"

        # Tokenize: "ACDEF" → tensor of integer token IDs, shape (1, L+2)
        inputs = self._tokenizer(
            sequence,
            return_tensors="pt",
            add_special_tokens=True,  # adds [CLS] and [EOS]
        )
        inputs = {k: v.to(self._device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = self._model(**inputs)

        # last_hidden_state: (batch=1, L+2, D)
        # Slice off [CLS] (index 0) and [EOS] (index -1) → (1, L, D)
        hidden = outputs.last_hidden_state[:, 1:-1, :]  # remove special tokens
        embeddings = hidden.squeeze(0)  # (L, D)

        return cast(np.ndarray, embeddings.cpu().numpy().astype(np.float32))


# ──────────────────────────────────────────────────────────────────────────
# Utility: sequence extraction from AtomArray
# ──────────────────────────────────────────────────────────────────────────


def _extract_sequence(structure: Any) -> str:
    """
    Extract the one-letter amino acid sequence from a biotite AtomArray.

    Iterates residues in chain/res_id order, maps three-letter codes to
    one-letter codes using _THREE_TO_ONE.  Unknown residues (HETATM, ligands,
    modified AAs) are mapped to 'X', which ESM-2 treats as a masked token.

    Args:
        structure: biotite.structure.AtomArray.

    Returns:
        str — single-letter sequence, e.g. "MQIFVKTLTG".
    """
    import biotite.structure as struc

    res_starts = struc.get_residue_starts(structure)
    sequence = []
    for start in res_starts:
        res_name = structure.res_name[start].strip()
        one_letter = _THREE_TO_ONE.get(res_name, "X")
        if one_letter != "X":
            sequence.append(one_letter)
        else:
            logger.debug(
                "Unknown residue '%s' at position %d → 'X'", res_name, structure.res_id[start]
            )
            sequence.append("X")

    return "".join(sequence)
