"""This module contains the `DamageDB` class."""

import warnings
from collections import deque
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np
from numpy import typing as npt
from numpy.lib.recfunctions import structured_to_unstructured as str2unstr
from scipy.spatial.transform import Rotation
from sklearn.decomposition import PCA

from irradiapy import dtypes, materials
from irradiapy.io.lammpsreader import LAMMPSReader


@dataclass
class DamageDB:
    """Class used to reconstruct the damage produced by a PKA from a database of MD debris.

    Attributes
    ----------
    dir_mddb : Path
        Directory of the MD debris database.
    compute_tdam : bool
        Whether to apply Lindhard's formula to the recoil energy. It should be `True` for
        MD simulations without electronic stopping.
    mat_pka : materials.Material
        PKA material.
    mat_target : materials.Material
        Target material.
    dpa_mode : materials.Material.DpaMode
        Mode for dpa calculation.
    tdam_mode : materials.Material.TdamMode
        Mode for PKA to damage energy calculation.
    seed : int, optional (default=0)
        Random seed for random number generator.
    """

    dir_mddb: Path
    compute_tdam: bool
    mat_pka: "materials.Material"
    mat_target: "materials.Material"
    dpa_mode: "materials.Material.DpaMode"
    tdam_mode: "materials.Material.TdamMode"
    seed: int = 0
    __rng: np.random.Generator = field(init=False)
    __calc_nd: Callable[[float], float] = field(init=False)
    __files: dict[float, list[Path]] = field(init=False)
    __energies: npt.NDArray[np.float64] = field(init=False)
    __nenergies: int = field(init=False)

    def __post_init__(self) -> None:
        self.__rng = np.random.default_rng(self.seed)
        # Scan the database
        self.__files = {
            float(folder.name): list(folder.iterdir())
            for folder in self.dir_mddb.iterdir()
            if folder.is_dir()
        }
        self.__energies = np.array(sorted(self.__files.keys(), reverse=True))
        self.__nenergies = len(self.__energies)
        # PKA energy to damage energy conversion
        self.__compute_damage_energy = lambda x: self.mat_pka.epka_to_tdam(
            mat_pka=self.mat_target, epka=x, mode=self.tdam_mode
        )
        # Select the dpa model for residual energy
        self.__calc_nd = lambda x: self.mat_target.tdam_to_dpa(
            tdam=x, mode=self.dpa_mode
        )

    def __get_files(self, pka_e: float) -> tuple[dict[float, list[Path]], int]:
        """Get cascade files and number of residual FP for a given PKA energy.

        Parameters
        ----------
        pka_e : float
            PKA energy.

        Returns
        -------
        tuple[dict[float, list[Path]], int]
            Dictionary of selected paths and number of residual FP.
        """
        # Decompose the PKA energy into cascades and residual energy
        residual_energy = (
            self.__compute_damage_energy(pka_e) if self.compute_tdam else pka_e
        )
        cascade_counts = np.zeros(self.__nenergies, dtype=np.int64)
        for i, energy in enumerate(self.__energies):
            cascade_counts[i], residual_energy = divmod(residual_energy, energy)
        # Select the files for each energy
        if residual_energy > 0:
            residual_energy = self.__compute_damage_energy(residual_energy)
        debris_files = {
            energy: self.__rng.choice(self.__files[energy], cascade_counts[i])
            for i, energy in enumerate(self.__energies)
        }
        # Get the number of residual FP
        nfp = np.round(self.__calc_nd(residual_energy)).astype(np.int64)
        return debris_files, nfp

    def get_pka_debris(
        self,
        pka_e: float,
        pka_pos: npt.NDArray[np.float64],
        pka_dir: npt.NDArray[np.float64],
    ) -> dtypes.Defect:
        """Get PKA debris from its energy position, and direction.

        Parameters
        ----------
        pka_e : float
            PKA energy.
        pka_pos : npt.NDArray[np.float64]
            PKA position.
        pka_dir : npt.NDArray[np.float64]
            PKA direction.

        Returns
        -------

        dtypes.Defect
            Defects after the cascades.
        """
        files, nfp = self.__get_files(pka_e)
        # Get the maximum energy available in the database for the given PKA.
        # If no energy is available, return zero to place only FP.
        db_emax = next(
            (energy for energy in self.__energies if len(files[energy])), 0.0
        )
        # Possible to get cascades from the database
        if db_emax > 0.0:
            defects = self.__process_highest_energy_cascade(
                files, db_emax, pka_pos, pka_dir
            )
            parallelepiped = self.__get_parallelepiped(defects)
            defects = self.__place_other_debris(files, defects, parallelepiped)
            if nfp:
                defects = self.__place_fps_in_parallelepiped(
                    defects, nfp, parallelepiped
                )
            return defects
        # If no energy is available, generate FP only
        if nfp:
            defects = self.__place_fps_in_sphere(nfp, pka_pos, pka_dir)
            return defects
        defects = np.empty(0, dtype=dtypes.defect)
        return defects

    def __process_highest_energy_cascade(
        self,
        files: dict,
        db_emax: float,
        pka_pos: npt.NDArray[np.float64],
        pka_dir: npt.NDArray[np.float64],
    ) -> dtypes.Defect:
        """Process the highest energy cascade.

        Parameters
        ----------
        files : dict
            Dictionary of files for each energy.
        db_emax : float
            Energy of the highest energy cascade.
        pka_pos : npt.NDArray[np.float64]
            PKA position.
        pka_dir : npt.NDArray[np.float64]
            PKA direction.

        Returns
        -------
        dtypes.Defect
            Defects after the highest energy cascade.
        """
        file = files[db_emax][0]
        files[db_emax] = np.delete(files[db_emax], 0)
        defects = deque(LAMMPSReader(file), maxlen=1).pop()["atoms"]
        xaxis = np.array([1.0, 0.0, 0.0])
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=UserWarning)
            transform = Rotation.align_vectors([pka_dir], [xaxis])[0]

        pos = str2unstr(defects[["x", "y", "z"]], dtype=np.float64, copy=False)
        pos = transform.apply(pos) + pka_pos
        defects["x"] = pos[:, 0]
        defects["y"] = pos[:, 1]
        defects["z"] = pos[:, 2]

        return defects

    def __place_other_debris(
        self,
        files: dict,
        defects: dtypes.Defect,
        parallelepiped: tuple[PCA, npt.NDArray, npt.NDArray],
    ) -> dtypes.Defect:
        """Place other debris in the parallelepiped.

        Parameters
        ----------
        files : dict
            Dictionary of files for each energy.
        defects : dtypes.Defect
            Defects after the highest energy cascade.
        parallelepiped : tuple[PCA, npt.NDArray, npt.NDArray]
            Parallelepiped definition.

        Returns
        -------
        dtypes.Defect
            Defects after placing the other debris.
        """
        for energy in self.__energies:
            for file0 in files[energy]:
                defects_ = deque(LAMMPSReader(file0), maxlen=1).pop()["atoms"]

                transform = Rotation.random(rng=self.__rng)
                pos = str2unstr(defects_[["x", "y", "z"]], dtype=np.float64, copy=False)
                pos0 = self.__get_parallelepiped_points(*parallelepiped, 1)
                pos = transform.apply(pos) + pos0
                defects_["x"] = pos[:, 0]
                defects_["y"] = pos[:, 1]
                defects_["z"] = pos[:, 2]

                defects = np.concatenate((defects, defects_))
        return defects

    def __place_fps_in_parallelepiped(
        self,
        defects: dtypes.Defect,
        nfp: int,
        parallelepiped: tuple[PCA, npt.NDArray, npt.NDArray],
    ) -> dtypes.Defect:
        """Place FPs anywhere in the parallelepiped.

        Parameters
        ----------
        defects : dtypes.Defect
            Defects after placing the other debris.
        nfp : int
            Number of FPs.
        parallelepiped : tuple[PCA, npt.NDArray, npt.NDArray]
            Parallelepiped definition.

        Returns
        -------
        dtypes.Defect
            Defects after placing the FPs.
        """
        defects_ = np.zeros(2 * nfp, dtype=dtypes.defect)
        defects_["type"][:nfp] = self.mat_target.atomic_number
        defects_["x"][:nfp] = self.mat_target.dist_fp / 2.0
        pos = str2unstr(defects_[["x", "y", "z"]], dtype=np.float64, copy=False)
        pos[:nfp] = Rotation.random(nfp, rng=self.__rng).apply(pos[:nfp])
        pos[nfp:] = -pos[:nfp]
        pos0 = self.__get_parallelepiped_points(*parallelepiped, nfp)
        pos[:nfp] += pos0
        pos[nfp:] += pos0
        defects_["x"] = pos[:, 0]
        defects_["y"] = pos[:, 1]
        defects_["z"] = pos[:, 2]

        return np.concatenate((defects, defects_))

    def __place_fps_in_sphere(
        self,
        nfp: int,
        pka_pos: npt.NDArray[np.float64],
        pka_dir: npt.NDArray[np.float64],
    ) -> dtypes.Defect:
        """Generate FPs in a sphere.

        Parameters
        ----------
        nfp : int
            Number of FPs.
        pka_pos : npt.NDArray[np.float64]
            PKA position.
        pka_dir : npt.NDArray[np.float64]
            PKA direction.

        Returns
        -------
        dtypes.Defect
            Defects after generating.
        """
        defects_ = np.zeros(2 * nfp, dtype=dtypes.defect)
        defects_["type"][:nfp] = self.mat_target.atomic_number
        defects_["x"][:nfp] = self.mat_target.dist_fp / 2
        pos = str2unstr(defects_[["x", "y", "z"]], dtype=np.float64, copy=False)
        pos[:nfp] = Rotation.random(nfp, rng=self.__rng).apply(pos[:nfp])
        pos[nfp:] = -pos[:nfp]

        random = self.__rng.random((nfp, 3))
        theta = np.arccos(2.0 * random[:, 0] - 1.0)
        phi = 2.0 * np.pi * random[:, 1]
        radius = nfp * self.mat_target.dist_fp / 2.0
        r = radius * np.cbrt(random[:, 2])
        points = np.empty((nfp, 3))
        points[:, 0] = r * np.sin(theta) * np.cos(phi)
        points[:, 1] = r * np.sin(theta) * np.sin(phi)
        points[:, 2] = r * np.cos(theta)

        pos[:nfp] += points
        pos[nfp:] += points
        pos += pka_pos + pka_dir * radius
        defects_["x"] = pos[:, 0]
        defects_["y"] = pos[:, 1]
        defects_["z"] = pos[:, 2]

        return defects_

    def __get_parallelepiped(self, atoms: dtypes.Atom) -> tuple:
        """
        Define a parallelepiped from the atomic positions using PCA.

        Parameters
        ----------
        atoms : dtypes.Atom
            Atomic positions.

        Returns
        -------
        tuple
            PCA object, minimum PCA coordinates, maximum PCA coordinates.
        """
        pos = str2unstr(atoms[["x", "y", "z"]], dtype=np.float64, copy=False)
        pca = PCA(n_components=3)
        pca.fit(pos)
        atoms_pca = pca.transform(pos)
        min_pca = np.min(atoms_pca, axis=0)
        max_pca = np.max(atoms_pca, axis=0)
        return pca, min_pca, max_pca

    def __get_parallelepiped_points(
        self,
        pca: PCA,
        min_pca: npt.NDArray[np.float64],
        max_pca: npt.NDArray[np.float64],
        npoints: int,
    ) -> npt.NDArray[np.float64]:
        """
        Generate random points within a parallelepiped.

        Parameters
        ----------
        pca : PCA
            PCA object.
        min_pca : npt.NDArray[np.float64]
            Minimum PCA coordinates.
        max_pca : npt.NDArray[np.float64]
            Maximum PCA coordinates.
        npoints : int
            Number of points to generate.

        Returns
        -------
        npt.NDArray[np.float64]
            Random points within the parallelepiped.
        """
        random_points_pca = self.__rng.uniform(min_pca, max_pca, size=(npoints, 3))
        return pca.inverse_transform(random_points_pca)
