fastplms/models/esmfold2/esmfold2_prepare_input.py

"""Translate typed sequence inputs into the tensors consumed by ESMFold2.

The conversion has four explicit stages: entity and chain assignment, residue
tokenization, structural feature construction, and atom-table padding.  Keeping
those stages separate makes the biological indexing rules testable without
loading model weights.
"""

from __future__ import annotations

import math
import warnings
from collections import defaultdict
from contextlib import suppress
from dataclasses import dataclass, field
from itertools import combinations
from typing import Any

import numpy as np
import torch

from .esmfold2_conformers import (
    get_ccd_leaving_atoms,
    get_idealized_atom_pos,
    get_ligand_ccd_atoms_with_charges,
    get_ligand_ccd_bonds,
    get_ligand_idealized_atom_pos,
)
from .esmfold2_constants import (
    CHARGED_ATOMS,
    DNA_1TO3,
    DNA_BACKBONE_ATOMS,
    DNA_HEAVY_ATOMS,
    DNA_RESIDUE_TO_RES_TYPE,
    DNA_RNA_LIGAND_INPUT_ID,
    DNA_UNK_RES_TYPE,
    ELEMENT_TO_ATOMIC_NUM,
    ESM_PROTEIN_VOCAB,
    MOL_TYPE_DNA,
    MOL_TYPE_NONPOLYMER,
    MOL_TYPE_PROTEIN,
    MOL_TYPE_RNA,
    MSA_GAP_TOKEN_ID,
    PROTEIN_1TO3,
    PROTEIN_3TO1,
    PROTEIN_HEAVY_ATOMS,
    PROTEIN_RESIDUE_TO_RES_TYPE,
    PROTEIN_UNK_RES_TYPE,
    RNA_1TO3,
    RNA_BACKBONE_ATOMS,
    RNA_HEAVY_ATOMS,
    RNA_RESIDUE_TO_RES_TYPE,
    RNA_UNK_RES_TYPE,
)
from .esmfold2_types import (
    MSA,
    DNAInput,
    LigandInput,
    Modification,
    ProteinInput,
    RNAInput,
    StructurePredictionInput,
)

_ZERO_POS = np.zeros(3, dtype=np.float32)
_ENCODE_ATOM_NAME_CACHE: dict[str, list[int]] = {}
_ELEMENT_ATOMIC_NUM_CACHE: dict[str, int] = {}
_TWO_LETTER_ELEMENTS = frozenset({"FE", "ZN", "MG", "MN", "CO", "NI", "CU", "SE", "BR"})


@dataclass
class AtomInfo:
    """One row in the unpadded atom table."""

    name: str
    element: str
    charge: int
    ref_pos: np.ndarray  # R has shape (3,).
    pos: np.ndarray  # X has shape (3,).
    token_index: int = -1
    atom_index: int = -1
    space_uid: int = -1
    is_valid: bool = True


@dataclass
class TokenInfo:
    """Biological and atom-span annotations for one model token."""

    token_index: int
    residue_index: int
    residue_name: str
    mol_type: int
    res_type: int
    input_id: int
    asym_id: int
    sym_id: int
    entity_id: int
    atom_start: int
    atom_count: int


@dataclass
class ChainInfo:
    """One input chain after entity and symmetry assignment."""

    chain_id: str
    asym_id: int
    entity_id: int
    sym_id: int
    mol_type: int
    tokens: list[TokenInfo] = field(default_factory=list)
    ligand_bonds: list[tuple[str, str]] = field(default_factory=list)


@dataclass
class _TokenizationState:
    """Mutable cursor shared by residue tokenizers."""

    token_index: int
    atom_index: int
    space_uid: int
    tokens: list[TokenInfo] = field(default_factory=list)
    atoms: list[AtomInfo] = field(default_factory=list)

    def _append_atom(
        self,
        name: str,
        element: str,
        charge: int,
        ref_pos: np.ndarray | None,
    ) -> None:
        self.atoms.append(
            AtomInfo(
                name=name,
                element=element,
                charge=charge,
                ref_pos=(ref_pos.copy() if ref_pos is not None else _ZERO_POS.copy()),
                pos=_ZERO_POS.copy(),
                token_index=self.token_index,
                atom_index=self.atom_index,
                space_uid=self.space_uid,
            )
        )
        self.atom_index += 1

    def _append_token(
        self,
        *,
        residue_index: int,
        residue_name: str,
        mol_type: int,
        res_type: int,
        input_id: int,
        asym_id: int,
        sym_id: int,
        entity_id: int,
        atom_start: int,
        atom_count: int,
    ) -> None:
        self.tokens.append(
            TokenInfo(
                token_index=self.token_index,
                residue_index=residue_index,
                residue_name=residue_name,
                mol_type=mol_type,
                res_type=res_type,
                input_id=input_id,
                asym_id=asym_id,
                sym_id=sym_id,
                entity_id=entity_id,
                atom_start=atom_start,
                atom_count=atom_count,
            )
        )
        self.token_index += 1

    def add_residue_token(
        self,
        atom_specs: list[tuple[str, str, int, np.ndarray | None]],
        **token_fields: Any,
    ) -> None:
        start = self.atom_index
        for atom_spec in atom_specs:
            self._append_atom(*atom_spec)
        self._append_token(
            atom_start=start,
            atom_count=len(atom_specs),
            **token_fields,
        )
        self.space_uid += 1

    def add_atom_tokens(
        self,
        atom_specs: list[tuple[str, str, int, np.ndarray | None]],
        **token_fields: Any,
    ) -> None:
        for atom_spec in atom_specs:
            start = self.atom_index
            self._append_atom(*atom_spec)
            self._append_token(atom_start=start, atom_count=1, **token_fields)
        self.space_uid += 1


