"""Test suite for the Element classes."""

from typing import Any
import numpy as np
import pytest

from llg3d.element import get_element_class, Cobalt, Iron, Nickel, gamma, mu_0
from llg3d.grid import Grid
from llg3d.solvers.math_utils import normalize


ElementDict = dict[str, Any]


@pytest.fixture
def element_params(grid_standard: Grid) -> ElementDict:
    """Fixture providing common element parameters."""
    return {
        "T": 1100,
        "H_ext": 0.0,
        "g": grid_standard,
        "dt": 1e-12,
    }


@pytest.fixture
def cobalt(element_params: ElementDict) -> Cobalt:
    """Fixture providing a Cobalt element instance."""
    return Cobalt(**element_params)


@pytest.fixture
def iron(element_params: ElementDict) -> Iron:
    """Fixture providing an Iron element instance."""
    return Iron(**element_params)


@pytest.fixture
def nickel(element_params: ElementDict) -> Nickel:
    """Fixture providing a Nickel element instance."""
    return Nickel(**element_params)


def test_cobalt_properties(cobalt: Cobalt):
    """Test that Cobalt element has correct properties."""
    assert cobalt.A == 30.0e-12
    assert cobalt.K == 520.0e3
    assert cobalt.gamma_0 == gamma * mu_0
    assert cobalt.lambda_G == 0.5
    assert cobalt.M_s == 1400.0e3
    assert cobalt.a_eff == 0.25e-9
    assert cobalt.anisotropy == "uniaxial"


def test_iron_to_dict(iron: Iron):
    """Test that element parameters can be converted to a dictionary."""
    d = iron.to_dict()

    assert d == {
        "coeff_1": iron.coeff_1,
        "coeff_2": iron.coeff_2,
        "coeff_3": iron.coeff_3,
        "coeff_4": iron.coeff_4,
        "lambda_G": iron.lambda_G,
        "anisotropy": 1,  # Cubic is mapped to 1
        "gamma_0": iron.gamma_0,
    }


def test_get_element_class():
    """Test that get_element_class returns Cobalt for 'Cobalt'."""
    assert get_element_class("Cobalt") == Cobalt
    assert get_element_class("Iron") == Iron
    assert get_element_class("Nickel") == Nickel


def test_get_element_class_invalid_raises():
    """Test that get_element_class raises ValueError for unknown element."""
    with pytest.raises(ValueError):
        get_element_class("UnknownElement")


def test_cast_to_dtype(nickel: Nickel):
    """Test that element coefficients are cast to specified dtype."""
    nickel._cast_to_dtype(np.dtype(np.float32))

    for attr in ["coeff_1", "coeff_2", "coeff_3", "coeff_4", "lambda_G", "gamma_0"]:
        val = getattr(nickel, attr)
        assert isinstance(val, np.ndarray)
        assert val.dtype == np.float32


def test_get_CFL(cobalt: Cobalt):
    """Test that get_CFL returns a positive float."""
    cfl = cobalt.get_CFL()
    assert isinstance(cfl, float)
    assert cfl > 0.0


def test_compute_H_anisotropy_uniaxial(
    grid_standard: Grid, element_params: ElementDict
):
    """Test the compute_H_anisotropy function."""

    m = np.random.rand(3, *grid_standard.dims)
    normalize(m)  # Make sure m is normalized
    H_aniso_uniaxial = np.empty_like(m)
    elem_uniaxial = Cobalt(**element_params)
    elem_uniaxial.compute_H_anisotropy(m, H_aniso_uniaxial)

    assert H_aniso_uniaxial.shape == m.shape

    # For uniaxial, only first component should be non-zero
    assert np.all(H_aniso_uniaxial[0] == elem_uniaxial.coeff_2 * m[0])
    assert np.all(H_aniso_uniaxial[1] == 0.0)
    assert np.all(H_aniso_uniaxial[2] == 0.0)


def test_compute_H_anisotropy_cubic(grid_standard: Grid, element_params: ElementDict):
    elem_cubic = Iron(**element_params)

    # H_aniso should be zero when m is aligned with x, y, or z axes
    for i in range(3):
        assert elem_cubic.anisotropy == "cubic"
        m = np.zeros((3, *grid_standard.dims))
        m[i] = 1.0  # Aligned along x-axis
        H_aniso_cubic = np.empty_like(m)
        elem_cubic.compute_H_anisotropy(m, H_aniso_cubic)
        assert H_aniso_cubic.shape == m.shape
        assert np.all(H_aniso_cubic == 0.0)

    m = np.random.rand(3, *grid_standard.dims)
    normalize(m)  # Make sure m is normalized
    H_aniso_cubic = np.empty_like(m)
    elem_cubic = Iron(**element_params)
    elem_cubic.compute_H_anisotropy(m, H_aniso_cubic)
    assert H_aniso_cubic.shape == m.shape
    # |H_aniso[i]| ≤ 1.5 × |coeff_2|
    coeff_2 = elem_cubic.coeff_2
    H_aniso_max = 1.5 * abs(coeff_2)
    for i in range(3):
        assert np.all(np.abs(H_aniso_cubic[i]) <= H_aniso_max)


def test_compute_H_anisotropy_invalid_anisotropy(grid_standard: Grid):
    """Test that compute_H_anisotropy raises ValueError for unknown anisotropy."""

    class UnknownAnisotropyElement(Cobalt):
        anisotropy = "unknown"  # type: ignore

    m = np.random.rand(3, *grid_standard.dims)

    H_aniso = np.empty_like(m)
    elem_unknown = UnknownAnisotropyElement(
        T=300.0, H_ext=0.0, g=grid_standard, dt=1e-13
    )

    with pytest.raises(ValueError, match="Unknown anisotropy type: unknown"):
        elem_unknown.compute_H_anisotropy(m, H_aniso)  # type: ignore
