"""Parser class for different molecule files."""

# ========================
# Python internal packages
# ========================
# I
import io

# L
import logging

# P
import pathlib

# T
import typing

# W
import warnings

# =================
# External packages
# =================
# M
import MDAnalysis
from MDAnalysis.guesser import default_guesser

# O
import openmm

# P
import pdbfixer

# N
import numpy

# R
import rdkit.Chem as Chem
import rdkit.Chem.rdDetermineBonds as rdDetermineBonds
import rdkit.Geometry.rdGeometry as rdGeometry

# ==============
# Module package
# ==============
# U
from ..utils.messenger import (
    Info,
    Warn,
    Error,
)


class Molecule:
    """Container class that bridges RDKit and MDAnalysis molecule
    representations.
    """

    def __init__(
        self,
        rdkit: Chem.rdchem.Mol,
        mda: MDAnalysis.Universe | None = None,
    ) -> None:
        """Initialize the molecule representations.

        Parameters
        ----------
        rdkit : `Chem.rdchem.Mol`
            The RDKit molecule object.

        mda : `MDAnalysis.Universe`, optional
            The MDAnalysis Universe object. By default, `None`.
        """
        self.rdkit: Chem.rdchem.Mol = rdkit
        self.mda: MDAnalysis.Universe | None = mda

    def get_index(self, coordinate: numpy.ndarray) -> int | None:
        """Return the atom index matching a given coordinate.

        Parameters
        ----------
        coordinate : `numpy.ndarray`
            A 3D coordinate array.

        Returns
        -------
        `int | None`
            The index of the matching atom. If no atom matches the provided
            coordinate, `None` is return.
        """
        difference_array: numpy.ndarray = numpy.round(
            numpy.abs(coordinate - self.mda.atoms.positions), 5
        )
        sum_array: numpy.ndarray = numpy.sum(difference_array, axis=1)

        if not numpy.isclose(numpy.min(sum_array), 0):
            return None

        return int(sum_array.argmin())

    def get_data(self, index: int) -> dict:
        """Retrieve atom data from the MDAnalysis Universe.

        Parameters
        ----------
        index : `int`
            The index of the atom to extract metadata from.

        Returns
        -------
        `dict`
            A dictionary containing atom information (atom name, type and ID,
            residue name and ID, chain name).
        """
        atom: MDAnalysis.AtomGroup = self.mda.atoms[index]

        return {
            "atom_name": atom.name,
            "atom_type": atom.element,
            "atom_id": int(atom.id),
            "residu_name": atom.resname,
            "residu_id": int(atom.resid),
            "chain": atom.segid,
        }