def encode_atom_name(name: str) -> list[int]:
    """Encode a four-character atom name with the model's ASCII offset."""
    cached = _ENCODE_ATOM_NAME_CACHE.get(name)
    if cached is None:
        cached = [0 if char == " " else ord(char) - 32 for char in name.ljust(4)[:4]]
        _ENCODE_ATOM_NAME_CACHE[name] = cached
    return cached


def get_element_atomic_num(element: str) -> int:
    """Map an element symbol to the model's atomic-number vocabulary."""
    cached = _ELEMENT_ATOMIC_NUM_CACHE.get(element)
    if cached is None:
        cached = ELEMENT_TO_ATOMIC_NUM.get(element.upper(), 0)
        _ELEMENT_ATOMIC_NUM_CACHE[element] = cached
    return cached


def _infer_element(atom_name: str) -> str:
    normalized = atom_name.strip()
    if not normalized:
        return "C"
    if normalized[0].isdigit():
        return normalized[1] if len(normalized) > 1 else "H"
    if len(normalized) == 2 and normalized in _TWO_LETTER_ELEMENTS:
        return normalized
    return normalized[0]


def _compute_res_type(name: str, mol_type: int) -> int:
    if mol_type == MOL_TYPE_PROTEIN:
        return PROTEIN_RESIDUE_TO_RES_TYPE.get(name, PROTEIN_UNK_RES_TYPE)
    if mol_type == MOL_TYPE_DNA:
        return DNA_RESIDUE_TO_RES_TYPE.get(
            name, RNA_RESIDUE_TO_RES_TYPE.get(name, DNA_UNK_RES_TYPE)
        )
    if mol_type == MOL_TYPE_RNA:
        return RNA_RESIDUE_TO_RES_TYPE.get(
            name, DNA_RESIDUE_TO_RES_TYPE.get(name, RNA_UNK_RES_TYPE)
        )
    return PROTEIN_UNK_RES_TYPE


def _compute_esm_input_id(name: str, mol_type: int) -> int:
    if mol_type != MOL_TYPE_PROTEIN:
        return DNA_RNA_LIGAND_INPUT_ID
    letter = PROTEIN_3TO1.get(name)
    return (
        DNA_RNA_LIGAND_INPUT_ID
        if letter is None
        else ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"])
    )


def _apply_modifications(residues: list[str], modifications: list[Modification] | None) -> set[int]:
    changed: set[int] = set()
    for modification in modifications or ():
        residues[modification.position] = modification.ccd
        changed.add(modification.position)
    return changed


def _ideal_atom_specs(
    residue_name: str,
    residue_type: int,
    atom_names: list[str],
    *,
    charges: bool = True,
) -> list[tuple[str, str, int, np.ndarray | None]]:
    return [
        (
            atom_name,
            _infer_element(atom_name),
            CHARGED_ATOMS.get((residue_name, atom_name), 0) if charges else 0,
            get_idealized_atom_pos(residue_type, atom_name),
        )
        for atom_name in atom_names
    ]


def _ccd_atom_specs(
    residue_name: str,
    atoms: list[tuple[str, str, int]],
    excluded: set[str],
    *,
    force_zero: bool = False,
) -> list[tuple[str, str, int, np.ndarray | None]]:
    return [
        (
            atom_name,
            element,
            charge,
            None if force_zero else get_ligand_idealized_atom_pos(residue_name, atom_name),
        )
        for atom_name, element, charge in atoms
        if atom_name not in excluded
    ]


def tokenize_protein(
    sequence: str,
    modifications: list[Modification] | None,
    entity_id: int,
    asym_id: int,
    sym_id: int,
    token_offset: int,
    atom_offset: int,
    space_uid_offset: int,
) -> tuple[list[TokenInfo], list[AtomInfo]]:
    """Tokenize protein residues, atom-tokenizing modified CCD components."""
    residues = [PROTEIN_1TO3.get(letter, "UNK") for letter in sequence]
    modified = _apply_modifications(residues, modifications)
    state = _TokenizationState(token_offset, atom_offset, space_uid_offset)

    for residue_index, residue_name in enumerate(residues):
        canonical_name = "MET" if residue_name == "MSE" else residue_name
        common_fields = {
            "residue_index": residue_index,
            "mol_type": MOL_TYPE_PROTEIN,
            "asym_id": asym_id,
            "sym_id": sym_id,
            "entity_id": entity_id,
        }
        if residue_index not in modified and canonical_name in PROTEIN_HEAVY_ATOMS:
            residue_type = _compute_res_type(canonical_name, MOL_TYPE_PROTEIN)
            state.add_residue_token(
                _ideal_atom_specs(
                    canonical_name,
                    residue_type,
                    PROTEIN_HEAVY_ATOMS[canonical_name],
                ),
                residue_name=canonical_name,
                res_type=residue_type,
                input_id=_compute_esm_input_id(canonical_name, MOL_TYPE_PROTEIN),
                **common_fields,
            )
            continue

        ccd_atoms = get_ligand_ccd_atoms_with_charges(residue_name)
        if ccd_atoms is None:
            ccd_atoms = [
                (_infer_element(name), _infer_element(name), 0) for name in ("N", "CA", "C", "O")
            ]
        excluded = (
            set() if residue_index == len(residues) - 1 else get_ccd_leaving_atoms(residue_name)
        )
        retained = [atom for atom in ccd_atoms if atom[0] not in excluded]
        state.add_atom_tokens(
            _ccd_atom_specs(
                residue_name,
                retained,
                set(),
                force_zero=len(retained) == 1,
            ),
            residue_name=residue_name,
            res_type=PROTEIN_UNK_RES_TYPE,
            input_id=DNA_RNA_LIGAND_INPUT_ID,
            **common_fields,
        )
    return state.tokens, state.atoms


