"""Detect and process pharmacophore using RDKit."""

# =======================
# Python internal package
# =======================
# C
import collections.abc

# D
import dataclasses

# P
import pathlib

# T
import tempfile
import textwrap
import typing

# ================
# External package
# ================
# N
import numpy

# R
import rdkit.Chem as Chem
import rdkit.Chem.rdMolChemicalFeatures as rdMolFeature
import rdkit.Geometry as Geometry
import rdkit.RDConfig as RDConfig

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

# P
from .parser.molecule import (
    Molecule,
)

# U
from .utils.generic import (
    DataClass,
    KeyIdentifier,
    PharmacophoreFamily,
)


@dataclasses.dataclass
class Pharmacophore:
    """Represent a pharmacophore point."""

    # Mandatory fields.
    key: KeyIdentifier
    id: int
    coordinate_x: float
    coordinate_y: float
    coordinate_z: float
    family: str
    subfamily: str

    # Optional fields.
    atom_name: str | None = None
    atom_type: str | None = None
    atom_id: str | None = None
    residu_name: str | None = None
    residu_id: str | None = None
    chain: str | None = None

    # Linking to the MDanalysis universe.
    mda_id: int | None = dataclasses.field(default=None, repr=False)
    # Keeping rdkit atom index.
    rdkit_atom_id: list[int] | None = dataclasses.field(
        default_factory=list, repr=False
    )


