"""Define the chemical elements."""

from abc import ABC
from typing import Literal

import numpy as np

from .grid import Grid

k_B = 1.38e-23  #: Boltzmann constant :math:`[J.K^{-1}]`
mu_0 = 4 * np.pi * 1.0e-7  #: Vacuum permeability :math:`[H.m^{-1}]`
gamma = 1.76e11  #: Gyromagnetic ratio :math:`[rad.s^{-1}.T^{-1}]`


class Element(ABC):
    """
    Abstract class for chemical elements.

    Args:
        T: Temperature in Kelvin
        H_ext: External magnetic field strength
        g: Grid object representing the simulation grid
        dt: Time step for the simulation
        dtype: Optional numpy dtype to cast scalar coefficients to
    """

    A: float = 0.0  #: Exchange constant :math:`[J.m^{-1}]`
    K: float = 0.0  #: Anisotropy constant :math:`[J.m^{-3}]`
    lambda_G: float = 0.0  #: Damping parameter :math:`[1]`
    M_s: float = 0.0  #: Saturation magnetization :math:`[A.m^{-1}]`
    a_eff: float = 0.0  #: Effective lattice constant :math:`[m]`
    anisotropy: Literal["uniaxial", "cubic"]  #: Type of anisotropy

    def __init__(
        self,
        T: float,
        H_ext: float,
        g: Grid,
        dt: float,
        dtype: np.dtype | None = None,
    ) -> None:
        self.H_ext = H_ext
        self.g = g
        self.dt = dt
        self.gamma_0 = gamma * mu_0  #: Rescaled gyromagnetic ratio [mA^-1.s^-1]
        self.d_0 = np.sqrt(self.A / np.abs(self.K))  #: Domain wall width [m]

        # --- Characteristic Scales ---
        self.coeff_1 = self.gamma_0 * 2.0 * self.A / (mu_0 * self.M_s)
        self.coeff_2 = self.gamma_0 * 2.0 * self.K / (mu_0 * self.M_s)
        self.coeff_3 = self.gamma_0 * H_ext

        # corresponds to the temperature actually put into the random field
        T_simu = T * self.g.dx / self.a_eff
        # calculation of the random field related to temperature
        # (we only take the volume over one mesh)
        h_random = np.sqrt(
            2 * self.lambda_G * k_B / (self.gamma_0 * mu_0 * self.M_s * self.g.dV)
        )
        H_random = h_random * np.sqrt(T_simu) * np.sqrt(1.0 / self.dt)
        self.coeff_4 = H_random * self.gamma_0

        # If a dtype is provided, cast scalar coefficients to that dtype
        if dtype is not None:
            self._cast_to_dtype(dtype)

    def _cast_to_dtype(self, dtype: np.dtype):
        """
        Cast all float attributes of the instance to the specified numpy dtype.

        Args:
            dtype: The numpy dtype to cast to (e.g., np.float32, np.float64)
        """
        # Collect instance attributes and class attributes (excluding private)
        names = set(vars(self)) | {
            k for k in vars(self.__class__) if not k.startswith("__")
        }
        for nm in names:
            val = getattr(self, nm)
            if isinstance(val, (float, np.floating)):
                setattr(self, nm, np.asarray(val, dtype=dtype))

    def get_CFL(self) -> float:
        r"""
        Returns the value of the CFL.

        .. math::

            CFL = \frac{dt \cdot 2 \gamma_0 A}{\mu_0 M_s dx^2}

        Returns:
            The CFL value
        """
        return self.dt * self.coeff_1 / self.g.dx**2

    def to_dict(self) -> dict:
        """
        Export element parameters to a dictionary for JAX JIT compatibility.

        Returns:
            Dictionary containing element parameters needed for computations
        """
        # Map anisotropy string to integer for JIT compatibility
        aniso_map = {"uniaxial": 0, "cubic": 1}

        return {
            "coeff_1": self.coeff_1,
            "coeff_2": self.coeff_2,
            "coeff_3": self.coeff_3,
            "coeff_4": self.coeff_4,
            "lambda_G": self.lambda_G,
            "anisotropy": aniso_map[self.anisotropy],
            "gamma_0": self.gamma_0,
        }

    def compute_H_anisotropy(self, m: np.ndarray, H_aniso: np.ndarray):
        r"""
        Compute the anisotropy field.

        For uniaxial anisotropy:

        .. math::

            \boldsymbol{H}_{\text{ani, uniaxial}}=\frac{2K}{\mu_0M_s}(\boldsymbol{e}_x
            \cdot\boldsymbol{m})\boldsymbol{e}_x,\label{uniaxial}


        For cubic anisotropy:

        .. math::

            \boldsymbol{H}_{\text{ani, cubic}}=-\frac{2K}{\mu_0M_s}\sum_{(i,j,k)\in I}
            \left((\boldsymbol{e}_j\cdot\boldsymbol{m})^2+(\boldsymbol{e}_k\cdot
            \boldsymbol{m})^2+(\boldsymbol{e}_j\cdot\boldsymbol{m})^2(\boldsymbol{e}_k
            \cdot\boldsymbol{m})^2\right)(\boldsymbol{e}_i\cdot\boldsymbol{m})\boldsymbol{e}_i,\label{cubic}


        Args:
            m: Magnetization array (shape (3, nx, ny, nz)).
            H_aniso: Pre-allocated output array (shape (3, nx, ny, nz)).

        Raises:
            ValueError: If the anisotropy type is unknown
        """
        if self.anisotropy == "uniaxial":
            H_aniso[0] = self.coeff_2 * m[0]
            H_aniso[1] = 0.0
            H_aniso[2] = 0.0
        elif self.anisotropy == "cubic":
            m1, m2, m3 = m
            m1m1 = m1 * m1
            m2m2 = m2 * m2
            m3m3 = m3 * m3
            H_aniso[0] = -self.coeff_2 * (1 - m1m1 + m2m2 * m3m3) * m1
            H_aniso[1] = -self.coeff_2 * (1 - m2m2 + m1m1 * m3m3) * m2
            H_aniso[2] = -self.coeff_2 * (1 - m3m3 + m1m1 * m2m2) * m3
        else:
            raise ValueError(f"Unknown anisotropy type: {self.anisotropy}")