def tokenize_nucleotide(
    sequence: str,
    modifications: list[Modification] | None,
    mol_type: int,
    entity_id: int,
    asym_id: int,
    sym_id: int,
    token_offset: int,
    atom_offset: int,
    space_uid_offset: int,
) -> tuple[list[TokenInfo], list[AtomInfo]]:
    """Tokenize DNA or RNA, retaining backbone atoms for unknown bases."""
    dna = mol_type == MOL_TYPE_DNA
    letter_map = DNA_1TO3 if dna else RNA_1TO3
    heavy_atoms = DNA_HEAVY_ATOMS if dna else RNA_HEAVY_ATOMS
    backbone_atoms = DNA_BACKBONE_ATOMS if dna else RNA_BACKBONE_ATOMS
    unknown_type = DNA_UNK_RES_TYPE if dna else RNA_UNK_RES_TYPE
    residues = [letter_map.get(letter, "UNK") for letter in sequence]
    modified = _apply_modifications(residues, modifications)
    state = _TokenizationState(token_offset, atom_offset, space_uid_offset)

    for residue_index, residue_name in enumerate(residues):
        common_fields = {
            "residue_index": residue_index,
            "residue_name": residue_name,
            "mol_type": mol_type,
            "asym_id": asym_id,
            "sym_id": sym_id,
            "entity_id": entity_id,
            "input_id": DNA_RNA_LIGAND_INPUT_ID,
        }
        if residue_index not in modified and residue_name in heavy_atoms:
            residue_type = _compute_res_type(residue_name, mol_type)
            state.add_residue_token(
                _ideal_atom_specs(residue_name, residue_type, heavy_atoms[residue_name]),
                res_type=residue_type,
                **common_fields,
            )
            continue
        if residue_index not in modified and residue_name == "UNK":
            state.add_residue_token(
                [(atom_name, _infer_element(atom_name), 0, None) for atom_name in backbone_atoms],
                res_type=unknown_type,
                **common_fields,
            )
            continue

        ccd_atoms = get_ligand_ccd_atoms_with_charges(residue_name)
        if ccd_atoms is None:
            ccd_atoms = [(_infer_element(name), _infer_element(name), 0) for name in backbone_atoms]
        excluded = (
            set() if residue_index == len(residues) - 1 else get_ccd_leaving_atoms(residue_name)
        )
        state.add_atom_tokens(
            _ccd_atom_specs(residue_name, ccd_atoms, excluded),
            res_type=PROTEIN_UNK_RES_TYPE,
            **common_fields,
        )
    return state.tokens, state.atoms


def tokenize_ligand_ccd(
    ccd_codes: list[str],
    entity_id: int,
    asym_id: int,
    sym_id: int,
    token_offset: int,
    atom_offset: int,
    space_uid_offset: int,
    has_covalent_bond: bool,
) -> tuple[list[TokenInfo], list[AtomInfo]]:
    """Tokenize CCD ligands with one model token per retained atom."""
    state = _TokenizationState(token_offset, atom_offset, space_uid_offset)
    for residue_index, code in enumerate(ccd_codes):
        ccd_atoms = get_ligand_ccd_atoms_with_charges(code)
        if ccd_atoms is None:
            raise ValueError(f"CCD component {code} not found")
        excluded = get_ccd_leaving_atoms(code) if has_covalent_bond else set()
        state.add_atom_tokens(
            _ccd_atom_specs(code, ccd_atoms, excluded),
            residue_index=residue_index,
            residue_name=code,
            mol_type=MOL_TYPE_NONPOLYMER,
            res_type=PROTEIN_UNK_RES_TYPE,
            input_id=DNA_RNA_LIGAND_INPUT_ID,
            asym_id=asym_id,
            sym_id=sym_id,
            entity_id=entity_id,
        )
    return state.tokens, state.atoms