class PharmacophoreFactory:
    """A factory class for extracting and organizing pharmacophoric features
    from molecules.
    """

    id: int = 0

    def __init__(
        self,
        miscellaneous: DataClass,
        feature_file: str | None = None,
        redefinition: dict[str, str | None | PharmacophoreFamily]
        | None = None,
    ) -> None:
        """Initializes a pharmacophore object.

        Parameters
        ----------
        miscellaneous : `DataClass`
            Configuration dataclass controlling feature activation.

        feature_file : `str | None`, optional
            Path to a custom RDKit feature definition file (`.fdef`). If
            `None`, the default RDKit feature definition is used and extended
            internally. By default, `None`.

        redefinition : `dict[str, str | None | PharmacophoreFamily] | None`,
                        optional
            Mapping from RDKit feature family names to internal
            `PharmacophoreFamily` values. A value of `None` inside the
            dictionary disables the corresponding feature family. If `None` is
            provided, instead of a dictionary, a default mapping is applied. By
            default, `None`.

        Raises
        ------
        `ValueError`
            - The given `fdef` file does not exist.
            - The given `fdef` file is not a file.
            - The file has not a `fdef` extension.
        """
        self.__miscellaneous: DataClass = miscellaneous

        self.redefinition: (
            dict[str, str | None | PharmacophoreFamily] | None
        ) = redefinition

        if self.redefinition is None:
            self.redefinition = {
                "Donor": PharmacophoreFamily.DONOR,
                "Acceptor": PharmacophoreFamily.ACCEPTOR,
                "NegIonizable": PharmacophoreFamily.NEGATIVE,
                "PosIonizable": PharmacophoreFamily.POSITIVE,
                "ZnBinder": None,
                "Aromatic": PharmacophoreFamily.AROMATIC,
                "Hydrophobe": PharmacophoreFamily.HYDROPHOBE,
                "LumpedHydrophobe": PharmacophoreFamily.LUMPED_HYDROPHOBE,
                "HalogenDonor": PharmacophoreFamily.HALOGEN_DONOR,
                "HalogenAcceptor": PharmacophoreFamily.HALOGEN_ACCEPTOR,
            }

        if feature_file is None:
            self.feature_factory: rdMolFeature.MolChemicalFeatureFactory = (
                self.__set_default_fdef()
            )
        else:
            self.feature_factory = rdMolFeature.BuildFeatureFactory(
                feature_file
            )

    def __set_default_fdef(self) -> rdMolFeature.MolChemicalFeatureFactory:
        """Create an extended default RDKit feature definition.

        Returns
        -------
        `rdMolFeature.MolChemicalFeatureFactory`
            A feature factory built from the extended feature definition.
        """
        fdef_content: str = ""

        with open(
            pathlib.Path(RDConfig.RDDataDir) / "BaseFeatures.fdef",
            "r",
            encoding="utf-8",
        ) as file:
            fdef_content = file.read()

        # `.fdef` part for detecting and assigning Halogen Pharmacophore
        # Features.
        fdef_content += textwrap.dedent("""\
        AtomType HalogenDonor [Cl,Br,I]

        DefineFeature SingleAtomDonor [{HalogenDonor}]
          Family HalogenDonor
          Weights 1.0
        EndFeature

        AtomType HalogenPlipAcceptor [O,N,S;D1;$(O-[#6,#7,#15,#16]),$(N-[#6,#7,#15,#16]),$(S-[#6,#7,#15,#16])]

        DefineFeature SingleAtomAcceptor [{Hydroxyl},{ChalcAcceptor},{NAcceptor},{HalogenAcceptor},{HalogenPlipAcceptor}]
          Family HalogenAcceptor
          Weights 1.0
        EndFeature
        """)

        with tempfile.NamedTemporaryFile(
            "w", suffix=".fdef", delete=False
        ) as file:
            _ = file.write(fdef_content)

        feature_factory: rdMolFeature.MolChemicalFeatureFactory = (
            rdMolFeature.BuildFeatureFactory(file.name)
        )

        return feature_factory

    def get_collection(self, molecule: Molecule) -> dict[str, list]:
        """Extracts pharmacophore features from a molecule and organizes them
        into a dictionnary with iterators.

        Parameters
        ----------
        molecule : `Molecule`
            The molecule object to extract features from.

        Returns
        -------
        `dict[str, list]`
            A dictionary mapping feature categories to their corresponding
            iterators.

        Raises
        ------
        `KeyError`
            - If a feature family found in the molecule is not defined in the
            redefinition mapping.

        `ValueError`
            - If an unsupported property category is encountered.
        """
        collection: dict = {}
        parsed_hydrophobe: list[int] = []

        Info.COMPUTE_PHARMACOPHORE(status="in progress…")
        iterator: collections.abc.Iterable = reversed(
            self.feature_factory.GetFeaturesForMol(molecule.rdkit)
        )
        Info.COMPUTE_PHARMACOPHORE(status="done!")

        for feature in iterator:
            family: str = feature.GetFamily()

            try:
                category: str | None = self.redefinition[family]
            except KeyError:
                Error.FAMILY_NOT_REDEFINE(
                    ValueError, family=feature.GetFamily()
                )

            if category is None:
                continue

            if not getattr(self.__miscellaneous, f"enable_{category}"):
                continue

            data: dict = {
                "feature": feature,
                "molecule": molecule,
                "id": self.id,
                "category": category,
                "family": family,
                "parsed_hydrophobe": parsed_hydrophobe,
            }

            self.id += 1

            coordinate: typing.Callable = PharmacophoreFactory.coordinate

            match category:
                case (
                    PharmacophoreFamily.ACCEPTOR
                    | PharmacophoreFamily.HALOGEN_ACCEPTOR
                    | PharmacophoreFamily.HALOGEN_DONOR
                ):
                    decorator = DecoratorNeighborsAtom(coordinate)
                    decorator = DecoratorCoordinateAtom(coordinate, decorator)
                case (
                    PharmacophoreFamily.HYDROPHOBE
                    | PharmacophoreFamily.LUMPED_HYDROPHOBE
                    | PharmacophoreFamily.NEGATIVE
                    | PharmacophoreFamily.POSITIVE
                ):
                    decorator = DecoratorCoordinateAtom(coordinate)
                    decorator = DecoratorCoordinateFeature(
                        coordinate, decorator
                    )
                case PharmacophoreFamily.DONOR:
                    decorator = DecoratorHydrogenDonor(coordinate)
                    decorator = DecoratorCoordinateAtom(coordinate, decorator)
                case PharmacophoreFamily.AROMATIC:
                    decorator = DecoratorCoordinateAtom(coordinate)
                    decorator = DecoratorNormalVector(coordinate, decorator)
                    decorator = DecoratorCoordinateFeature(
                        coordinate, decorator
                    )
                case _:
                    Error.WRONG_CATEGORY(ValueError, category=category)

            # Special case redefinition. Lumped hydrophobe are hydrophobe. But,
            # we want to filter out the one define in lumped hydrophobe outside
            # the classical hydrophobe. We do that mostly due to PDB files.
            if category == PharmacophoreFamily.LUMPED_HYDROPHOBE:
                category = PharmacophoreFamily.HYDROPHOBE

            if category not in collection:
                collection[category] = []

            decorator.operation(collection[category], data)

        return collection

    @staticmethod
    def coordinate(point: Geometry.Point3D) -> dict[str, float]:
        """Converts a `Geometry.Point3D` object to a dictionary.

        Parameters
        ----------
        point : `Geometry.Point3D`
            The RDKit 3D point to convert.

        Returns
        -------
        `dict[str, float]`
            A dictionnary with coordinates.
        """
        return {
            "coordinate_x": point.x,
            "coordinate_y": point.y,
            "coordinate_z": point.z,
        }


