"""Add a way to visualize using PyMOL."""

# =======================
# Python internal package
# =======================
# D
import dataclasses

# ==============
# Module package
# ==============
# U
from ..utils.generic import (
    DataClass,
    InteractionType,
    KeyIdentifier,
    PharmacophoreFamily,
)


class Pymol:
    """Generate a PyMOL (.pml) visualization script from pharmacophore and
    interaction data using CGOs.
    """

    def __init__(
        self,
        path: str,
        pharmacophore_collection_a: dict[PharmacophoreFamily, list[DataClass]],
        pharmacophore_collection_b: dict[PharmacophoreFamily, list[DataClass]],
        interaction_collection: dict[InteractionType, list[DataClass]],
        drawning_configuration: DataClass,
    ) -> None:
        """Build and export a PyMOL visualization script (.pml).

        Parameters
        ----------
        path : `str`
            Output file path where the .pml script will be written.

        pharmacophore_collection_a : `dict[PharmacophoreFamily, list[DataClass]]`
            First set of pharmacophore features.

        pharmacophore_collection_b : `dict[PharmacophoreFamily, list[DataClass]]`
            Second set of pharmacophore features.

        interaction_collection : `dict[InteractionType, list[DataClass]]`
            Molecular interactions.

        drawning_configuration : `DataClass`
            Configuration dataclass defining colors and radii.
        """
        self.__drawning_configuration: DataClass = drawning_configuration
        self.__script_content: list[str] = []

        # Initialize the PML file with necessary imports
        self.__script_content.append("python")
        self.__script_content.append("from pymol import cgo, cmd")
        self.__script_content.append("")

        # 1. Process Pharmacophores
        key_list: set = set(
            list(pharmacophore_collection_a.keys())
            + list(pharmacophore_collection_b.keys())
        )

        for key in key_list:
            pharmacophore: list = []

            if key in pharmacophore_collection_a:
                pharmacophore += pharmacophore_collection_a[key]
            if key in pharmacophore_collection_b:
                pharmacophore += pharmacophore_collection_b[key]

            self.__draw_pharmacophore(pharmacophore, key)

        # 2. Process Interactions
        for key, interaction in interaction_collection.items():
            self.__draw_interaction(interaction, key)

        # Close the python block in the PML file
        self.__script_content.append("python end")

        # Write to file
        with open(path, "w", encoding="utf-8") as file:
            file.write("\n".join(self.__script_content))

    def __fetch_fields(self, key: str) -> dict:
        """Extract visualization parameters associated with a given key.

        Parameters
        ----------
        key : `str`
            Identifier prefix used to match configuration fields.

        Returns
        -------
        `dict`
            Dictionary containing 'color' (as hex) and 'radius'.
        """
        fetched_field: dict = {}

        for field_i in dataclasses.fields(self.__drawning_configuration):
            if not field_i.name.startswith(f"{key}__"):
                continue

            property_name: str = field_i.name.split("__")[-1]
            fetched_field[property_name] = getattr(
                self.__drawning_configuration, field_i.name
            )

        return fetched_field

    def __hex_to_rgb(self, hexadecimal_color: str) -> tuple[float, ...]:
        """Convert a hexadecimal color string to a RGB tuple.

        Parameters
        ----------
        hexadecimal_color : `str`
            Color code in hexadecimal format.

        Returns
        -------
        `tuple[float]`
            Color code in RGB tuple format.
        """
        rgb_color: list[float] = []

        for index in range(1, 7, 2):
            rgb_color.append(
                int(hexadecimal_color[index : index + 2], 16) / 255
            )

        return tuple(rgb_color)

    def __draw_pharmacophore(
        self,
        collection: list[DataClass],
        key: PharmacophoreFamily,
    ) -> None:
        """Generate CGO spheres for pharmacophores.

        Parameters
        ----------
        collection : `list[DataClass]`
            List of pharmacophore dataclass instances to draw.

        key : `PharmacophoreFamily`
            Pharmacophore family associated with the collection.
        """
        field: dict = self.__fetch_fields(str(key))
        color: tuple[float, ...] = self.__hex_to_rgb(field["color"])

        # Start building the CGO list string.
        cgo_name: str = f"pharmacophore_{key}"
        cgo_line: list[str] = [f"# {cgo_name}"]
        cgo_line.append(f"{cgo_name}_obj = [")

        # Set color once for the whole group to save space.
        cgo_line.append(
            f"    cgo.COLOR, {color[0]:.3f}, {color[1]:.3f}, {color[2]:.3f},"
        )

        for pharmacophore in collection:
            if pharmacophore.key != KeyIdentifier.PHARMACOPHORE:
                continue

            cgo_line.append(
                "    cgo.SPHERE, "
                + f"{pharmacophore.coordinate_x:.3f}, "
                + f"{pharmacophore.coordinate_y:.3f}, "
                + f"{pharmacophore.coordinate_z:.3f}, "
                + f"{field['radius']},"
            )

        cgo_line.append("]")
        cgo_line.append(f"cmd.load_cgo({cgo_name}_obj, '{cgo_name}')")
        cgo_line.append("")

        self.__script_content.extend(cgo_line)

    def __draw_interaction(
        self,
        collection: list[DataClass],
        key: InteractionType,
    ) -> None:
        """Generate CGO cylinders for interactions.

        Parameters
        ----------
        collection : `list[DataClass]`
            List of interaction dataclass instances to draw.

        key : `InteractionType`
            Interaction type associated with the collection.
        """
        field: dict = self.__fetch_fields(str(key))
        color: tuple[float, ...] = self.__hex_to_rgb(field["color"])

        cgo_name = f"interaction_{key}"
        cgo_line = [f"# {cgo_name}"]
        cgo_line.append(f"{cgo_name}_obj = [")

        for interaction in collection:
            cgo_line.append(
                "    cgo.CYLINDER, "
                + f"{interaction.molecule_a.coordinate_x:.3f}, "
                + f"{interaction.molecule_a.coordinate_y:.3f}, "
                + f"{interaction.molecule_a.coordinate_z:.3f}, "
                + f"{interaction.molecule_b.coordinate_x:.3f}, "
                + f"{interaction.molecule_b.coordinate_y:.3f}, "
                + f"{interaction.molecule_b.coordinate_z:.3f}, "
                + f"{field['radius']}, "
                + f"{color[0]:.3f}, {color[1]:.3f}, {color[2]:.3f}, "
                + f"{color[0]:.3f}, {color[1]:.3f}, {color[2]:.3f}, "
            )

        cgo_line.append("]")
        cgo_line.append(f"cmd.load_cgo({cgo_name}_obj, '{cgo_name}')")
        cgo_line.append("")

        self.__script_content.extend(cgo_line)
