"""
LLG3D solver using MPI.

The parallelization is done in the x direction.
"""

import argparse
import sys
import time

import mpi4py
import numpy as np

from .. import solvers
from .base import BaseSolver
from .math_utils import cross_product
from .profiling import timeit

mpi4py.rc.initialize = False
mpi4py.rc.finalize = True

from mpi4py import MPI  # noqa: E402

comm: MPI.Comm
rank: int
size: int
status: MPI.Status

Boundaries = tuple[np.ndarray, np.ndarray, MPI.Request | None, MPI.Request | None]
Boundaries_3d = tuple[Boundaries, Boundaries, Boundaries]


def initialize_mpi():
    """Initialize MPI variables."""
    if not MPI.Is_initialized():
        MPI.Init()

    global comm, rank, size, status
    # Set MPI variables with the actual MPI values
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    size = comm.Get_size()
    status = MPI.Status()
    # Set these variables in the solver module as well
    solvers.comm = comm
    solvers.rank = rank
    solvers.size = size
    solvers.status = status
    solvers.mpi_initialized = True


class MPISolver(BaseSolver):
    """MPI LLG3D solver."""

    solver_type = "mpi"

    def __post_init__(self):
        """Initialize MPI and check parameters."""
        initialize_mpi()
        super().__post_init__()
        if self.Jx % size != 0:
            raise ValueError(
                f"Jx ({self.Jx}) must be divisible by the number of processes ({size})"
            )

    def _get_boundaries_x(self, m: np.ndarray) -> Boundaries:
        """
        Returns the boundaries asynchronously.

        Allows overlapping communication time of boundaries with calculations
        if non-blocking communication is used.

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

        Returns:
            - m_i_x_start: Start boundary in x direction
            - m_i_x_end: End boundary in x direction
            - request_start: Request for start boundary or None
            - request_end: Request for end boundary or None
        """
        # Extract slices for Neumann boundary conditions

        # Create boundary buffers with the same dtype as `m` to avoid
        # MPI packing/unpacking issues when precision is `single`.
        m_i_x_start = np.empty((1, self.grid.Jy, self.grid.Jz), dtype=self.np_float)
        m_i_x_end = np.empty_like(m_i_x_start)

        # Prepare ring communication:
        # Even if procs 0 and size - 1 shouldn't receive anything from left
        # and right respectively, it's simpler to express it like this
        right = (rank + 1) % size
        left = (rank - 1 + size) % size

        if self.blocking:
            # Wait for boundaries to be available
            comm.Sendrecv(
                m[:1, :, :], dest=left, sendtag=0, recvbuf=m_i_x_end, source=right
            )
            comm.Sendrecv(
                m[-1:, :, :], dest=right, sendtag=1, recvbuf=m_i_x_start, source=left
            )
            return m_i_x_start, m_i_x_end, None, None
        else:
            request_start = comm.Irecv(m_i_x_start, source=left, tag=201)
            request_end = comm.Irecv(m_i_x_end, source=right, tag=202)
            comm.Isend(m[-1:, :, :], dest=right, tag=201)
            comm.Isend(m[:1, :, :], dest=left, tag=202)

            return m_i_x_start, m_i_x_end, request_start, request_end

    @timeit
    def _get_boundaries_3d(self, m: np.ndarray) -> Boundaries_3d:
        """
        Returns the boundaries for all components.

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

        Returns:
            Boundaries for all three components
        """
        return (
            self._get_boundaries_x(m[0]),
            self._get_boundaries_x(m[1]),
            self._get_boundaries_x(m[2]),
        )

    def laplacian_3d(
        self,
        m_i: np.ndarray,
        m_i_x_start: np.ndarray,
        m_i_x_end: np.ndarray,
        request_start: MPI.Request | None,
        request_end: MPI.Request | None,
    ) -> np.ndarray:
        """
        Returns the Laplacian of m_i in 3D.

        We start by calculating contributions in y and z, to wait
        for the end of communications in x.

        Args:
            m_i: i-component of the magnetization array (shape (nx, ny, nz))
            m_i_x_start: start boundary in x direction
            m_i_x_end: end boundary in x direction
            request_start: request for start boundary
            request_end: request for end boundary

        Returns:
            Laplacian of m_i (shape (nx, ny, nz))
        """
        # Extract slices for Neumann boundary conditions
        m_i_y_start = m_i[:, 1:2, :]
        m_i_y_end = m_i[:, -2:-1, :]

        m_i_z_start = m_i[:, :, 1:2]
        m_i_z_end = m_i[:, :, -2:-1]

        m_i_y_plus = np.concatenate((m_i[:, 1:, :], m_i_y_end), axis=1)
        m_i_y_minus = np.concatenate((m_i_y_start, m_i[:, :-1, :]), axis=1)
        m_i_z_plus = np.concatenate((m_i[:, :, 1:], m_i_z_end), axis=2)
        m_i_z_minus = np.concatenate((m_i_z_start, m_i[:, :, :-1]), axis=2)

        if self.grid.uniform:  # Uniform grid spacing: limit truncature errors
            laplacian = self.grid.inv_dx2 * (
                m_i_y_plus + m_i_y_minus + m_i_z_plus + m_i_z_minus - 6 * m_i
            )
        else:
            laplacian = (
                self.grid.inv_dy2 * (m_i_y_plus + m_i_y_minus)
                + self.grid.inv_dz2 * (m_i_z_plus + m_i_z_minus)
                + self.grid.center_coeff * m_i
            )

        # Wait for x-boundaries to be available (communications completed)
        if request_start is not None:  # non blocking communication
            request_start.Wait(status)
        if request_end is not None:  # non blocking communication
            request_end.Wait(status)

        # For extreme procs, apply Neumann boundary conditions in x
        if rank == size - 1:
            m_i_x_end = m_i[-2:-1, :, :]
        if rank == 0:
            m_i_x_start = m_i[1:2, :, :]

        m_i_x_plus = np.concatenate((m_i[1:, :, :], m_i_x_end), axis=0)
        m_i_x_minus = np.concatenate((m_i_x_start, m_i[:-1, :, :]), axis=0)

        laplacian += self.grid.inv_dx2 * (m_i_x_plus + m_i_x_minus)

        return laplacian

    @timeit
    def compute_laplacian(self, m: np.ndarray, boundaries: Boundaries_3d) -> np.ndarray:
        """
        Compute the laplacian of m in 3D.

        Args:
            m: Magnetization array (shape (3, nx, ny, nz))
            boundaries: Boundaries for x direction

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

    @timeit
    def compute_slope(
        self,
        m: np.ndarray,
        R_random: np.ndarray,
        H_aniso: np.ndarray,
        boundaries: Boundaries_3d,
    ) -> np.ndarray:
        """
        Compute the slope of the LLG equation.

        Args:
            m: Magnetization array (shape (3, nx, ny, nz))
            R_random: Random field array (shape (3, nx, ny, nz))
            H_aniso: Pre-allocated buffer for anisotropy field.
            boundaries: Boundaries for x direction

        Returns:
            Slope array (shape (3, nx, ny, nz))
        """
        # Precalculate terms used multiple times

        self.elem.compute_H_anisotropy(m, H_aniso)

        laplacian_m = self.compute_laplacian(m, boundaries)
        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, R_eff)
        m_cross_m_cross_R_eff = cross_product(m, m_cross_R_eff)

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

        return s

    @timeit
    def _xyz_average(self, m: np.ndarray) -> float:
        """
        Returns the spatial average of m with shape (g.dims) using the midpoint method.

        Performs the local sum on each process and then reduces it to process 0.

        Args:
            m: Array to be integrated (shape (x, y, z))

        Returns:
            Spatial average of m
        """
        # Make a copy of m to avoid modifying its value
        mm = m.copy()

        # On y and z edges, divide the contribution by 2
        mm[:, 0, :] /= 2
        mm[:, -1, :] /= 2
        mm[:, :, 0] /= 2
        mm[:, :, -1] /= 2

        # On x edges (only on extreme procs), divide the contribution by 2
        if rank == 0:
            mm[0] /= 2
        if rank == size - 1:
            mm[-1] /= 2
        local_sum = mm.sum()

        # Sum across all processes gathered by process 0
        global_sum = comm.reduce(local_sum)

        # Spatial average is the global sum divided by the number of cells
        return float(global_sum) / self.grid.ncell if global_sum is not None else 0.0

    @timeit
    def _yz_average(self, m_i: np.ndarray) -> np.ndarray:
        """
        Compute the integral of m_i over y and z directions.

        Gather the local profiles from all processes to form the global profile.

        Args:
            m_i: Magnetization component array (shape (nx, ny, nz))

        Returns:
            1D array of the integral over y and z (shape (nx,))
        """
        local_yz_average = super()._yz_average(m_i)
        # global coordinates
        global_yz_average = np.empty(self.grid.Jx, dtype=self.np_float)
        comm.Gather(local_yz_average, global_yz_average)
        return global_yz_average

    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

            # Prediction phase
            x_boundaries: Boundaries_3d = self._get_boundaries_3d(m_n)

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

            s_pre = self.compute_slope(m_n, R_random, H_aniso, x_boundaries)
            m_pre = m_n + self.dt * s_pre

            # Correction phase
            x_boundaries = self._get_boundaries_3d(m_pre)

            s_cor = self.compute_slope(m_pre, R_random, H_aniso, x_boundaries)
            m_n += self.dt * 0.5 * (s_pre + s_cor)

            # 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


class ArgumentParser(argparse.ArgumentParser):
    """An argument parser compatible with MPI."""

    def _print_message(self, message, file=None):
        if rank == 0 and message:
            if file is None:
                file = sys.stderr
            file.write(message)

    def exit(self, status: int = 0, message: str | None = None):
        """
        Exit the program using MPI finalize.

        Args:
            status: Exit status code
            message: Optional exit message
        """
        if message:
            self._print_message(message, sys.stderr)
        comm.barrier()
        MPI.Finalize()
        exit(status)
