"""Utility functions related to collisions' debris from SRIM data."""

import pathlib
from typing import Optional, Union

import numpy as np

from irradiapy import dpa, dtypes, materials
from irradiapy.damagedb import DamageDB
from irradiapy.io.xyzwriter import XYZWriter
from irradiapy.srimpy.srimdb import SRIMDB


def generate_debris(
    srimdb: SRIMDB,
    dir_mddb: pathlib.Path,
    compute_tdam: bool,
    path_collisions: pathlib.Path,
    dpa_mode: dpa.DpaMode,
    add_injected: bool,
    outsiders: bool,
    seed: Optional[int] = None,
    depth_offset: Union[int, float, np.number] = 0.0,
) -> None:
    """Turns SRIM's collisions into `.xyz` files for the given database of cascades' debris.

    Warning
    -------
        Assumes a monolayer monoatomic target and same element for all ions.

    Parameters
    ----------
    srimdb : SRIMDB
        SRIM database class.
    dir_mddb : pathlib.Path
        Directory where the database of cascades' debris is stored.
    compute_tdam : bool
        Whether to transform the PKA energies into damage energies. It should be `True` for
        MD simulations without electronic stopping.
    path_collisions : pathlib.Path
        Directory where the ions debris will be stored as `.xyz` files.
    dpa_mode : dpa.DpaMode
        Formula to convert the residual energy into Frenkel pairs.
    add_injected : bool
        Whether to add the injected interstitial.
    outsiders : bool
        Whether to include defects generated outside the material (PKAs close to surfaces).
    seed : int, optional
        Random seed for reproducibility.
    depth_offset : float, optional
        Offset to add to the depth of the defects, by default 0.0.
    """
    depth_min, depth_max = 0.0, srimdb.target.layers[0].width
    mat_pka = materials.get_material_by_atomic_number(
        next(
            iter(srimdb.trimdat.read(what="atom_numb", condition="WHERE ion_numb = 1"))
        )[0]
    )
    mat_target = materials.get_material_by_atomic_number(
        srimdb.target.layers[0].elements[0].atomic_number
    )
    nions = srimdb.nions
    damagedb = DamageDB(
        dir_mddb=dir_mddb,
        compute_tdam=compute_tdam,
        mat_pka=mat_pka,
        mat_target=mat_target,
        dpa_mode=dpa_mode,
        seed=seed,
    )
    with XYZWriter(path_collisions) as writer:
        for nion in range(1, nions + 1):
            writer.save(
                __generate_ion_defects(
                    srimdb,
                    nion,
                    damagedb,
                    add_injected,
                    depth_min,
                    depth_max,
                    outsiders,
                    depth_offset,
                ),
                extra_comment=f"Projectile={nion}",
            )


def __generate_ion_defects(
    srimdb: SRIMDB,
    nion: int,
    damagedb: DamageDB,
    add_injected: bool,
    depth_min: float,
    depth_max: float,
    outsiders: bool,
    depth_offset: Union[int, float, np.number] = 0.0,
) -> np.ndarray:
    """Generates the defects for a specific ion in the SRIM simulation.

    Parameters
    ----------
    srimdb : SRIMDB
        SRIM database class.
    nion : int
        Ion number.
    damagedb : DamageDB
        DamageDB class that will choose MD debris.
    add_injected : bool
        Whether to add the injected interstitial.
    depth_min : float
        Minimum depth to consider for defects.
    depth_max : float
        Maximum depth to consider for defects.
    outsiders : bool
        Whether to include defects generated outside the material (PKAs close to surface).
    depth_offset : float, optional
        Offset to add to the depth of the defects, by default 0.0.

    Returns
    -------
    np.ndarray
        An array containing the defects generated by a single ion.
    """
    defects = np.empty(0, dtype=dtypes.defect)
    for depth, y, z, cosx, cosy, cosz, pka_e in srimdb.collision.read(
        what="depth, y, z, cosx, cosy, cosz, recoil_energy",
        condition=f"WHERE ion_numb = {nion}",
    ):
        pka_pos = np.array([depth, y, z])
        pka_dir = np.array([cosx, cosy, cosz])
        defects_ = damagedb.get_pka_debris(
            pka_e=pka_e, pka_pos=pka_pos, pka_dir=pka_dir
        )
        defects = np.concatenate((defects, defects_))

    if add_injected:
        injected = list(
            srimdb.range3d.read(
                what="depth, y, z", condition=f"WHERE ion_numb = {nion} LIMIT 1"
            )
        )
        if injected:
            atom_number = list(
                srimdb.trimdat.read(
                    what="atom_numb", condition=f"WHERE ion_numb = {nion}"
                )
            )[0][0]
            injected = np.array([(atom_number, injected[0])], dtype=dtypes.defect)
            defects = np.concatenate((defects, injected))
    if not outsiders:
        defects = defects[
            (defects["pos"][:, 0] >= depth_min) & (defects["pos"][:, 0] <= depth_max)
        ]
    defects["pos"][:, 0] += depth_offset
    return defects