def tokenize_ligand_smiles(
    smiles: str,
    entity_id: int,
    asym_id: int,
    sym_id: int,
    token_offset: int,
    atom_offset: int,
    space_uid_offset: int,
    seed: int | None = None,
) -> tuple[list[TokenInfo], list[AtomInfo], list[tuple[str, str]]]:
    """Generate a conformer and tokenize each heavy atom of a SMILES ligand."""
    from rdkit import Chem
    from rdkit.Chem import AllChem

    molecule = Chem.MolFromSmiles(smiles)
    if molecule is None:
        raise ValueError(f"Failed to parse SMILES: {smiles}")
    molecule = Chem.AddHs(molecule)
    canonical_order = AllChem.CanonicalRankAtoms(molecule)  # type: ignore[attr-defined]
    for atom, canonical_index in zip(molecule.GetAtoms(), canonical_order, strict=True):
        name = atom.GetSymbol().upper() + str(canonical_index + 1)
        if len(name) > 4:
            raise ValueError(f"SMILES {smiles} has atom name longer than 4 chars: {name}")
        atom.SetProp("name", name)

    options = AllChem.ETKDGv3()  # type: ignore[attr-defined]
    options.clearConfs = False
    if seed is not None:
        options.randomSeed = seed
    conformer_id = AllChem.EmbedMolecule(molecule, options)  # type: ignore[attr-defined]
    if conformer_id == -1:
        options.useRandomCoords = True
        conformer_id = AllChem.EmbedMolecule(molecule, options)  # type: ignore[attr-defined]
    if conformer_id != -1:
        with suppress(RuntimeError, ValueError):
            AllChem.UFFOptimizeMolecule(  # type: ignore[attr-defined]
                molecule, confId=conformer_id, maxIters=1000
            )

    molecule = Chem.RemoveHs(molecule)
    if molecule.GetNumConformers() == 0:
        raise ValueError(f"Failed to generate conformer for SMILES: {smiles}")
    conformer = molecule.GetConformer(0)
    atom_specs: list[tuple[str, str, int, np.ndarray | None]] = []
    for atom in molecule.GetAtoms():
        position = conformer.GetAtomPosition(atom.GetIdx())
        atom_specs.append(
            (
                atom.GetProp("name"),
                atom.GetSymbol(),
                atom.GetFormalCharge(),
                np.asarray([position.x, position.y, position.z], dtype=np.float32),
            )
        )
    state = _TokenizationState(token_offset, atom_offset, space_uid_offset)
    state.add_atom_tokens(
        atom_specs,
        residue_index=0,
        residue_name="LIG",
        mol_type=MOL_TYPE_NONPOLYMER,
        res_type=PROTEIN_UNK_RES_TYPE,
        input_id=DNA_RNA_LIGAND_INPUT_ID,
        asym_id=asym_id,
        sym_id=sym_id,
        entity_id=entity_id,
    )
    bonds = [
        (
            bond.GetBeginAtom().GetProp("name"),
            bond.GetEndAtom().GetProp("name"),
        )
        for bond in molecule.GetBonds()
    ]
    return state.tokens, state.atoms, bonds


def _get_sequence_key(item: Any) -> str:
    if isinstance(item, ProteinInput):
        return f"PROTEIN:{item.sequence}"
    if isinstance(item, DNAInput):
        return f"DNA:{item.sequence}"
    if isinstance(item, RNAInput):
        return f"RNA:{item.sequence}"
    if isinstance(item, LigandInput):
        return f"LIGAND_CCD:{','.join(item.ccd)}" if item.ccd else f"LIGAND_SMILES:{item.smiles}"
    raise ValueError(f"Unknown input type: {type(item)}")


def _tokenize_chain(
    item: Any,
    chain_id: str,
    *,
    entity_id: int,
    asym_id: int,
    sym_id: int,
    token_offset: int,
    atom_offset: int,
    space_uid_offset: int,
    covalent_chains: set[str],
    seed: int | None,
) -> tuple[list[TokenInfo], list[AtomInfo], list[tuple[str, str]]]:
    common = {
        "entity_id": entity_id,
        "asym_id": asym_id,
        "sym_id": sym_id,
        "token_offset": token_offset,
        "atom_offset": atom_offset,
        "space_uid_offset": space_uid_offset,
    }
    if isinstance(item, ProteinInput):
        if item.msa is None:
            warnings.warn(
                f"No MSA provided for {item.id}, using single sequence mode",
                stacklevel=2,
            )
        tokens, atoms = tokenize_protein(item.sequence, item.modifications, **common)
        return tokens, atoms, []
    if isinstance(item, (DNAInput, RNAInput)):
        mol_type = MOL_TYPE_DNA if isinstance(item, DNAInput) else MOL_TYPE_RNA
        tokens, atoms = tokenize_nucleotide(
            item.sequence, item.modifications, mol_type=mol_type, **common
        )
        return tokens, atoms, []
    if not isinstance(item, LigandInput):
        raise ValueError(f"Unknown input type: {type(item)}")
    if item.ccd is not None:
        if item.smiles is not None:
            warnings.warn("Both ccd and smiles provided, using ccd", stacklevel=2)
        tokens, atoms = tokenize_ligand_ccd(
            item.ccd,
            has_covalent_bond=chain_id in covalent_chains,
            **common,
        )
        return tokens, atoms, []
    if item.smiles is not None:
        return tokenize_ligand_smiles(item.smiles, seed=seed, **common)
    raise ValueError("LigandInput must have either ccd or smiles")


