# Copyright 2026 AlQuraishi Laboratory
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Contains code related to parsing Query objects into AtomArrays and processed reference
molecules.
"""

import logging
from collections.abc import Iterable
from functools import lru_cache
from typing import NamedTuple

import biotite.structure as struc
import numpy as np
from biotite.interface.rdkit import from_mol, to_mol
from biotite.structure import AtomArray
from rdkit import Chem

from openfold3.core.data.pipelines.sample_processing.conformer import (
    ProcessedReferenceMolecule,
)
from openfold3.core.data.primitives.structure.cleanup import remove_hydrogens
from openfold3.core.data.primitives.structure.component import set_atomwise_annotation
from openfold3.core.data.primitives.structure.conformer import (
    multistrategy_compute_conformer,
)
from openfold3.core.data.resources.residues import (
    DNA_RESTYPE_1TO3,
    MOLECULE_TYPE_TO_LEAVING_ATOMS,
    MOLECULE_TYPE_TO_UKNOWN_RESIDUES_3,
    PROTEIN_RESTYPE_1TO3,
    RNA_RESTYPE_1TO3,
    MoleculeType,
)
from openfold3.projects.of3_all_atom.config.inference_query_format import Query

logger = logging.getLogger(__name__)


class StructureWithReferenceMolecules(NamedTuple):
    """Central object required for structure feature-creation in inference.

    Attributes:
        atom_array (struc.AtomArray):
            AtomArray parsed from the input Query for which coordinates will be
            predicted.
        processed_reference_mols (list[ProcessedReferenceMolecule]):
            List of processed reference molecules (RDKit mol objects with atom names and
            computed conformers) that are required for feature construction.
    """

    atom_array: struc.AtomArray
    processed_reference_mols: list[ProcessedReferenceMolecule]


get_residue_cached = lru_cache(maxsize=500)(struc.info.residue)
"""Cached residue information retrieval from Biotite to speed up preprocessing."""


def get_leaving_atoms(ccd_code: str) -> np.ndarray:
    """Returns the leaving atoms for a given CCD code.

    Args:
        ccd_code (str):
            The CCD code of the residue to get the leaving atoms for.

    Returns:
        np.ndarray:
            An array of leaving atom names for the given CCD code.
    """
    leaving_atom_flag = struc.info.get_from_ccd(
        category_name="chem_comp_atom",
        comp_id=ccd_code,
        column_name="pdbx_leaving_atom_flag",
    ).as_array()
    atom_names = struc.info.get_from_ccd(
        category_name="chem_comp_atom",
        comp_id=ccd_code,
        column_name="atom_id",
    ).as_array()

    leaving_atoms = atom_names[leaving_atom_flag == "Y"]

    return leaving_atoms


def atom_array_from_ccd_code(
    ccd_code: str,
    chain_id: str,
    res_id: int = 1,
    molecule_type: MoleculeType | None = None,
) -> AtomArray:
    """Creates an AtomArray from a CCD code.

    Fetches the residue information from Biotite and constructs an AtomArray with the
    specified chain ID, residue ID, and molecule type.

    Args:
        ccd_code (str):
            The CCD code of the residue to create the AtomArray from.
        chain_id (str):
            The chain ID to assign to the created AtomArray.
        res_id (int):
            The residue ID to assign to the created AtomArray. Defaults to 1.
        molecule_type (MoleculeType | None):
            The MoleculeType of the molecule. If None, no molecule type annotation will
            be set. Defaults to None.
    """
    res_array = get_residue_cached(ccd_code)
    res_array = remove_hydrogens(res_array)

    res_array.res_id[:] = res_id
    res_array.chain_id[:] = chain_id

    if molecule_type is not None:
        res_array.set_annotation(
            "molecule_type_id", np.repeat(molecule_type, len(res_array))
        )

    return res_array


def atom_array_from_mol(
    mol: Chem.Mol,
    atom_names: Iterable[str],
    chain_id: str,
    molecule_type: MoleculeType = MoleculeType.LIGAND,
    res_id: int = 1,
    res_name: str = "LIG",
) -> AtomArray:
    """Creates an AtomArray from an RDKit mol object.

    Args:
        mol (Chem.Mol):
            The RDKit molecule to create the AtomArray from.
        atom_names (Iterable[str]):
            Iterable of atom names to set in the AtomArray.
        chain_id (str):
            The chain ID to assign to the created AtomArray.
        molecule_type (MoleculeType):
            The MoleculeType of the molecule. Defaults to MoleculeType.LIGAND.
        res_id (int):
            The residue ID to assign to the created AtomArray. Defaults to 1.
        res_name (str):
            The residue name to assign to the created AtomArray. Defaults to "LIG".

    Returns:
        AtomArray:
            An AtomArray containing the atoms from the RDKit mol with the specified atom
            names, chain ID, residue ID, and residue name, and molecule type. The order
            of atoms in the AtomArray will match the order of atom names provided.
    """
    atom_array = from_mol(mol, conformer_id=0, add_hydrogen=False)

    # Set global annotations
    atom_array.chain_id[:] = chain_id
    atom_array.hetero[:] = True
    atom_array.res_id[:] = res_id
    atom_array.set_annotation(
        "molecule_type_id", np.repeat(molecule_type, len(atom_array))
    )

    # Set specific annotations
    atom_array.set_annotation("res_name", np.repeat(res_name, len(atom_array)))
    atom_array.set_annotation("atom_name", atom_names)

    return atom_array


def processed_reference_molecule_from_atom_array(
    atom_array: struc.AtomArray,
    atoms_to_mask: Iterable[str] = None,
) -> ProcessedReferenceMolecule:
    """Creates a processed reference molecule from an AtomArray.

    Args:
        atom_array (struc.AtomArray):
            The AtomArray to create the processed reference molecule from. Atom names of
            the processed reference molecule will be set to the atom names of the
            AtomArray.
        atoms_to_mask (Iterable[str] | None):
            Optional iterable of atom names to mask in the processed reference molecule.
            If provided, these atoms will not be included in the final structure, but
            will still be part of the RDKit mol object to retain chemical validity and
            generate the correct conformer. If None, which is the default, no atoms will
            be masked.

    Returns:
        ProcessedReferenceMolecule:
            A processed reference molecule containing the RDKit mol with a computed
            conformer and the atom mask.
    """
    # Mask certain atoms that should not be present in the final structure
    if atoms_to_mask is not None:
        atom_mask = ~np.isin(atom_array.atom_name, atoms_to_mask)
    else:
        atom_mask = np.ones(len(atom_array), dtype=bool)

    # Convert to RDKit mol
    mol = to_mol(atom_array, kekulize=True)
    Chem.SanitizeMol(mol)
    mol.RemoveConformer(0)

    return processed_reference_molecule_from_mol(
        mol=mol,
        atom_names=atom_array.atom_name,
        atom_mask=atom_mask,
    )


def processed_reference_molecule_from_mol(
    mol: Chem.Mol,
    atom_names: Iterable[str] | None = None,
    atom_mask: np.ndarray | None = None,
) -> ProcessedReferenceMolecule:
    """Creates a processed reference molecule from an RDKit mol object.

    Args:
        mol (Chem.Mol):
            The RDKit molecule to create the processed reference molecule from.
        atom_names (Iterable[str] | None):
            Optional atom names to set for the atoms in the RDKit mol. If None, the atom
            names will be set to a simple pattern like C1, C2, N1, N2, etc.
        atom_mask (np.ndarray | None):
            Optional mask for atoms in the processed reference molecule that should not
            be included in the feature creation, e.g. leaving atoms. Those atoms will
            still be part of the rdkit.Mol object to retain chemical validity of the
            molecule and generate the correct conformer. If None, which is the default,
            no atoms will be masked.

    Returns:
        ProcessedReferenceMolecule:
            A processed reference molecule containing the RDKit mol with a computed
            conformer and the atom mask.
    """
    # Compute conformer (note that we call this before creating the annotations, as this
    # function will remove all hydrogens in the input mol and can therefore change the
    # mask length)
    mol, conf_id, _ = multistrategy_compute_conformer(
        mol, remove_hs=True, timeout_standard=120, timeout_rand_init=120
    )
    assert conf_id == 0

    # Assume all atoms are in the structure if no special mask is given
    if atom_mask is None:
        atom_mask = np.ones(mol.GetNumAtoms(), dtype=bool)

    # Set atom names if provided, otherwise renumber to C1, C2, N1, N2, etc.
    if atom_names is not None:
        mol = set_atomwise_annotation(mol, "atom_name", atom_names)
    else:
        elements = [atom.GetSymbol().upper() for atom in mol.GetAtoms()]
        atom_names = struc.create_atom_names(elements)
        mol = set_atomwise_annotation(mol, "atom_name", atom_names)

    # This is a different mask only required for fallback conformers in the training
    # script where some coordinates are not defined
    mol = set_atomwise_annotation(mol, "used_atom_mask", [True] * mol.GetNumAtoms())

    return ProcessedReferenceMolecule(
        mol=mol,
        in_crop_mask=atom_mask,
        permutations=None,
    )


def structure_with_ref_mols_from_sequence(
    sequence: str,
    poly_type: MoleculeType,
    chain_id: str,
    non_canonical_residues: dict[int, str] | None = None,
) -> StructureWithReferenceMolecules:
    """Builds an AtomArray and processed reference molecules from a sequence.

    Will read the entire sequence into an AtomArray and create reference molecule
    objects with separate conformers for each residue. Currently only supports standard
    residues, any non-canonical residue will be treated as an unknown residue.

    Args:
        sequence (str):
            The sequence of the polymeric molecule as a string of 1-letter residue
            codes.
        poly_type (MoleculeType):
            The MoleculeType of the polymeric molecule. Should be one of
            MoleculeType.PROTEIN, MoleculeType.DNA, or MoleculeType.RNA.
        chain_id (str):
            The chain ID to assign to the created AtomArray.
        non_canonical_residues (dict[int, str] | None):
            A dictionary mapping residue IDs to non-canonical residue names. Defaults to
            None.

    Returns:
        StructureWithReferenceMolecules:
            A named tuple containing the AtomArray and a list of processed reference
            molecules, each corresponding to a residue in the sequence.
    """
    if non_canonical_residues is None:
        non_canonical_residues = {}

    # Figure out 3-letter code mapping
    match poly_type:
        case MoleculeType.PROTEIN:
            resname_1_to_3 = PROTEIN_RESTYPE_1TO3
        case MoleculeType.DNA:
            resname_1_to_3 = DNA_RESTYPE_1TO3
        case MoleculeType.RNA:
            resname_1_to_3 = RNA_RESTYPE_1TO3
        case _:
            raise ValueError(f"Unsupported molecule type: {poly_type}")

    # Figure out the unknown residue 3-letter identifier and leaving atom names
    unk_res = MOLECULE_TYPE_TO_UKNOWN_RESIDUES_3[poly_type]
    base_leaving_atoms = MOLECULE_TYPE_TO_LEAVING_ATOMS[poly_type]

    atom_array = None
    processed_reference_mols = []

    for res_id, resname_1 in enumerate(sequence, start=1):
        # Swap to non-standard residue if necessary
        if res_id in non_canonical_residues:
            resname_3 = non_canonical_residues[res_id]
            leaving_atoms = get_leaving_atoms(resname_3)

            if len(leaving_atoms) == 0:
                logger.warning(
                    f"Non-canonical residue {resname_3} at position {res_id} has no "
                    "leaving atoms defined. Using default leaving atoms."
                )
                leaving_atoms = base_leaving_atoms

        else:
            leaving_atoms = base_leaving_atoms

            # Get 3-letter code of standard residue
            if resname_1 in resname_1_to_3:
                resname_3 = resname_1_to_3[resname_1]

            # Set unknown placeholder residue
            else:
                logger.warning(
                    f"Unknown residue {resname_1} at position {res_id} in sequence. "
                    f"Using placeholder residue {unk_res}."
                )
                resname_3 = unk_res

        # Construct atom array for the residue
        res_array = atom_array_from_ccd_code(
            resname_3,
            chain_id=chain_id,
            res_id=res_id,
            molecule_type=poly_type,
        )

        # Parse into RDKit mol and compute conformer
        processed_ref_mol = processed_reference_molecule_from_atom_array(
            res_array, atoms_to_mask=leaving_atoms
        )
        processed_reference_mols.append(processed_ref_mol)

        # Remove the leaving atoms from the atom array
        leaving_atom_mask = np.isin(res_array.atom_name, leaving_atoms)
        if not leaving_atom_mask.any():
            logger.warning(
                f"Residue {resname_3} at position {res_id} has no leaving atoms to "
                "remove. This could cause issues with incorrect polymer linkage."
            )
        res_array = res_array[~leaving_atom_mask]

        # Initialize atom array
        if atom_array is None:
            atom_array = res_array

        # Append to atom array
        else:
            atom_array += res_array

    # Auto-connect bonds
    atom_array.bonds = struc.connect_via_residue_names(atom_array)

    # Force coordinates to 0 for consistency
    atom_array.coord[:] = 0.0

    return StructureWithReferenceMolecules(
        atom_array=atom_array,
        processed_reference_mols=processed_reference_mols,
    )


def structure_with_ref_mol_from_mol(
    mol: Chem.Mol,
    chain_id: str,
    atom_mask: np.ndarray | None = None,
    res_name: str = "LIG",
) -> StructureWithReferenceMolecules:
    """Creates a single AtomArray and processed reference molecule from an RDKit mol.

    Args:
        mol (Chem.Mol):
            The RDKit molecule to create the AtomArray and processed reference molecule
            from.
        chain_id (str):
            The chain ID to assign to the created AtomArray.
        atom_mask (np.ndarray | None):
            Optional mask for atoms to include in the processed reference molecule. If
            None, all atoms will be included.
        res_name (str):
            The residue name to assign to the created AtomArray. Defaults to "LIG".
    Returns:
        StructureWithReferenceMolecules:
            A named tuple containing the AtomArray and a list with a single processed
            reference molecule. The residue ID will be set to 1.
    """

    # Build the ligand molecule
    proc_ref_mol = processed_reference_molecule_from_mol(mol, atom_mask=atom_mask)

    # Get the processed mol that now will have a computed conformer
    mol = proc_ref_mol.mol

    # Retrieve atom names from special annotation in the reference mol
    atom_names = [atom.GetProp("annot_atom_name") for atom in mol.GetAtoms()]

    # Convert to AtomArray
    atom_array = atom_array_from_mol(
        mol, atom_names=atom_names, chain_id=chain_id, res_name=res_name
    )

    # Force coordinates to 0 for consistency
    atom_array.coord[:] = 0.0

    return StructureWithReferenceMolecules(
        atom_array=atom_array, processed_reference_mols=[proc_ref_mol]
    )


def structure_with_ref_mol_from_ccd_code(
    ccd_code: str,
    chain_id: str,
) -> StructureWithReferenceMolecules:
    """Creates a single AtomArray and processed reference molecule from a CCD code.

    Args:
        ccd_code (str):
            The CCD code of the molecule to create.
        chain_id (str):
            The chain ID to assign to the created AtomArray.

    Returns:
        StructureWithReferenceMolecules:
            A named tuple containing the AtomArray and a list with a single processed
            reference molecule. The residue ID will be set to 1.
    """

    # Build ligand AtomArray
    atom_array = atom_array_from_ccd_code(
        ccd_code,
        chain_id=chain_id,
        res_id=1,
        molecule_type=MoleculeType.LIGAND,
    )

    # Get processed reference molecule
    proc_ref_mol = processed_reference_molecule_from_atom_array(atom_array)

    # Force coordinates to 0 for consistency
    atom_array.coord[:] = 0.0

    return StructureWithReferenceMolecules(
        atom_array=atom_array, processed_reference_mols=[proc_ref_mol]
    )


def structure_with_ref_mol_from_smiles(
    smiles: str,
    chain_id: str,
    res_name: str = "LIG",
) -> StructureWithReferenceMolecules:
    """Creates a single AtomArray and processed ref molecule from a SMILES string.

    Args:
        smiles (str):
            The SMILES string of the molecule to create.
        chain_id (str):
            The chain ID to assign to the created AtomArray.
        res_name (str):
            The residue name to assign to the created AtomArray. Defaults to "LIG".

    Returns:
        StructureWithReferenceMolecules:
            A named tuple containing the AtomArray and a list with a single processed
            reference molecule. The residue ID will be set to 1. Atom names of the
            molecule will be set to follow the pattern C1, C2, N1, N2, etc.
    """
    mol = Chem.MolFromSmiles(smiles)

    return structure_with_ref_mol_from_mol(
        mol,
        chain_id=chain_id,
        res_name=res_name,
    )


def structure_with_ref_mols_from_query(query: Query) -> StructureWithReferenceMolecules:
    """Builds an AtomArray and processed reference molecules from a Query object.

    Parses the Query object into a full AtomArray and processed reference molecules
    (RDKit mol objects with atom names and computed conformers).

    The returned AtomArray follows the chain IDs given in the Query object. If a chain
    specifies multiple chain IDs, repeated identical chains with those IDs will be
    constructed and given the same entity ID.

    Residue names will be inferred from the sequence or CCD codes. If a ligand is
    specified through a SMILES string, it will be named as "LIG".

    Args:
        query (Query):
            The Query object containing the chains to construct the structure from.

    Returns:
        StructureWithReferenceMolecules:
            A named tuple containing the AtomArray and a list of processed reference
            molecules.
    """
    # Initialize eventually returned objects
    atom_array = None
    processed_reference_mols: list[ProcessedReferenceMolecule] = []

    # Create entity mapping
    all_entities = set()
    for chain in query.chains:
        if chain.sequence is not None:
            all_entities.add(chain.sequence)
        elif chain.ccd_codes is not None:
            for ccd in chain.ccd_codes:
                all_entities.add(ccd)
        elif chain.smiles is not None:
            all_entities.add(chain.smiles)
    all_entities = sorted(all_entities)
    entity_to_id = {e: i + 1 for i, e in enumerate(all_entities)}

    # Create smiles->comp ID mapping to allow for distinguishing different ligands
    # specified by smiles
    all_smiles = set()
    for chain in query.chains:
        if chain.smiles is not None:
            all_smiles.add(chain.smiles)
    if len(all_smiles) > 0:
        all_smiles = sorted(all_smiles)
        smiles_to_comp_id = {s: f"LIG{i}" for i, s in enumerate(all_smiles)}
    else:
        smiles_to_comp_id = {}

    # Build the structure segment-wise from all chains in the query.
    for chain in query.chains:
        for chain_id in chain.chain_ids:
            match chain.molecule_type:
                # Build polymeric segment
                case MoleculeType.PROTEIN | MoleculeType.DNA | MoleculeType.RNA:
                    segment_atom_array, segment_ref_mols = (
                        structure_with_ref_mols_from_sequence(
                            sequence=chain.sequence,
                            poly_type=chain.molecule_type,
                            chain_id=chain_id,
                            non_canonical_residues=chain.non_canonical_residues,
                        )
                    )
                    representation = chain.sequence

                # Build ligand molecule
                case MoleculeType.LIGAND:
                    # Build ligand from SMILES
                    if chain.smiles is not None:
                        segment_atom_array, segment_ref_mols = (
                            structure_with_ref_mol_from_smiles(
                                smiles=chain.smiles,
                                chain_id=chain_id,
                                res_name=smiles_to_comp_id[chain.smiles],
                            )
                        )
                        representation = chain.smiles

                    # Build ligand from CCD code
                    elif chain.ccd_codes is not None:
                        # TODO: add multi-residue ligand support
                        if len(chain.ccd_codes) > 1:
                            raise NotImplementedError(
                                "Multiple CCD codes for a single chain are not yet "
                                "supported."
                            )

                        segment_atom_array, segment_ref_mols = (
                            structure_with_ref_mol_from_ccd_code(
                                ccd_code=chain.ccd_codes[0],
                                chain_id=chain_id,
                            )
                        )
                        representation = chain.ccd_codes[0]

                    # Build ligand from SDF file
                    elif chain.sdf_file_path is not None:
                        # TODO: add SDF support
                        raise NotImplementedError(
                            "SDF format for ligands is not yet supported."
                        )

                    else:
                        raise ValueError("No valid molecule specification found.")

            # Add processed reference molecules
            processed_reference_mols.extend(segment_ref_mols)

            segment_atom_array.set_annotation(
                "entity_id",
                np.repeat(entity_to_id[representation], len(segment_atom_array)),
            )

            # Append atom array to end
            if atom_array is None:
                atom_array = segment_atom_array
            else:
                atom_array += segment_atom_array

    # Force coordinates to 0 for consistency
    atom_array.coord[:] = 0.0

    return StructureWithReferenceMolecules(
        atom_array=atom_array, processed_reference_mols=processed_reference_mols
    )