# pylint: disable=too-few-public-methods
# Component class, so without many methods.


class DecoratorNeighborsAtom:
    """A decorator class that create `Pharmacophore()` bases on atoms
    neighbors.
    """

    def __init__(
        self,
        coordinate: typing.Callable,
        component: typing.Callable | None = None,
    ) -> None:
        """Initializes a decorator.

        Parameters
        ----------
        coordinate : `typing.Callable`
            A function that processes extract coordinates.

        component : `typing.Callable | None`, optional
            A subsequent component to execute in the chain, if provided.
        """
        self.__component: typing.Callable | None = component
        self.__cordinate: typing.Callable = coordinate

    def operation(self, iterator: list, data: dict) -> None:
        """Processes atomic neightbor coordinates and assigns them to a
        pharmacophore iterator.

        **The decorator:**
        - Marks the previously added pharmacophore as a `PHARMACOPHORE`.
        - Identifies the first neighboring atom of the feature atom.
        - Extracts coordinates and atomic metadata.
        - Appends a new `ATOM` pharmacophore entry.
        - Propagates execution to the next decorator if present.

        Parameters
        ----------
        iterator : `list`
            The iterator to which pharmacophoric features will be added.

        data : `dict`
            A dictionary containing molecular feature information, including:
            - RDKit feature and molecule (`Chem.Mol()`).
            - MDAnalysis molecule.
            - Feature category and family.
        """
        position: typing.Callable = (
            data["molecule"].rdkit.GetConformer().GetAtomPosition
        )

        iterator[-1].key = KeyIdentifier.PHARMACOPHORE

        index: int = data["feature"].GetAtomIds()[0]
        atom: Chem.Atom = data["molecule"].rdkit.GetAtomWithIdx(index)

        neighbor_atom: Chem.Atom = atom.GetNeighbors()[0]
        neighbor_index: int = neighbor_atom.GetIdx()

        coordinate: dict = self.__cordinate(position(neighbor_index))

        extra_data: dict = {
            "atom_name": neighbor_atom.GetSymbol(),
            "atom_type": neighbor_atom.GetSymbol(),
            "atom_id": neighbor_index,
            "rdkit_atom_id": [neighbor_index],
        }

        if data["molecule"].mda is not None:
            mda_id: int = data["molecule"].get_index(
                numpy.fromiter(coordinate.values(), dtype=float)
            )

            extra_data |= data["molecule"].get_data(mda_id)
            extra_data["mda_id"] = mda_id

        pharmacophore: Pharmacophore = Pharmacophore(
            key=KeyIdentifier.ATOM,
            id=data["id"],
            family=data["family"],
            subfamily=data["feature"].GetType(),
            **coordinate,
            **extra_data,
        )

        iterator += [pharmacophore]

        if self.__component is not None:
            self.__component.operation(iterator, data)