def build_chains_from_input(
    input: StructurePredictionInput, seed: int | None = None
) -> tuple[list[ChainInfo], list[TokenInfo], list[AtomInfo]]:
    """Assign entities and symmetry copies, then tokenize every input chain."""
    chains: list[ChainInfo] = []
    tokens: list[TokenInfo] = []
    atoms: list[AtomInfo] = []
    entity_for_sequence: dict[str, int] = {}
    next_symmetry: dict[int, int] = {}
    covalent_chains = {
        chain_id
        for bond in input.covalent_bonds or ()
        for chain_id in (bond.chain_id1, bond.chain_id2)
    }
    space_uid_offset = 0

    for item in input.sequences:
        key = _get_sequence_key(item)
        entity_id = entity_for_sequence.setdefault(key, len(entity_for_sequence))
        chain_ids = [item.id] if isinstance(item.id, str) else item.id
        for chain_id in chain_ids:
            sym_id = next_symmetry.get(entity_id, 0)
            next_symmetry[entity_id] = sym_id + 1
            asym_id = len(chains)
            new_tokens, new_atoms, ligand_bonds = _tokenize_chain(
                item,
                chain_id,
                entity_id=entity_id,
                asym_id=asym_id,
                sym_id=sym_id,
                token_offset=len(tokens),
                atom_offset=len(atoms),
                space_uid_offset=space_uid_offset,
                covalent_chains=covalent_chains,
                seed=seed,
            )
            chains.append(
                ChainInfo(
                    chain_id=chain_id,
                    asym_id=asym_id,
                    entity_id=entity_id,
                    sym_id=sym_id,
                    mol_type=(new_tokens[0].mol_type if new_tokens else MOL_TYPE_PROTEIN),
                    tokens=new_tokens,
                    ligand_bonds=ligand_bonds,
                )
            )
            tokens.extend(new_tokens)
            atoms.extend(new_atoms)
            space_uid_offset += len({atom.space_uid for atom in new_atoms})
    return chains, tokens, atoms


def _atom_indices_by_name(atoms: list[AtomInfo]) -> dict[int, dict[str, int]]:
    result: dict[int, dict[str, int]] = defaultdict(dict)
    for atom in atoms:
        if atom.is_valid:
            result[atom.token_index][atom.name] = atom.atom_index
    return result


def _ligand_frames(
    tokens: list[TokenInfo],
    atoms: list[AtomInfo],
    atom_indices: dict[int, dict[str, int]],
) -> dict[int, tuple[int, int, int]]:
    atom_for_token: dict[int, int] = {}
    tokens_by_residue: dict[tuple[int, int], list[int]] = defaultdict(list)
    for token in tokens:
        if token.mol_type != MOL_TYPE_NONPOLYMER:
            continue
        named_atoms = atom_indices.get(token.token_index)
        if named_atoms:
            atom_for_token[token.token_index] = next(iter(named_atoms.values()))
        tokens_by_residue[(token.asym_id, token.residue_index)].append(token.token_index)

    frames: dict[int, tuple[int, int, int]] = {}
    for residue_tokens in tokens_by_residue.values():
        residue_atoms = [
            atom_for_token[token] for token in residue_tokens if token in atom_for_token
        ]
        if len(residue_atoms) < 3:
            for token in residue_tokens:
                if token in atom_for_token:
                    atom_index = atom_for_token[token]
                    frames[token] = (atom_index, atom_index, atom_index)
            continue
        R = np.asarray([atoms[index].ref_pos for index in residue_atoms])
        distances = np.sqrt(((R[:, None] - R[None]) ** 2).sum(-1))
        nearest = np.argsort(distances, axis=1)
        local = np.column_stack((nearest[:, 1], nearest[:, 0], nearest[:, 2]))
        local_index = {atom_index: index for index, atom_index in enumerate(residue_atoms)}
        for token in residue_tokens:
            atom_index = atom_for_token.get(token)
            if atom_index is None:
                continue
            selected = local[local_index[atom_index]]
            frames[token] = tuple(residue_atoms[int(index)] for index in selected)
    return frames


def _frame_for_token(
    token: TokenInfo,
    named_atoms: dict[str, int],
    ligand_frames: dict[int, tuple[int, int, int]],
) -> tuple[int, int, int]:
    fallback = next(iter(named_atoms.values()), 0)
    if token.mol_type == MOL_TYPE_PROTEIN:
        return (
            (fallback, fallback, fallback)
            if token.res_type == PROTEIN_UNK_RES_TYPE
            else (
                named_atoms.get("N", 0),
                named_atoms.get("CA", 0),
                named_atoms.get("C", 0),
            )
        )
    if token.mol_type in (MOL_TYPE_DNA, MOL_TYPE_RNA):
        return (
            (fallback, fallback, fallback)
            if token.res_type == PROTEIN_UNK_RES_TYPE
            else (
                named_atoms.get("C1'", 0),
                named_atoms.get("C3'", 0),
                named_atoms.get("C4'", 0),
            )
        )
    if token.mol_type == MOL_TYPE_NONPOLYMER:
        return ligand_frames.get(token.token_index, (fallback, fallback, fallback))
    return fallback, fallback, fallback


def _resolved_frames(
    frames: np.ndarray, tokens: list[TokenInfo], atoms: list[AtomInfo]
) -> np.ndarray:
    if not tokens:
        return np.zeros(0, dtype=bool)
    X = (
        np.asarray([atom.pos for atom in atoms], dtype=np.float32)
        if atoms
        else np.zeros((0, 3), dtype=np.float32)
    )
    valid_atoms = (
        np.asarray([atom.is_valid for atom in atoms], dtype=bool)
        if atoms
        else np.zeros(0, dtype=bool)
    )
    resolved_atoms = valid_atoms & np.any(X != 0, axis=1)
    origin = X[frames[:, 1]]
    left = X[frames[:, 0]] - origin
    right = X[frames[:, 2]] - origin
    left_norm = np.linalg.norm(left, axis=1)
    right_norm = np.linalg.norm(right, axis=1)
    valid_norms = (left_norm >= 1e-6) & (right_norm >= 1e-6)
    cosine = np.zeros(len(tokens), dtype=np.float32)
    if np.any(valid_norms):
        cosine[valid_norms] = np.sum(left[valid_norms] * right[valid_norms], axis=1) / (
            left_norm[valid_norms] * right_norm[valid_norms]
        )
    angle = np.degrees(np.arccos(np.abs(np.clip(cosine, -1, 1))))
    all_resolved = resolved_atoms[frames].all(axis=1)
    repeated = (frames[:, 0] == frames[:, 1]) & (frames[:, 1] == frames[:, 2])
    return all_resolved & ~repeated & valid_norms & (angle >= 25)