class MDAnalysisUniverse:
    """Wrapper class around MDAnalysis to handle molecule loading,
    selection, guessing, and conversion to RDKit.
    """

    def __init__(
        self,
        file: pathlib.Path,
        selection: str,
        vdw_radius: dict[str, float] | None = None,
    ) -> None:
        """Initialize the MDAnalysis universe wrapper.

        Parameters
        ----------
        file : `pathlib.Path`
            Path to the molecular structure file.

        selection : `str`
            Atom selection string to apply to the molecule.

        vdw_radius : `dict[str, float]`, optional
            Van der Waals radii used for bond guessing. By default, use
            `MDAnalysis` defined radius.
        """
        self.file: pathlib.Path = file
        self.selection: str = selection
        self.vdw_radius: dict

        if vdw_radius is None:
            self.vdw_radius = MDAnalysis.guesser.tables.vdwradii.copy()
        else:
            self.vdw_radius = vdw_radius

        self.molecule: MDAnalysis.Universe | MDAnalysis.AtomGroup | None = None

    def __call__(self, **k_w_arg) -> typing.Self:
        """Create the MDAnalysis Universe.

        Parameters
        ----------
        **k_w_arg
            Keyword arguments forwarded to `MDAnalysis.Universe`.

        Returns
        -------
        `typing.Self`
            The current instance with an initialized MDAnalysis universe.

        Raises
        ------
        `ValueError`
            - If the input file format is not supported by MDAnalysis.
        """
        try:
            # To ignore MDAnalysis that ask us to guess things, as we do it
            # later on.
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", category=UserWarning)
                self.molecule = MDAnalysis.Universe(**k_w_arg)
        except ValueError:
            Error.MDA_FORMAT(
                ValueError,
                extension=self.file.suffix,
            )

        return self

    def select(self, selection: str) -> typing.Self:
        """Apply an atom selection to the universe.

        Parameters
        ----------
        selection : `str`
            MDAnalysis selection string.

        Returns
        -------
        `typing.Self`
            The current instance with the selected atoms.

        Raises
        ------
        `ValueError`
            - If the selection string is invalid.
            - If the selection results in an empty atom group.
        """
        if selection == "":
            self.molecule = self.molecule.atoms
            return self

        try:
            molecule: MDAnalysis.AtomGroup = self.molecule.select_atoms(
                selection
            )
        except MDAnalysis.exceptions.SelectionError:
            Error.WRONG_SELECTION(ValueError, selection=selection)

        if len(molecule.atoms) <= 0:
            Error.EMPTY_SELECTION(ValueError, selection=selection)

        self.molecule = molecule

        return self

    def guesser(self) -> typing.Self:
        """Infer missing topology information and guess bonds.

        Returns
        -------
        `typing.Self`
            The current instance with updated topology information.

        Warn
        ----
        `UserWarning`
            - Catch `MDAnalysis` error and throw a warning instead.
        """
        guesser: default_guesser.DefaultGuesser = (
            default_guesser.DefaultGuesser(self.molecule)
        )

        to_guess: dict[str, typing.Callable] = {
            "types": guesser.guess_types,
            "elements": guesser.guess_types,
        }

        for key_guess, method in to_guess.items():
            try:
                self.molecule.add_TopologyAttr(key_guess, method())
            except ValueError as value_error:
                Warn.MDA_GUESS(warning=value_error)

        self.molecule.atoms.guess_bonds(vdwradii=self.vdw_radius)

        return self

    def mda_to_rdkit(self) -> Molecule:
        """Convert the MDAnalysis object into an RDKit molecule.

        Returns
        -------
        `Molecule`
            A molecule object bridging RDKit and MDAnalysis.

        Warn
        ----
        `UserWarning`
            - If implicit hydrogens are missing and explicit hydrogens
              must be used for conversion.
        """
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=UserWarning)

            try:
                rdkit = self.molecule.atoms.convert_to("RDKIT")
            except AttributeError:
                Warn.NO_IMPLICIT(
                    file=self.file,
                    selection=self.selection,
                )

                rdkit = self.molecule.atoms.convert_to(
                    "RDKIT", NoImplicit=False
                )

        return Molecule(rdkit, self.molecule)