class DecoratorCoordinateAtom:
    """A decorator class that create `Pharmacophore()` bases on atoms
    coordinates that are inside a given feature.
    """

    def __init__(
        self,
        coordinate: typing.Callable,
        component: typing.Callable | None = None,
    ) -> None:
        """Initializes a decorator.

        Parameters
        ----------
        coordinate : `typing.Callable`
            A function that processes extract coordinates.

        component : `typing.Callable | None`, optional
            A subsequent component to execute in the chain, if provided.
        """
        self.__component: typing.Callable | None = component
        self.__cordinate: typing.Callable = coordinate

    def operation(self, iterator: list, data: dict) -> None:
        """Processes atomic coordinates and assigns them to a pharmacophore
        iterator.

        **The decorator:**
        - Applies hydrophobe and lumped-hydrophobe filtering rules.
        - Extracts atomic coordinates and metadata.
        - Appends a new `ATOM` pharmacophore to the iterator.
        - Propagates execution to the next decorator if present.

        Parameters
        ----------
        iterator : `list`
            The iterator to which pharmacophoric features will be added.

        data : `dict`
            A dictionary containing molecular feature information, including:
            - RDKit feature and molecule (`Chem.Mol()`).
            - MDAnalysis molecule.
            - Feature category and family.
            - Parsed hydrophobe indice.
        """
        position: typing.Callable = (
            data["molecule"].rdkit.GetConformer().GetAtomPosition
        )

        for index in data["feature"].GetAtomIds():
            # Special case redefinition. Lumped hydrophobe are hydrophobe. But,
            # we want to filter out the one define in lumped hydrophobe outside
            # the classical hydrophobe.
            if data["category"] == PharmacophoreFamily.LUMPED_HYDROPHOBE:
                data["parsed_hydrophobe"] += [index]
            elif data["category"] == PharmacophoreFamily.HYDROPHOBE:
                if index in data["parsed_hydrophobe"]:
                    _ = iterator.pop()
                    continue

            coordinate: dict = self.__cordinate(position(index))

            atom: Chem.Atom = data["molecule"].rdkit.GetAtomWithIdx(index)
            extra_data: dict = {
                "atom_name": atom.GetSymbol(),
                "atom_type": atom.GetSymbol(),
                "atom_id": index,
                "rdkit_atom_id": [index],
            }

            if data["molecule"].mda is not None:
                mda_id: int = data["molecule"].get_index(
                    numpy.fromiter(coordinate.values(), dtype=float)
                )

                extra_data |= data["molecule"].get_data(mda_id)
                extra_data["mda_id"] = mda_id

            pharmacophore: Pharmacophore = Pharmacophore(
                key=KeyIdentifier.ATOM,
                id=data["id"],
                family=data["family"],
                subfamily=data["feature"].GetType(),
                **coordinate,
                **extra_data,
            )

            iterator += [pharmacophore]

            if self.__component is not None:
                data["index"] = index
                self.__component.operation(iterator, data)