def compute_frame_indices(
    tokens: list[TokenInfo], atoms: list[AtomInfo]
) -> tuple[np.ndarray, np.ndarray]:
    """Return frame atom indices F with shape (l, 3) and validity M with shape (l,)."""
    named_atoms = _atom_indices_by_name(atoms)
    ligand_frames = _ligand_frames(tokens, atoms, named_atoms)
    frames = np.asarray(
        [
            _frame_for_token(token, named_atoms.get(token.token_index, {}), ligand_frames)
            for token in tokens
        ],
        dtype=np.int64,
    )
    return frames, _resolved_frames(frames, tokens, atoms)


def _atom_tokenized_residues(
    tokens: list[TokenInfo], atoms: list[AtomInfo]
) -> dict[tuple[int, int], list[tuple[str, int]]]:
    grouped: dict[tuple[int, int], list[tuple[str, int]]] = defaultdict(list)
    for atom in atoms:
        if not atom.is_valid or atom.token_index >= len(tokens):
            continue
        token = tokens[atom.token_index]
        if token.mol_type == MOL_TYPE_NONPOLYMER or token.res_type == PROTEIN_UNK_RES_TYPE:
            grouped[(token.asym_id, token.residue_index)].append((atom.name, atom.token_index))
    return grouped


def _backbone_token(
    residue_tokens: list[TokenInfo], atom_name: str, atoms: list[AtomInfo]
) -> int | None:
    if len(residue_tokens) == 1 and residue_tokens[0].res_type != PROTEIN_UNK_RES_TYPE:
        return residue_tokens[0].token_index
    for token in residue_tokens:
        for atom_index in range(token.atom_start, token.atom_start + token.atom_count):
            if atom_index < len(atoms) and atoms[atom_index].name == atom_name:
                return token.token_index
    return residue_tokens[0].token_index if residue_tokens else None


def compute_token_bonds(
    tokens: list[TokenInfo],
    atoms: list[AtomInfo],
    input: StructurePredictionInput,
    chains: list[ChainInfo],
) -> torch.Tensor:
    """Build the symmetric token-bond matrix M with shape (l, l, 1)."""
    edges: set[tuple[int, int]] = set()

    def connect(left: int | None, right: int | None) -> None:
        if left is not None and right is not None and left != right:
            edges.add((min(left, right), max(left, right)))

    explicit_bonds = {
        (chain.asym_id, 0): chain.ligand_bonds for chain in chains if chain.ligand_bonds
    }
    for residue_key, atom_list in _atom_tokenized_residues(tokens, atoms).items():
        if not atom_list:
            continue
        residue_name = tokens[atom_list[0][1]].residue_name
        token_for_name = {name: token_index for name, token_index in atom_list}
        bonds = explicit_bonds.get(residue_key)
        if bonds is None:
            bonds = get_ligand_ccd_bonds(residue_name)
        if bonds:
            for left_name, right_name in bonds:
                if left_name in token_for_name and right_name in token_for_name:
                    connect(token_for_name[left_name], token_for_name[right_name])
        else:
            for left, right in combinations([token_index for _, token_index in atom_list], 2):
                connect(left, right)

    if input.covalent_bonds:
        chain_for_id = {chain.chain_id: chain for chain in chains}
        residue_atoms: dict[tuple[int, int], list[AtomInfo]] = defaultdict(list)
        for atom in atoms:
            if atom.is_valid and atom.token_index < len(tokens):
                token = tokens[atom.token_index]
                residue_atoms[(token.asym_id, token.residue_index)].append(atom)
        for bond in input.covalent_bonds:
            left_chain = chain_for_id.get(bond.chain_id1)
            right_chain = chain_for_id.get(bond.chain_id2)
            if left_chain is None or right_chain is None:
                continue
            left_atoms = residue_atoms.get((left_chain.asym_id, bond.res_idx1), [])
            right_atoms = residue_atoms.get((right_chain.asym_id, bond.res_idx2), [])
            if bond.atom_idx1 < len(left_atoms) and bond.atom_idx2 < len(right_atoms):
                connect(
                    left_atoms[bond.atom_idx1].token_index,
                    right_atoms[bond.atom_idx2].token_index,
                )

    protein_residues: dict[tuple[int, int], list[TokenInfo]] = defaultdict(list)
    for token in tokens:
        if token.mol_type == MOL_TYPE_PROTEIN:
            protein_residues[(token.asym_id, token.residue_index)].append(token)
    for (asym_id, residue_index), residue_tokens in protein_residues.items():
        if not any(token.res_type == PROTEIN_UNK_RES_TYPE for token in residue_tokens):
            continue
        previous = protein_residues.get((asym_id, residue_index - 1))
        following = protein_residues.get((asym_id, residue_index + 1))
        if previous:
            connect(
                _backbone_token(previous, "C", atoms),
                _backbone_token(residue_tokens, "N", atoms),
            )
        if following:
            connect(
                _backbone_token(residue_tokens, "C", atoms),
                _backbone_token(following, "N", atoms),
            )

    matrix = torch.zeros(len(tokens), len(tokens), 1, dtype=torch.float32)
    for left, right in edges:
        matrix[left, right, 0] = 1.0
        matrix[right, left, 0] = 1.0
    return matrix