class MoleculeParser:
    """Factory class for parsing molecular structure files and returning
    a `Molecule` instance.
    """

    file: pathlib.Path
    index: int
    selection: str
    add_hydrogen: bool

    def __new__(
        cls,
        file: pathlib.Path,
        index: int = 0,
        add_hydrogen: bool = False,
        threshold: float = 0.05,
        selection: str = "",
        vdw_radius: dict[str, float] | None = None,
    ) -> Molecule:
        """Parse a molecular structure file and return a `Molecule`.

        Parameters
        ----------
        file : `pathlib.Path`
            Path to the molecular structure file.

        index : `int`, optional
            Index of the molecule to retrieve when multiple molecules are
            present in the file. By default, `0`.

        add_hydrogen : `bool`, optional
            Whether missing hydrogens should be added when parsing
            `.pdb` or `.pdbqt` files.

        treshold : `float`, optional
            Planarity threshold above which a ring is not considered
            aromatic.

        selection : `str`, optional
            Atom selection applied through `MDAnalysis`.

        vdw_radius : `dict[str, float]`, optional
            Van der Waals radii used by MDAnalysis to guess bonds.

        Returns
        -------
        `Molecule`
            The parsed molecule.

        Warns
        -----
        - When aromaticity must be reassigned manually for `.pdb`
          or `.pdbqt` files.
        - When atom selection is not supported for the detected format.

        Raises
        ------
        `ValueError`
            - If the molecule cannot be built.
            - If the molecule is only partially built.
        """
        # Cleaning weird `RDKit` errors.
        log: logging.Logger = logging.getLogger("rdkit")
        log.setLevel(logging.ERROR)

        cls.file: pathlib.Path = file
        cls.index: int = index
        cls.selection: str = selection
        cls.add_hydrogen: bool = add_hydrogen
        cls.MDAnalysisUniverse: MDAnalysisUniverse = MDAnalysisUniverse(
            selection=selection, file=file, vdw_radius=vdw_radius
        )

        file_extension: str = file.suffix.lower()[1:]

        # ==========================
        # Try to build the molecule.
        # ==========================
        if hasattr(cls, file_extension):
            molecule: Molecule = getattr(cls, file_extension)()
        # Else, using MDAnalysis.
        else:
            molecule = cls.default()

        if selection != "" and molecule.mda is None:
            Warn.SELECTION_NOT_SUPPORTED(extension=file.suffix.lower())

        if molecule.mda is not None:
            Warn.AROMATICITY(file=file, selection=selection)
            cls.assign_cycle(molecule, threshold)

        # =========================================
        # Check if the molecule is correctly built.
        # =========================================
        if molecule.rdkit is None:
            Error.BUILD_FAIL(ValueError, file=file)

        if not isinstance(molecule.rdkit, Chem.rdchem.Mol):
            Error.PARTIAL_BUILD(ValueError, file=file)

        return molecule

    @classmethod
    def assign_cycle(cls, molecule: Molecule, threshold: float) -> None:
        """Assigns aromaticity to cyclic structures in a molecule based on if
        the cycle is plan or not.

        Parameters
        ----------
        molecule : `Molecule`
            An RDKit molecule object containing atom and ring information.

        threshold : `float`
            The threshold for determining planarity. Rings with maximum atom
            deviation above this value will not be marked as aromatic.

        Warns
        -----
        - When skipping a cycle of 3 or 4 atoms.
        """
        position: typing.Callable[[int], rdGeometry.Point3D] = (
            molecule.rdkit.GetConformer().GetAtomPosition
        )

        for ring in molecule.rdkit.GetRingInfo().AtomRings():
            if len(ring) in [3, 4]:
                Warn.SKIP_CYCLE(length=len(ring))
                continue

            coord_list: None | numpy.ndarray = None

            for atom_id in ring:
                coordinate: rdGeometry.Point3D = position(atom_id)

                if coord_list is None:
                    coord_list = numpy.array(
                        [[coordinate.x, coordinate.y, coordinate.z]]
                    )
                else:
                    coord_list = numpy.vstack(
                        (
                            coord_list,
                            [coordinate.x, coordinate.y, coordinate.z],
                        )
                    )

            max_overlap: numpy.ndarray = numpy.max(
                planarity_computation(coord_list)
            )

            if max_overlap > threshold:
                continue

            for i, atom_id in enumerate(ring):
                molecule.rdkit.GetAtomWithIdx(atom_id).SetIsAromatic(True)

                bond: Chem.rdchem.BondType = (
                    molecule.rdkit.GetBondBetweenAtoms(
                        atom_id, ring[(i + 1) % len(ring)]
                    )
                )

                bond.SetIsAromatic(True)
                bond.SetBondType(Chem.rdchem.BondType.AROMATIC)

    @classmethod
    def mol(cls) -> Molecule:
        """Parse a `.mol` file.

        Returns
        -------
        `Molecule`
            The parsed molecule object or `None` if parsing fails.
        """
        molecule: Chem.rdchem.Mol | None = Chem.rdmolfiles.MolFromMolFile(
            str(cls.file), removeHs=False
        )

        return Molecule(molecule)

    @classmethod
    def mol2(cls) -> Molecule:
        """Parse a `.mol2` file and return a molecule object.

        Returns
        -------
        `Molecule`
            The parsed molecule object.

        Raises
        ------
        `ValueError`
            If the file cannot be parsed or the atom-typing scheme is incorrect.
        """
        molecule: Chem.rdchem.Mol | None = Chem.rdmolfiles.MolFromMol2File(
            str(cls.file), removeHs=False
        )

        # The `.mol2` format need a specific atom-typing scheme.
        if molecule is None:
            Error.MOL_2(ValueError, cls.file)

        return Molecule(molecule)

    @classmethod
    def tpl(cls) -> Molecule:
        """
        Parse a `.tpl` file and return a molecule object.

        Returns
        -------
        `Molecule`
            The parsed molecule object.
        """
        molecule: Chem.rdchem.Mol | None = Chem.rdmolfiles.MolFromTPLFile(
            str(cls.file), removeHs=False
        )

        return Molecule(molecule)

    @classmethod
    def xyz(cls) -> Molecule:
        """Parse a `.xyz` file, determine connectivity, and return a molecule
        object.

        Returns
        -------
        `Molecule`
            The parsed molecule object.
        """
        molecule: Chem.rdchem.Mol | None = Chem.rdmolfiles.MolFromXYZFile(
            str(cls.file)
        )
        rdDetermineBonds.DetermineConnectivity(molecule)

        return Molecule(molecule)

    @classmethod
    def sdf(cls) -> Molecule:
        """Parse an `.sdf` file and return the specified molecule object.

        Returns
        -------
        `Molecule`
            The molecule object at the specified index in the file.

        Raises
        ------
        `IndexError`
            If the index is out of range for the molecules in the file.
        """
        supplier: list[Chem.rdchem.Mol | None] = list(
            Chem.rdmolfiles.SDMolSupplier(str(cls.file), removeHs=False)
        )

        if cls.index >= len(supplier):
            Error.INDEX_OVER_LEN(
                IndexError,
                index=cls.index,
                length=len(supplier),
                file=cls.file,
            )

        return Molecule(supplier[cls.index])

    @classmethod
    def pdb(cls) -> Molecule | None:
        """Parse a `.pdb` file using PDBFixer and MDAnalysis.

        Returns
        -------
        `Molecule | None`
            The parsed molecule, or `None` if parsing fails.

        Raises
        ------
        `RuntimeError`
            - If hydrogen addition or conversion fails.
        """
        fixer: pdbfixer.PDBFixer = pdbfixer.PDBFixer(filename=str(cls.file))

        if cls.add_hydrogen:
            Info.ADD_HYDROGEN(file=cls.file)

            fixer.findMissingResidues()
            fixer.findMissingAtoms()
            fixer.findNonstandardResidues()
            fixer.addMissingAtoms()

            try:
                fixer.addMissingHydrogens(pH=7.0)
            except ValueError:
                Error.SAVE_PDB_FAIL(RuntimeError, file=cls.file)

        # Write to a string buffer instead of a file
        buffer: io.StringIO = io.StringIO()
        openmm.app.PDBFile.writeFile(
            fixer.topology, fixer.positions, buffer, keepIds=True
        )
        _ = buffer.seek(0)

        try:
            return (
                cls.MDAnalysisUniverse(topology=buffer, format="PDB")
                .guesser()
                .select(cls.selection)
                .mda_to_rdkit()
            )
        except Chem.rdchem.AtomValenceException:
            Error.SAVE_PDB_FAIL(RuntimeError, file=cls.file)

        return None

    @classmethod
    def default(cls) -> Molecule:
        """Fallback parser using MDAnalysis for unsupported file formats.

        Returns
        -------
        `Molecule`
            The parsed molecule.
        """
        return (
            cls.MDAnalysisUniverse(topology=cls.file)
            .guesser()
            .select(cls.selection)
            .mda_to_rdkit()
        )


