"""This module provides a class to find and analyze defects in crystalline structures."""

from collections import defaultdict
from dataclasses import dataclass, field

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

from irradiapy import dtypes


@dataclass
class DefectsIdentifier:
    """Class to identify  defects in crystalline structures.

    This class provides methods to identify point defects (vacancies and interstitials) in a
    body-centered cubic (bcc) lattice based on atomic positions from simulation data. It supports
    optional rescaling and recentering of defect positions, and can align the system with a
    specified primary knock-on atom (PKA) direction.

    Parameters
    ----------
    lattice : str
        Lattice type. Currently only "bcc" is supported.
    a0 : float
        Lattice parameter, in angstroms.
    debug : bool, optional (default=False)
        If True, enables debug mode for additional output.
    """

    lattice: str
    a0: float
    debug: bool = False
    __perx: bool = field(init=False)
    __pery: bool = field(init=False)
    __perz: bool = field(init=False)
    __nxhi: int = field(init=False)
    __nyhi: int = field(init=False)
    __nzhi: int = field(init=False)
    __nxlo: int = field(init=False)
    __nylo: int = field(init=False)
    __nzlo: int = field(init=False)
    __nx: int = field(init=False)
    __ny: int = field(init=False)
    __nz: int = field(init=False)
    __sub_count: int = 2  # Number of atoms per primitive unit cell (bcc).

    def __post_init__(self) -> None:
        if self.lattice == "bcc":
            self.__sub_count = 2
            self.__pos_ucell = np.array(
                [
                    [0.0, 0.0, 0.0],
                    [1.0, 0.0, 0.0],
                    [1.0, 1.0, 0.0],
                    [0.0, 1.0, 0.0],
                    [0.0, 0.0, 1.0],
                    [1.0, 0.0, 1.0],
                    [1.0, 1.0, 1.0],
                    [0.0, 1.0, 1.0],
                    [0.5, 0.5, 0.5],
                ],
                dtype=np.float64,
            )
            self.__idx_ucell = np.array(
                [
                    [0, 0, 0, 0],
                    [1, 0, 0, 0],
                    [1, 1, 0, 0],
                    [0, 1, 0, 0],
                    [0, 0, 1, 0],
                    [1, 0, 1, 0],
                    [1, 1, 1, 0],
                    [0, 1, 1, 0],
                    [0, 0, 0, 1],
                ],
                dtype=np.int64,
            )
            # Scale positions by a0
            self.__pos_ucell *= self.a0
        else:
            raise ValueError("Only bcc lattice is supported.")

    def __rescale_translate_rotate(
        self,
        atoms: dtypes.Atom,
        a1: float | np.number,
        pos_pka: npt.NDArray[np.float64],
        theta_pka: float,
        phi_pka: float,
    ) -> None:
        """Rescales, translates, and rotates the positions of atoms.

        Parameters
        ----------
        atoms : dtypes.Atom
            Structured array of atomic positions (fields: x, y, z).
        a1 : float
            Final lattice parameter to rescale positions.
        pos_pka : npt.NDArray[np.float64]
            Position vector of the PKA for translation.
        theta_pka : float
            Polar angle (in radians) for the PKA direction (for rotation).
        phi_pka : float
            Azimuthal angle (in radians) for the PKA direction (for rotation).
        """
        # Translate
        atoms["x"] -= pos_pka[0]
        atoms["y"] -= pos_pka[1]
        atoms["z"] -= pos_pka[2]

        # Rotation matrix, align with PKA initial direction
        xaxis = np.array([1.0, 0.0, 0.0])
        pka_dir = np.array(
            [
                np.sin(theta_pka) * np.cos(phi_pka),
                np.sin(theta_pka) * np.sin(phi_pka),
                np.cos(theta_pka),
            ]
        )
        transform = R.align_vectors([xaxis], [pka_dir])[0].as_matrix()

        # Scaling matrix
        scaling = a1 / self.a0
        scaling_matrix = (
            np.diag([scaling] * 3) if isinstance(scaling, float) else np.diag(scaling)
        )

        # Combined scale + rotation
        transform = transform @ scaling_matrix

        # Apply transformations
        pos = str2unstr(atoms[["x", "y", "z"]])
        pos = transform.apply(pos)
        atoms["x"] = pos[:, 0]
        atoms["y"] = pos[:, 1]
        atoms["z"] = pos[:, 2]

    def __apply_boundary_conditions(self, idx_atoms: npt.NDArray[np.int64]) -> None:
        """Applies periodic boundary conditions to atomic index coordinates.

        Only one unit cell around the simulation box is considered for periodic boundary conditions.

        Parameters
        ----------
        idx_atoms : npt.NDArray[np.int64]
            Array of atomic positions with index coordinates (shape: [N, 4]).
        """
        if self.__perx:
            idx_atoms[:, 0][idx_atoms[:, 0] == self.__nxhi] = self.__nxlo
            idx_atoms[:, 0][idx_atoms[:, 0] == self.__nxlo - 1] = self.__nxhi - 1
        if self.__pery:
            idx_atoms[:, 1][idx_atoms[:, 1] == self.__nyhi] = self.__nylo
            idx_atoms[:, 1][idx_atoms[:, 1] == self.__nylo - 1] = self.__nyhi - 1
        if self.__perz:
            idx_atoms[:, 2][idx_atoms[:, 2] == self.__nzhi] = self.__nzlo
            idx_atoms[:, 2][idx_atoms[:, 2] == self.__nzlo - 1] = self.__nzhi - 1

    def __site_id_to_indices(self, i: int) -> tuple[int, int, int, int]:
        """Decodes a site ID into its corresponding lattice indices.

        Parameters
        ----------
        i : int
            Site ID to decode.

        Returns
        -------
        tuple of int
            (ix, iy, iz, ia) indices corresponding to the site ID.
        """
        ia = i % self.__sub_count
        tmp = i // self.__sub_count
        iz = tmp % self.__nz
        tmp //= self.__nz
        iy = tmp % self.__ny
        ix = tmp // self.__ny
        ix += self.__nxlo
        iy += self.__nylo
        iz += self.__nzlo
        return ix, iy, iz, ia

    def __site_id_to_cartesian(self, i: int) -> npt.NDArray[np.float64]:
        """Converts a site ID to its corresponding cartesian coordinates.

        Parameters
        ----------
        i : int
            Site ID to decode.

        Returns
        -------
        npt.NDArray[np.float64]
            Cartesian coordinates (x, y, z) corresponding to the site ID.
        """
        ix, iy, iz, ia = self.__site_id_to_indices(i)
        x = (ix + 0.5 * ia) * self.a0
        y = (iy + 0.5 * ia) * self.a0
        z = (iz + 0.5 * ia) * self.a0
        return np.array([x, y, z], dtype=np.float64)

    def __defect_identification(
        self,
        data_atoms: defaultdict,
        idx_atoms: npt.NDArray[np.int64],
    ) -> dtypes.Defect:
        """Identifies defects based on the assignments.

        Parameters
        ----------
        data_atoms : defaultdict
            Dictionary containing simulation data as given by the LAMMPSReader and similar readers.
        idx_atoms : npt.NDArray[np.int64]
            Array of atomic index coordinates (shape: [N, 4]).

        Returns
        -------
        dtypes.Defect
            Array of defects (structured array with fields: type, x, y, z).
        """
        # Build unique site IDs from idx_atoms
        ix = idx_atoms[:, 0] - self.__nxlo
        iy = idx_atoms[:, 1] - self.__nylo
        iz = idx_atoms[:, 2] - self.__nzlo
        s = idx_atoms[:, 3]
        nsites = self.__nx * self.__ny * self.__nz * self.__sub_count
        id_sites = ((ix * self.__ny + iy) * self.__nz + iz) * self.__sub_count + s

        # Sort and group
        sort_idx = np.argsort(id_sites)
        site_sorted = id_sites[sort_idx]
        split_points = np.cumsum(np.bincount(site_sorted, minlength=nsites))[:-1]
        grouped = np.split(sort_idx, split_points)

        defects = np.empty(0, dtype=dtypes.defect)
        for i, grp in enumerate(grouped):
            if len(grp) == 0:
                # Vacancy
                vac = np.array(
                    [(0, *self.__site_id_to_cartesian(i))], dtype=dtypes.defect
                )
                defects = np.concatenate((defects, vac))
            elif len(grp) > 1:
                # Interstitials
                xyz = self.__site_id_to_cartesian(i)
                coords = str2unstr(data_atoms["atoms"][["x", "y", "z"]][grp])
                dist2 = np.sum(np.square(coords - xyz), axis=1)
                keep_idx = np.argmin(dist2)
                inters_idx = np.ones(len(grp), dtype=bool)
                inters_idx[keep_idx] = False
                if np.any(inters_idx):
                    count = np.count_nonzero(inters_idx)
                    inters = np.zeros(count, dtype=dtypes.defect)
                    inters["type"] = data_atoms["atoms"]["element"][grp][inters_idx]
                    inters["x"] = coords[inters_idx, 0]
                    inters["y"] = coords[inters_idx, 1]
                    inters["z"] = coords[inters_idx, 2]
                    defects = np.concatenate((defects, inters))
        return defects

    def identify(
        self,
        data_atoms: defaultdict,
        a1: None | float = None,
        pos_pka: None | npt.NDArray[np.float64] = None,
        theta_pka: None | float = None,
        phi_pka: None | float = None,
        transform: None | bool = False,
    ) -> defaultdict:
        """Identify defects in the crystalline structure based on atomic positions.

        Parameters
        ----------
        data_atoms : defaultdict
            Dictionary containing simulation data as given by the LAMMPSReader and similar readers.
            Must include keys: 'atoms', 'natoms', 'boundary', 'xlo', 'xhi', 'ylo', 'yhi', 'zlo',
            'zhi', 'timestep'.
        a1 : float, optional
            Final lattice parameter. If provided, defect positions are rescaled to this value
            (independently of the `transform` value).
        pos_pka : npt.NDArray[np.float64], optional
            Position vector of the PKA. If provided with theta_pka and phi_pka, defects are
            recentered and aligned.
        theta_pka : float, optional
            Polar angle (in radians) for the PKA direction.
        phi_pka : float, optional
            Azimuthal angle (in radians) for the PKA direction.
        transform : bool, optional
            If True, defects are recentered and aligned with the PKA direction (if provided). If
            True but no PKA parameters are provided, defects are recentered based on their
            average position. Note that the box boundaries are not modified for visualization
            purposes, only the atomic positions are transformed.

        Returns
        -------
        data_defects : defaultdict
            Dictionary containing the defects found in the structure. Keys are the same as in
            `data_atoms`, but the 'atoms' key contains only defects and
            'natoms' reflects the number of defects found.
        """
        self.__nxlo = round(data_atoms["xlo"] / self.a0)
        self.__nxhi = round(data_atoms["xhi"] / self.a0)
        self.__nylo = round(data_atoms["ylo"] / self.a0)
        self.__nyhi = round(data_atoms["yhi"] / self.a0)
        self.__nzlo = round(data_atoms["zlo"] / self.a0)
        self.__nzhi = round(data_atoms["zhi"] / self.a0)
        self.__nx = self.__nxhi - self.__nxlo
        self.__ny = self.__nyhi - self.__nylo
        self.__nz = self.__nzhi - self.__nzlo
        self.__perx = data_atoms["boundary"][0] == "pp"
        self.__pery = data_atoms["boundary"][1] == "pp"
        self.__perz = data_atoms["boundary"][2] == "pp"

        idx_atoms = np.zeros((data_atoms["natoms"], 4), dtype=np.int64)
        mod_atoms = np.zeros((data_atoms["natoms"], 3), dtype=np.float64)
        idx_atoms[:, 0], mod_atoms[:, 0] = np.divmod(data_atoms["atoms"]["x"], self.a0)
        idx_atoms[:, 1], mod_atoms[:, 1] = np.divmod(data_atoms["atoms"]["y"], self.a0)
        idx_atoms[:, 2], mod_atoms[:, 2] = np.divmod(data_atoms["atoms"]["z"], self.a0)
        # idx_atoms[:, 3] is initially zero

        # Compute closest unit-cell positions
        cell_xyz = self.__pos_ucell
        dist = np.sum((mod_atoms[:, None, :] - cell_xyz[None, :, :]) ** 2, axis=2)
        closest_cell_indices = np.argmin(dist, axis=1)
        idx_atoms[:, 0] += self.__idx_ucell[closest_cell_indices, 0]
        idx_atoms[:, 1] += self.__idx_ucell[closest_cell_indices, 1]
        idx_atoms[:, 2] += self.__idx_ucell[closest_cell_indices, 2]
        idx_atoms[:, 3] += self.__idx_ucell[closest_cell_indices, 3]

        # Apply boundaries on idx_atoms
        self.__apply_boundary_conditions(idx_atoms)

        # Identify defects
        defects = self.__defect_identification(data_atoms, idx_atoms)

        # Recenter / scaling if requested
        if len(defects):
            if transform:
                if a1 is None:
                    a1 = self.a0
                if (
                    pos_pka is not None
                    and theta_pka is not None
                    and phi_pka is not None
                ):
                    self.__rescale_translate_rotate(
                        defects, a1, pos_pka, theta_pka, phi_pka
                    )
                else:
                    defects[["x", "y", "z"]] -= np.mean(
                        defects[["x", "y", "z"]], axis=1
                    )
                    defects[["x", "y", "z"]] *= a1
            elif a1 is not None:
                defects[["x", "y", "z"]] *= a1

        data_defects = defaultdict(None)
        data_defects["time"] = data_atoms["time"]
        data_defects["timestep"] = data_atoms["timestep"]
        data_defects["natoms"] = len(defects)
        data_defects["boundary"] = data_atoms["boundary"]
        data_defects["xlo"] = data_atoms["xlo"]
        data_defects["xhi"] = data_atoms["xhi"]
        data_defects["ylo"] = data_atoms["ylo"]
        data_defects["yhi"] = data_atoms["yhi"]
        data_defects["zlo"] = data_atoms["zlo"]
        data_defects["zhi"] = data_atoms["zhi"]
        data_defects["atoms"] = defects

        nvacs = np.count_nonzero(defects["type"] == 0)
        nsias = len(defects) - nvacs
        if self.debug:
            print(f"Number of interstitials: {nsias}")
            print(f"Number of vacancies: {nvacs}")
        return data_defects