def compute_representative_atoms(tokens: list[TokenInfo], atoms: list[AtomInfo]) -> torch.Tensor:
    """Choose one distogram atom per token and return indices I with shape (l,)."""
    named_atoms = _atom_indices_by_name(atoms)
    representatives = torch.zeros(len(tokens), dtype=torch.int64)
    for token in tokens:
        names = named_atoms.get(token.token_index, {})
        fallback = next(iter(names.values()), 0)
        if token.mol_type == MOL_TYPE_PROTEIN:
            representative = names.get("CB", names.get("CA", fallback))
        elif token.mol_type in (MOL_TYPE_DNA, MOL_TYPE_RNA):
            if token.res_type in (27, 32):
                representative = names.get("C1'", fallback)
            elif token.res_type in (23, 24, 28, 29):
                representative = names.get("C4", names.get("C1'", fallback))
            else:
                representative = names.get("C2", names.get("C1'", fallback))
        else:
            representative = fallback
        representatives[token.token_index] = representative
    return representatives


def _msa_assignments(
    input: StructurePredictionInput, chains: list[ChainInfo]
) -> dict[int, MSA | None]:
    chain_msas: dict[int, MSA | None] = {}
    chain_index = 0
    for item in input.sequences:
        chain_ids = [item.id] if isinstance(item.id, str) else list(item.id)
        for _ in chain_ids:
            chain = chains[chain_index]
            if isinstance(item, ProteinInput):
                chain_msas[chain.asym_id] = (
                    MSA.from_sequences([item.sequence]) if item.msa is None else item.msa
                )
            else:
                chain_msas[chain.asym_id] = None
            chain_index += 1
    return chain_msas


def compute_msa_features(
    input: StructurePredictionInput,
    chains: list[ChainInfo],
    tokens: list[TokenInfo],
    max_seqs: int = 16384,
) -> dict[str, torch.Tensor]:
    """Pair per-chain MSAs and return row features with shape (m, l)."""
    from .esmfold2_paired_msa import (
        construct_paired_msa,
        protein_letter_to_res_type,
    )

    chain_msas = _msa_assignments(input, chains)
    query_types = {
        chain.asym_id: np.asarray(
            [token.res_type for token in tokens if token.asym_id == chain.asym_id],
            dtype=np.int64,
        )
        for chain in chains
    }
    msa_residues, deletion_counts, _ = construct_paired_msa(
        chain_msas,
        query_types,
        np.asarray([token.asym_id for token in tokens], dtype=np.int64),
        np.asarray([token.residue_index for token in tokens], dtype=np.int64),
        letter_to_res_type=protein_letter_to_res_type(),
        max_seqs=max_seqs,
    )
    for token in tokens:
        if chain_msas.get(token.asym_id) is None:
            msa_residues[:, token.token_index] = MSA_GAP_TOKEN_ID
            msa_residues[0, token.token_index] = token.res_type
    if msa_residues.shape[0] == 0:
        msa_residues = np.full((1, len(tokens)), MSA_GAP_TOKEN_ID, dtype=np.int64)
        deletion_counts = np.zeros((1, len(tokens)), dtype=np.float32)

    msa = torch.from_numpy(msa_residues)
    deletion_count = torch.from_numpy(deletion_counts)
    deletion_value = (np.pi / 2) * torch.arctan(deletion_count / 3)
    return {
        "msa": msa,
        "deletion_value": deletion_value,
        "has_deletion": deletion_count > 0,
        "deletion_mean": deletion_value.mean(dim=0),
        "msa_attention_mask": torch.ones_like(msa, dtype=torch.bool),
    }