class Cobalt(Element):
    """Cobalt element."""

    A = 30.0e-12  #: Exchange constant :math:`[J.m^{-1}]`
    K = 520.0e3  #: Anisotropy constant :math:`[J.m^{-3}]`
    lambda_G = 0.5  #: Damping parameter :math:`[1]`
    M_s = 1400.0e3  #: Saturation magnetization :math:`[A.m^{-1}]`
    a_eff = 0.25e-9  #: Effective lattice constant :math:`[m]`
    anisotropy = "uniaxial"  #: Type of anisotropy (e.g., "uniaxial", "cubic")


class Iron(Element):
    """Iron element."""

    A = 21.0e-12  #: Exchange constant :math:`[J.m^{-1}]`
    K = 48.0e3  #: Anisotropy constant :math:`[J.m^{-3}]`
    lambda_G = 0.5  #: Damping parameter :math:`[1]`
    M_s = 1700.0e3  #: Saturation magnetization :math:`[A.m^{-1}]`
    a_eff = 0.286e-9  #: Effective lattice constant :math:`[m]`
    anisotropy = "cubic"  #: Type of anisotropy (e.g., "uniaxial", "cubic")


class Nickel(Element):
    """Nickel element."""

    A = 9.0e-12  #: Exchange constant :math:`[J.m^{-1}]`
    K = -5.7e3  #: Anisotropy constant :math:`[J.m^{-3}]`
    lambda_G = 0.5  #: Damping parameter :math:`[1]`
    M_s = 490.0e3  #: Saturation magnetization :math:`[A.m^{-1}]`
    a_eff = 0.345e-9  #: Effective lattice constant :math:`[m]`
    anisotropy = "cubic"  #: Type of anisotropy (e.g., "uniaxial", "cubic")


def get_element_class(element_name: str) -> type[Element]:
    """
    Get the class of the chemical element by its name.

    Example:
        >>> cls = get_element_class("Cobalt")
        >>> cls.__name__
        'Cobalt'

    Args:
        element_name: The name of the element

    Returns:
        The class of the element

    Raises:
        ValueError: If the element is not found
    """
    for cls in Element.__subclasses__():
        if cls.__name__ == element_name:
            return cls
    raise ValueError(f"Element '{element_name}' not found in {__file__}.")