class DecoratorCoordinateFeature:
    """A decorator class that create `Pharmacophore()` bases on coordinates
    of a given feature.
    """

    def __init__(
        self,
        coordinate: typing.Callable,
        component: typing.Callable | None = None,
    ) -> None:
        """Initializes a decorator.

        Parameters
        ----------
        coordinate : `typing.Callable`
            A function that processes extract coordinates.

        component : `Optinal[typing.Callable]`, optional
            A subsequent operation to execute in the chain, if provided.
        """
        self.__component: typing.Callable | None = component
        self.__cordinate: typing.Callable = coordinate

    def operation(self, iterator: list, data: dict) -> None:
        """Processes features coordinates and assigns them to a pharmacophore
        iterator.

        **The decorator:**
        - Extracts the feature centroid coordinates.
        - Builds a `PHARMACOPHORE` object.
        - Optionally enriches it with MDAnalysis chain/residue metadata.
        - Appends it to the pharmacophore iterator.
        - Propagates execution to the next decorator if present.

        Parameters
        ----------
        iterator : `list`
            The iterator to which pharmacophoric features will be added.

        data : `dict`
            A dictionary containing molecular feature information, including:
            - RDKit feature and molecule (`Chem.Mol()`).
            - MDAnalysis molecule.
            - Feature category and family.
        """
        extra_data: dict = {
            "rdkit_atom_id": list(data["feature"].GetAtomIds()),
        }

        if data["molecule"].mda is not None:
            # Fetch coordinates.
            position: typing.Callable = (
                data["molecule"].rdkit.GetConformer().GetAtomPosition
            )
            coordinate: dict = self.__cordinate(
                position(data["feature"].GetAtomIds()[0])
            )

            # Fetch MDA data.
            mda_id: int = data["molecule"].get_index(
                numpy.fromiter(coordinate.values(), dtype=float)
            )

            extra_data |= data["molecule"].get_data(mda_id)

            # Getting only chains / residus data.
            del extra_data["atom_name"]
            del extra_data["atom_type"]
            del extra_data["atom_id"]

            # For aromatic feature.
            data["mda_id"] = mda_id

        pharmacophore: Pharmacophore = Pharmacophore(
            key=KeyIdentifier.PHARMACOPHORE,
            id=data["id"],
            family=data["family"],
            subfamily=data["feature"].GetType(),
            **self.__cordinate(data["feature"].GetPos()),
            **extra_data,
        )

        iterator += [pharmacophore]

        if self.__component is not None:
            self.__component.operation(iterator, data)


class DecoratorHydrogenDonor:
    """A decorator class to process hydrogen donor pharmacophoric features."""

    def __init__(self, coordinate: typing.Callable) -> None:
        """Initializes a decorator.

        Parameters
        ----------
        coordinate : `typing.Callable`
            A function that processes extract coordinates.
        """
        self.__cordinate: typing.Callable = coordinate

    def operation(self, iterator: list, data: dict) -> None:
        """Identifies hydrogen donor atoms and associates hydrogen atom
        coordinates with them.

        **The decorator:**
        - Requires an existing pharmacophore in the iterator.
        - Identifies bonded hydrogen atoms.
        - Creates one `ATOM` pharmacophore per hydrogen.
        - Optionally enriches each pharmacophore with MDAnalysis metadata.

        Parameters
        ----------
        iterator : `list`
            The iterator to which pharmacophoric features will be added.

        data : `dict`
            A dictionary containing molecular feature information, including:
            - RDKit feature and molecule (`Chem.Mol()`).
            - MDAnalysis molecule.
            - Feature category and family.

        Raises
        ------
        `IndexError`
            - If the iterator does not contain at least one pharmacophore.
        """
        if len(iterator) == 0:
            Error.HYDROGEN_ITERATOR(IndexError)

        pharmacophore: Pharmacophore = iterator[-1]
        setattr(pharmacophore, "key", KeyIdentifier.PHARMACOPHORE)

        position: typing.Callable = (
            data["molecule"].rdkit.GetConformer().GetAtomPosition
        )

        atom: Chem.Atom = data["molecule"].rdkit.GetAtomWithIdx(data["index"])

        vars_pharmacophore: dict = dataclasses.asdict(pharmacophore)
        parameter: dict = {}

        # Skip the data_class fields with no default value.
        for field_i in dataclasses.fields(pharmacophore):
            if field_i.default is None:
                continue

            parameter[field_i.name] = vars_pharmacophore[field_i.name]

        parameter["key"] = KeyIdentifier.ATOM

        for neighbour in atom.GetNeighbors():
            if neighbour.GetSymbol() != "H":
                continue

            coordinate: dict = self.__cordinate(position(neighbour.GetIdx()))
            parameter |= coordinate

            parameter["atom_name"] = neighbour.GetSymbol()
            parameter["atom_type"] = neighbour.GetSymbol()
            parameter["atom_id"] = neighbour.GetIdx()
            parameter["rdkit_atom_id"] = [neighbour.GetIdx()]

            if data["molecule"].mda is not None:
                mda_id: int = data["molecule"].get_index(
                    numpy.fromiter(coordinate.values(), dtype=float)
                )

                if mda_id is not None:
                    parameter |= data["molecule"].get_data(mda_id)
                    parameter["mda_id"] = mda_id

            iterator += [Pharmacophore(**parameter)]