def compute_distogram_conditioning(
    input: StructurePredictionInput,
    chains: list[ChainInfo],
    tokens: list[TokenInfo],
    disto_center: torch.Tensor,
    min_dist: float = 2.0,
    max_dist: float = 22.0,
    num_bins: int = 64,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Bin user distances into D and return D plus its Boolean mask M."""
    del disto_center
    n_tokens = len(tokens)
    bins = torch.zeros((n_tokens, n_tokens), dtype=torch.long)
    mask = torch.zeros((n_tokens, n_tokens), dtype=torch.bool)
    if not input.distogram_conditioning:
        return bins, mask
    asym_for_chain = {chain.chain_id: chain.asym_id for chain in chains}
    tokens_for_asym: dict[int, list[int]] = defaultdict(list)
    for token in tokens:
        tokens_for_asym[token.asym_id].append(token.token_index)
    boundaries = torch.linspace(min_dist, max_dist, num_bins + 1)

    for conditioning in input.distogram_conditioning:
        asym_id = asym_for_chain.get(conditioning.chain_id)
        if asym_id is None:
            continue
        indices = tokens_for_asym[asym_id]
        distances = torch.as_tensor(conditioning.distogram, dtype=torch.float32)
        expected_shape = (len(indices), len(indices))
        if distances.shape != expected_shape:
            raise ValueError(
                f"Distogram shape {distances.shape} doesn't match chain length {len(indices)}"
            )
        selected = torch.bucketize(distances, boundaries[:-1]).sub(1).clamp(0, num_bins - 1)
        token_indices_tensor = torch.as_tensor(indices, dtype=torch.long)
        bins[token_indices_tensor[:, None], token_indices_tensor[None, :]] = selected
        mask[token_indices_tensor[:, None], token_indices_tensor[None, :]] = True
    return bins, mask


def _padded_atoms(atoms: list[AtomInfo]) -> list[AtomInfo]:
    target = math.ceil(len(atoms) / 32) * 32 if atoms else 32
    padding = [
        AtomInfo(
            name="",
            element="",
            charge=0,
            ref_pos=_ZERO_POS.copy(),
            pos=_ZERO_POS.copy(),
            token_index=0,
            atom_index=index,
            space_uid=0,
            is_valid=False,
        )
        for index in range(len(atoms), target)
    ]
    return [*atoms, *padding]


def _token_tensors(tokens: list[TokenInfo]) -> dict[str, torch.Tensor]:
    fields = {
        "token_index": "token_index",
        "residue_index": "residue_index",
        "asym_id": "asym_id",
        "sym_id": "sym_id",
        "entity_id": "entity_id",
        "mol_type": "mol_type",
        "res_type": "res_type",
        "input_ids": "input_id",
    }
    return {
        output_name: torch.from_numpy(
            np.asarray([getattr(token, attribute) for token in tokens], dtype=np.int64)
        )
        for output_name, attribute in fields.items()
    }


def _atom_tensors(atoms: list[AtomInfo]) -> dict[str, torch.Tensor]:
    n_atoms = len(atoms)
    ref_pos = np.zeros((n_atoms, 3), dtype=np.float32)
    ref_element = np.zeros(n_atoms, dtype=np.int64)
    ref_charge = np.zeros(n_atoms, dtype=np.int8)
    ref_name = np.zeros((n_atoms, 4), dtype=np.int64)
    ref_space = np.zeros(n_atoms, dtype=np.int64)
    atom_mask = np.zeros(n_atoms, dtype=np.bool_)
    atom_to_token = np.zeros(n_atoms, dtype=np.int64)
    positions = np.zeros((n_atoms, 3), dtype=np.float64)
    valid = np.zeros(n_atoms, dtype=np.bool_)
    for index, atom in enumerate(atoms):
        if atom.ref_pos is not None:
            ref_pos[index] = atom.ref_pos
        ref_charge[index] = atom.charge
        ref_space[index] = atom.space_uid if atom.space_uid >= 0 else atom.token_index
        atom_mask[index] = atom.is_valid
        valid[index] = atom.is_valid
        positions[index] = atom.pos
        if atom.is_valid:
            ref_element[index] = get_element_atomic_num(atom.element)
            ref_name[index] = encode_atom_name(atom.name)
            atom_to_token[index] = atom.token_index

    resolved = valid & np.any(positions != 0, axis=1)
    X = torch.from_numpy(positions)
    resolved_mask = torch.from_numpy(resolved)
    valid_mask = torch.from_numpy(valid)
    if resolved_mask.any():
        X = X - X[resolved_mask].mean(dim=0, keepdim=True)
        X[~valid_mask] = 0.0
    return {
        "ref_pos": torch.from_numpy(ref_pos),
        "ref_element": torch.from_numpy(ref_element),
        "ref_charge": torch.from_numpy(ref_charge),
        "ref_atom_name_chars": torch.from_numpy(ref_name),
        "ref_space_uid": torch.from_numpy(ref_space),
        "gt_coords": X.float().unsqueeze(0),
        "atom_attention_mask": torch.from_numpy(atom_mask),
        "atom_to_token": torch.from_numpy(atom_to_token),
        "is_resolved": torch.tensor(resolved, dtype=torch.bool),
    }


def build_feature_tensors(
    chains: list[ChainInfo],
    tokens: list[TokenInfo],
    atoms: list[AtomInfo],
    input: StructurePredictionInput,
) -> dict[str, torch.Tensor]:
    """Assemble the complete unbatched ESMFold2 feature dictionary."""
    token_features = _token_tensors(tokens)
    atom_features = _atom_tensors(_padded_atoms(atoms))
    frames, _ = compute_frame_indices(tokens, atoms)
    msa_features = compute_msa_features(input, chains, tokens)
    distogram, distogram_mask = compute_distogram_conditioning(
        input,
        chains,
        tokens,
        torch.zeros(len(tokens), 3, dtype=torch.float32),
    )
    return {
        **token_features,
        "token_bonds": compute_token_bonds(tokens, atoms, input, chains),
        "token_attention_mask": torch.ones(len(tokens), dtype=torch.bool),
        "pocket_feature": torch.zeros(len(tokens), dtype=torch.long),
        **atom_features,
        "distogram_atom_idx": compute_representative_atoms(tokens, atoms),
        "frames_idx": torch.from_numpy(frames).to(torch.int64),
        "disto_cond": distogram,
        "disto_cond_mask": distogram_mask,
        **msa_features,
    }


def prepare_esmfold2_input(
    input: StructurePredictionInput, seed: int | None = None
) -> tuple[dict[str, torch.Tensor], list[ChainInfo]]:
    """Convert one typed request to model features and output-chain metadata."""
    chains, tokens, atoms = build_chains_from_input(input, seed)
    return build_feature_tensors(chains, tokens, atoms, input), chains