"""LLG3D solver using NumPy."""

from typing import ClassVar
import time

import numpy as np

from .base import BaseSolver, timeit
from .math_utils import cross_product


class NumpySolver(BaseSolver):
    """NumPy-based LLG3D solver."""

    solver_type: ClassVar[str] = "numpy"  #: Solver type name

    @timeit
    def laplacian_3d(self, m_i: np.ndarray) -> np.ndarray:
        """
        Returns the laplacian of m_i in 3D with Neumann boundary conditions.

        Args:
            m_i: Magnetization direction (shape (nx, ny, nz))

        Returns:
            Laplacian of m_i (shape (nx, ny, nz))
        """
        m_i_padded = np.pad(m_i, ((1, 1), (1, 1), (1, 1)), mode="reflect")

        laplacian = (
            self.grid.inv_dx2
            * (m_i_padded[2:, 1:-1, 1:-1] + m_i_padded[:-2, 1:-1, 1:-1])
            + self.grid.inv_dy2
            * (m_i_padded[1:-1, 2:, 1:-1] + m_i_padded[1:-1, :-2, 1:-1])
            + self.grid.inv_dz2
            * (m_i_padded[1:-1, 1:-1, 2:] + m_i_padded[1:-1, 1:-1, :-2])
            + self.grid.center_coeff * m_i
        )
        return laplacian

    @timeit
    def compute_laplacian(self, m: np.ndarray) -> np.ndarray:
        """
        Compute the laplacian of m in 3D with Neumann boundary conditions.

        Args:
            m: Magnetization array (shape (3, nx, ny, nz))

        Returns:
            Laplacian of m (shape (3, nx, ny, nz))
        """
        return np.stack(
            [
                self.laplacian_3d(m[0]),
                self.laplacian_3d(m[1]),
                self.laplacian_3d(m[2]),
            ],
            axis=0,
        )

    @timeit
    def _compute_slope(
        self, m_n: np.ndarray, R_random: np.ndarray, H_aniso: np.ndarray
    ) -> np.ndarray:
        """
        Compute the slope of the LLG equation.

        Args:
            m_n: Magnetization array at time step n (shape (3, nx, ny, nz)).
            R_random: Random field array (shape (3, nx, ny, nz)).
            H_aniso: Optional pre-allocated buffer for anisotropy field.

        Returns:
            Slope array (shape (3, nx, ny, nz))
        """
        self.elem.compute_H_anisotropy(m_n, H_aniso)

        laplacian_m = self.compute_laplacian(m_n)
        R_eff = self.elem.coeff_1 * laplacian_m + R_random + H_aniso
        R_eff[0] += self.elem.coeff_3

        m_cross_R_eff = cross_product(m_n, R_eff)
        m_cross_m_cross_R_eff = cross_product(m_n, m_cross_R_eff)

        s = -(m_cross_R_eff + self.elem.lambda_G * m_cross_m_cross_R_eff)

        return s

    def _simulate(self) -> float:
        """
        Simulates the system for N iterations.

        Returns:
            The time taken for the simulation
        """
        m_n = self._init_m_n()
        H_aniso = np.empty_like(m_n)
        t = 0.0
        self._record(m_n, t, 0)  # Record the initial solution

        start_time = time.perf_counter()

        for n in self._progress_bar():
            t += self.dt

            # Adding randomness: temperature effect
            R_random = self._get_R_random()

            # Prediction phase
            s_pre = self._compute_slope(m_n, R_random, H_aniso)
            m_pre = m_n + self.dt * s_pre

            # Correction phase
            s_cor = self._compute_slope(m_pre, R_random, H_aniso)
            m_n += self.dt * 0.5 * (s_pre + s_cor)

            # We renormalize to verify the constraint of being on the sphere
            self._normalize(m_n)

            self._record(m_n, t, n)

        total_time = time.perf_counter() - start_time

        self._finalize()

        return total_time