class DecoratorNormalVector:
    """A decorator class to compute norm from a aromatic cycle."""

    def __init__(
        self,
        coordinate: typing.Callable,
        component: typing.Callable | None = None,
    ) -> None:
        """Initializes a decorator.

        Parameters
        ----------
        coordinate : `typing.Callable`
            A function that processes extract coordinates.

        component : `typing.Callable | None`, optional
            A subsequent component to execute in the chain, if provided.
        """
        self.__component: typing.Callable | None = component
        self.__cordinate: typing.Callable = coordinate

    def operation(self, iterator: list, data: dict) -> None:
        """Compute a normal vector from the given aromatic cycle feature.

        **The decorator:**
        - Extracts atom coordinates from an aromatic feature.
        - Compute a plane normal.
        - Normalizes the resulting vector.
        - Appends a `NORM` pharmacophore to the iterator.
        - Optionally propagates MDAnalysis metadata.

        Parameters
        ----------
        iterator : `list`
            The iterator to which pharmacophoric features will be added.

        data : `dict`
            A dictionary containing molecular feature information, including:
            - RDKit feature and molecule (`Chem.Mol()`).
            - MDAnalysis molecule.
            - Feature category and family.
        """
        position: typing.Callable = (
            data["molecule"].rdkit.GetConformer().GetAtomPosition
        )

        coord_list: list = []

        for index in data["feature"].GetAtomIds():
            coordinate: numpy.ndarray = numpy.fromiter(
                self.__cordinate(position(index)).values(), dtype=float
            )
            coord_list += [coordinate]

            if len(coord_list) >= 3:
                break

        centre: numpy.ndarray = numpy.array(
            [
                iterator[-1].coordinate_x,
                iterator[-1].coordinate_y,
                iterator[-1].coordinate_z,
            ]
        )

        normal_vector: numpy.ndarray = self.__normal(centre, *coord_list[:2])
        extra_data: dict = {}

        if "mda_id" in data:
            extra_data = data["molecule"].get_data(data["mda_id"])

            # Getting only chains / residus data.
            del extra_data["atom_name"]
            del extra_data["atom_type"]
            del extra_data["atom_id"]

        pharmacophore: Pharmacophore = Pharmacophore(
            key=KeyIdentifier.NORM,
            id=data["id"],
            family=data["family"],
            subfamily=data["feature"].GetType(),
            coordinate_x=normal_vector[0],
            coordinate_y=normal_vector[1],
            coordinate_z=normal_vector[2],
            **extra_data,
        )

        iterator += [pharmacophore]

        if self.__component is not None:
            self.__component.operation(iterator, data)

    def __normal(
        self,
        centre: numpy.ndarray,
        edge_1: numpy.ndarray,
        edge_2: numpy.ndarray,
    ) -> numpy.ndarray:
        """Compute a normal vector from a list of three set of coordinates.

        Parameters
        ----------
        centre : `numpy.ndarray`
            Coordinates of the aromatic feature centre.

        edge_1 : `numpy.ndarray`
            Coordinates of the first atom defining the aromatic plane.

        edge_2 : `numpy.ndarray`
            Coordinates of the second atom defining the aromatic plane.

        Returns
        -------
        `numpy.ndarray`
            A normalized vector orthogonal to the aromatic ring plane.
        """
        normal: numpy.ndarray = numpy.cross(edge_1 - centre, edge_2 - centre)
        normal /= numpy.linalg.norm(normal)

        return normal