def planarity_computation(coord_list: numpy.ndarray) -> numpy.array:
    r"""Compute the overlap values used to assess ring planarity.

    Mathematical used equations
    ---------------------------
    **Geometric centre:**

    > **Note 📝**
    >
    > This equation is applied for coordinates (x, y, z).

    $$
    \text{geometric centre}_x = \dfrac{1}{\text{number atoms}}%%
    \sum _{i = 1}^{\text{number atoms}} \mathbf{x}_i
    $$

    **Vectors from geometrical centre of a cycle and each of its atoms:**

    $$
    \hat{v}_i =%%
    \dfrac{\vec{\text{geometric centre}} - \vec{\text{coordinate}}_i}%%
    {\left\lVert \vec{\text{geometric centre}}%%
    - \vec{\text{coordinate}_i} \right\rVert}
    $$

    **Normal vector:**

    $$
    \hat{\mathbf{n}} = \dfrac{\vec{v}_1 \times \vec{v}}%%
    {\left\lVert \vec{v}_1 \times \vec{v}_2 \right\rVert}
    $$

    **Vectors overlapping:**

    $$
    {\text{overlap}}_i = \hat{v}_i \cdot \hat{\mathbf{n}}
    $$


    Parameters
    ----------
    coord_list : `numpy.ndarray`
        Array of shape `(n_atoms, 3)` containing the 3D coordinates of
        atoms forming a ring.

    Returns
    -------
    `numpy.ndarray`
        Array of overlap values between each atom-centred vector and
        the ring normal vector.
    """
    geometric_centre: numpy.ndarray = numpy.mean(coord_list, axis=0)

    vectors: numpy.ndarray = geometric_centre - coord_list
    vectors /= numpy.linalg.norm(vectors, axis=1, keepdims=True)

    normal_vector: numpy.ndarray = numpy.cross(vectors[0], vectors[1])
    normal_vector /= numpy.linalg.norm(normal_vector)

    return numpy.dot(vectors, normal_vector)
