"""Defines a class used for subdividing a grid for parallel execution using MPI.

.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
"""

from __future__ import annotations

import math
from enum import IntEnum
from typing import TYPE_CHECKING, Any, Literal, TypeVar

import numpy as np

from ..fields import FieldCollection
from ..fields.base import FieldBase
from ..fields.datafield_base import DataFieldBase
from ..tools import mpi
from ..tools.cache import cached_method
from ..tools.plotting import plot_on_axes
from .boundaries.axes import BoundariesBase, BoundariesList
from .boundaries.axis import BoundaryAxisBase, BoundaryPair
from .boundaries.local import _MPIBC

if TYPE_CHECKING:
    from collections.abc import Sequence

    from ..tools.typing import NumericArray
    from .base import GridBase


class MPIFlags(IntEnum):
    """Enum that contains flags for MPI communication."""

    field_split = 1  # split full field onto nodes
    field_combine = 2  # combine full field from subfields on nodes
    _boundary_lower = 8  # exchange with lower boundary of node with lower id
    _boundary_upper = 9  # exchange with upper boundary of node with lower id

    @classmethod
    def boundary_lower(cls, my_id: int, other_id: int) -> int:
        """Flag for connection between my lower boundary and `other_id`"""
        if my_id <= other_id:
            return 2 * my_id + cls._boundary_lower
        return 2 * other_id + cls._boundary_upper

    @classmethod
    def boundary_upper(cls, my_id: int, other_id: int) -> int:
        """Flag for connection between my upper boundary and `other_id`"""
        if my_id <= other_id:
            return 2 * my_id + cls._boundary_upper
        return 2 * other_id + cls._boundary_lower


def _get_optimal_decomposition(shape: Sequence[int], mpi_size: int) -> list[int]:
    """Determine optimal decomposition of a grid into several chunks.

    Args:
        shape (list of int):
            Size of each axis of the grid
        mpi_size (int):
            Number of nodes for decomposition

    Results:
        list of int: The number of decompositions of each axis
    """
    shape_arr = np.asarray(shape)
    decomposition = [-1] * len(shape)
    order = np.argsort(shape_arr)
    size_left = mpi_size  # size of axes that have been determined
    # iterate through dimensions along increasing axes length
    for dim_count, dim in enumerate(order):
        # get length of all axes that are left to be determined
        shape_left = shape_arr[order[dim_count:]]
        # determine the number of cells in these axes
        size_undetermined = np.prod(shape_left)
        # estimate optimal number of grid points on each node
        node_size_estimate = size_undetermined // size_left
        # estimate optimal length of each remaining axis
        node_axis_len_est = node_size_estimate ** (1 / len(shape_left))
        node_axis_len = max(1, int(np.floor(node_axis_len_est)))
        # get estimate for the number of subdivision for current axis
        decomp_axis_est = shape[dim] // node_axis_len
        # restrict the number of subdivisions to sensible range
        decomposition[dim] = np.clip(decomp_axis_est, 1, size_left)
        # calculate the number of nodes that are left
        size_left //= decomposition[dim]
    assert np.prod(np.asarray(decomposition)) <= mpi_size
    return decomposition


def _subdivide(num: int, chunks: int) -> NumericArray:
    r"""Subdivide `num` intervals in `chunk` chunks.

    Args:
        num (int):
            Number of intervals
        chunks (int):
            Number of chunks

    Returns:
        list: The number of intervals per chunk
    """
    if chunks > num:
        msg = "Cannot divide in more chunks than support points"
        raise RuntimeError(msg)
    return np.diff(np.linspace(0, num, chunks + 1).astype(int))


def _subdivide_along_axis(grid: GridBase, axis: int, chunks: int) -> list[GridBase]:
    """Subdivide the grid along a given axis.

    Args:
        axis (int):
            The axis along which the subdivision will happen
        chunks (int):
            The number of chunks along this axis

    Returns:
        list: A list of subgrids
    """
    if chunks <= 0:
        msg = "Chunks must be a positive Integer"
        raise ValueError(msg)
    if chunks == 1:
        return [grid]  # no subdivision necessary

    def replace_in_axis(arr, value):
        if isinstance(arr, tuple):
            return (*arr[:axis], value, *arr[axis + 1 :])
        res = arr.copy()
        res[axis] = value
        return res

    subgrids = []
    start = 0
    for size in _subdivide(grid.shape[axis], chunks):
        # determine new sub region
        end = start + size
        shape = replace_in_axis(grid.shape, size)
        periodic = replace_in_axis(grid.periodic, False)

        # determine bounds of the new grid
        axis_bounds = grid.axes_bounds[axis]
        cell_bounds = np.linspace(*axis_bounds, grid.shape[axis] + 1)
        bounds = replace_in_axis(
            grid.axes_bounds, (cell_bounds[start], cell_bounds[end])
        )

        # create new subgrid
        subgrid = grid.__class__.from_bounds(bounds, shape, periodic)
        subgrids.append(subgrid)
        start = end  # for next iteration

    return subgrids


TField = TypeVar("TField", bound=FieldBase)
TData = TypeVar("TData")


class GridMesh:
    """Handles a collection of subgrids arranged in a regular mesh.

    This class provides methods for managing MPI simulations of multiple connected
    subgrids. Each subgrid is also called a cell and identified with a unique number.
    """

    def __init__(self, basegrid: GridBase, subgrids: Sequence):
        """
        Args:
            basegrid (:class:`~pde.grids.base.GridBase`):
                The grid of the entire domain
            subgrids (nested lists of :class:`~pde.grids.base.GridBase`):
                The nested grids representing the subdivision
        """
        if basegrid._mesh is not None:
            msg = "Cannot subdivide a subgrid further"
            raise ValueError(msg)
        self.basegrid = basegrid
        self.subgrids = np.asarray(subgrids)
        for subgrid in self.subgrids.flat:
            subgrid._mesh = self

        assert basegrid.num_axes == self.subgrids.ndim

    @classmethod
    def from_grid(
        cls, grid: GridBase, decomposition: Literal["auto"] | int | list[int] = "auto"
    ) -> GridMesh:
        """Subdivide the grid into subgrids.

        Args:
            grid (:class:`~pde.grids.base.GridBase`):
                Grid that will be subdivided
            decomposition (list of ints):
                Number of subdivision in each direction. Should be a list of length
                `grid.num_axes` specifying the number of nodes for this axis. If one
                value is `-1`, its value will be determined from the number of available
                nodes. The default value `auto` tries to determine an optimal
                decomposition by minimizing communication between nodes.

        Returns:
            :class:`GridMesh`: The grid mesh created from the grid
        """
        if decomposition == "auto":
            # determine decomposition automatically
            decomposition = _get_optimal_decomposition(grid.shape, mpi.size)
        else:
            # parse `decomposition`
            try:
                decomposition = [int(d) for d in decomposition]  # type: ignore
            except TypeError:
                decomposition = [int(decomposition)]  # type: ignore

            size, var_index = 1, None
            for i, num in enumerate(decomposition):
                if num == -1:
                    if var_index is None:
                        var_index = i
                    else:
                        msg = "Can only specify one unknown dimension"
                        raise ValueError(msg)
                elif num > 0:
                    size *= num
                else:
                    msg = f"Unknown size `{num}`"
                    raise RuntimeError(msg)

            # replace potential variable index with correct value
            if var_index is not None:
                dim = mpi.size // size
                if dim > 0:
                    decomposition[var_index] = dim
                else:
                    msg = "Not enough nodes to satisfy decomposition"
                    raise RuntimeError(msg)

            # fill up with 1s until the grid size is met
            decomposition += [1] * (grid.num_axes - len(decomposition))

        # check compatibility with number of nodes
        if mpi.size > 1 and math.prod(decomposition) != mpi.size:
            msg = (
                f"Node count ({mpi.size}) incompatible with decomposition "
                f"({decomposition})"
            )
            raise RuntimeError(msg)

        # subdivide the base grid according to the decomposition
        subgrids = np.empty(decomposition, dtype=object)
        subgrids.flat[0] = grid.copy()  # seed the initial grid at the top-left
        idx_set: list[Any] = [0] * subgrids.ndim  # indices to extract all grids
        for axis, chunks in enumerate(decomposition):
            # iterate over all grids that have been determined already
            for idx, subgrid in np.ndenumerate(subgrids[tuple(idx_set)]):
                i = idx + (slice(None, None),) + (0,) * (subgrids.ndim - axis - 1)
                divison = _subdivide_along_axis(subgrid, axis=axis, chunks=chunks)
                subgrids[i] = divison

            # mark this axis as set
            idx_set[axis] = slice(None, None)

        return cls(basegrid=grid, subgrids=subgrids)  # type: ignore

    @property
    def num_axes(self) -> int:
        """int: the number of axes that the grids possess"""
        return self.subgrids.ndim

    @property
    def shape(self) -> tuple[int, ...]:
        """tuple: the number of subgrids along each axis"""
        return self.subgrids.shape

    def __len__(self) -> int:
        """Total number of subgrids."""
        return self.subgrids.size

    @property
    def current_node(self) -> int:
        """int: current MPI node"""
        return mpi.rank

    def __getitem__(self, node_id: int | None) -> GridBase:
        """Extract one subgrid from the mesh.

        Args:
            node_id (int):
                Index identifying the subgrid

        Returns:
            :class:`~pde.grids.base.GridBase`: The sub grid of the given cell
        """
        if node_id is None:
            node_id = self.current_node
        return self.subgrids[self._id2idx(node_id)]  # type: ignore

    @property
    def current_grid(self) -> GridBase:
        """:class:`~pde.grids.base.GridBase`:subgrid of current MPI node."""
        return self[self.current_node]

    def _id2idx(self, node_id: int) -> tuple[int, ...]:
        """Convert linear id into node index.

        Args:
            node_id (int):
                Linear index identifying the node

        Returns:
            tuple: Full index with `num_axes` entries.
        """
        return np.unravel_index(node_id, self.shape)  # type: ignore

    def _idx2id(self, node_idx: Sequence[int]) -> int:
        """Convert node index to linear index.

        Args:
            node_idx (tuple):
                Full index with `num_axes` entries.

        Returns:
            int: Linear index identifying the node
        """
        return np.ravel_multi_index(node_idx, self.shape)  # type: ignore

    @cached_method()
    def _get_data_indices_1d(self, with_ghost_cells: bool = False) -> list[list[slice]]:
        """Indices to extract valid field data for each subgrid.

        Args:
            with_ghost_cells (bool):
                Indicates whether the ghost cells are included in `field_data`

        Returns:
            A list of slices for each axis
        """
        # create indices for each axis
        i_add = 2 if with_ghost_cells else 0
        indices_1d = []
        for axis in range(self.num_axes):
            grid_ids: list[Any] = [0] * self.num_axes
            grid_ids[axis] = slice(None, None)
            data, last = [], 0
            for grid in self.subgrids[tuple(grid_ids)]:
                n = grid.shape[axis]
                data.append(slice(last, last + n + i_add))
                last += n
            indices_1d.append(data)
        return indices_1d

    @cached_method()
    def _get_data_indices(self, with_ghost_cells: bool = False) -> NumericArray:
        """Indices to extract valid field data for each subgrid.

        Args:
            with_ghost_cells (bool):
                Indicates whether the ghost cells are included in `field_data`

        Returns:
            :class:`~numpy.ndarray` an array of indices to access data of the sub grids
        """
        indices_1d = self._get_data_indices_1d(with_ghost_cells)

        # combine everything into a full indices
        indices: NumericArray = np.empty(self.shape, dtype=object)
        for idx in np.ndindex(self.shape):
            indices[idx] = tuple(indices_1d[n][i] for n, i in enumerate(idx))

        return indices

    def get_boundary_flag(self, neighbor: int, upper: bool) -> int:
        """Get MPI flag indicating the boundary between this node and its neighbor.

        Args:
            node_id (int):
                The ID of the neighboring node
            upper (bool):
                Flag indicating the boundary of the current node

        Returns:
            int: the unique flag associated with this boundary
        """
        if upper:
            # my upper boundary (lower boundary of right cell)
            return MPIFlags.boundary_upper(self.current_node, neighbor)
        # my lower boundary (upper boundary of left cell)
        return MPIFlags.boundary_lower(self.current_node, neighbor)

    def get_neighbor(
        self, axis: int, upper: bool, *, node_id: int | None = None
    ) -> int | None:
        """Get node id of the neighbor along the given axis and direction.

        Args:
            axis (int):
                The axis of the grid
            upper (bool):
                Flag deciding which boundary to consider
            node_id (int):
                The ID of the considered node. Derived from MPI.rank() if omitted.

        Returns:
            int: The id of the neighboring node
        """
        size = self.shape[axis]
        if size == 1:
            return None  # there are no other nodes along this axis

        if node_id is None:
            node_id = self.current_node

        idx = list(self._id2idx(node_id))

        if upper:
            # my upper boundary
            if idx[axis] < size - 1:
                idx[axis] = idx[axis] + 1  # proper cell neighbor on upper side
            elif self.basegrid.periodic[axis]:
                idx[axis] = 0  # last upper cell, but periodic conditions
            else:
                return None  # no neighbor

        else:
            # my lower boundary
            if idx[axis] > 0:
                idx[axis] = idx[axis] - 1  # proper cell neighbor on lower side
            elif self.basegrid.periodic[axis]:
                idx[axis] = size - 1  # last lower cell, but periodic conditions
            else:
                return None  # no neighbor

        return self._idx2id(idx)

    def extract_field_data(
        self,
        field_data: NumericArray,
        node_id: int | None = None,
        *,
        with_ghost_cells: bool = False,
    ) -> NumericArray:
        """Extract one subfield from a global one.

        Args:
            field_data (:class:`~numpy.ndarray`):
                The field data that will be split
            node_id (int):
                Index identifying the subgrid
            with_ghost_cells (bool):
                Indicates whether the ghost cells are included in `field_data`

        Returns:
            :class:`~numpy.ndarray`: Field data defined on the subgrid. The field
            contains ghost cells if `with_ghost_cells == True`.
        """
        # consistency check
        if with_ghost_cells:
            assert field_data.shape[-self.num_axes :] == self.basegrid._shape_full
        else:
            assert field_data.shape[-self.num_axes :] == self.basegrid.shape

        if node_id is None:
            node_id = self.current_node

        node_idx = self._id2idx(node_id)
        idx = self._get_data_indices_1d(with_ghost_cells)
        i = (..., *tuple(idx[n][j] for n, j in enumerate(node_idx)))
        return field_data[i]  # type: ignore

    def extract_subfield(
        self,
        field: TField,
        node_id: int | None = None,
        *,
        with_ghost_cells: bool = False,
    ) -> TField:
        """Extract one subfield from a global field.

        Args:
            field (:class:`~pde.fields.base.DataFieldBase`):
                The field that will be split
            node_id (int):
                Index identifying the subgrid
            with_ghost_cells (bool):
                Indicates whether the ghost cells are included in data

        Returns:
            :class:`~pde.fields.base.DataFieldBase`: The sub field of the node
        """
        if node_id is None:
            node_id = self.current_node

        if isinstance(field, DataFieldBase):
            # extract data from a single field
            data = self.extract_field_data(
                field._data_full if with_ghost_cells else field.data,
                node_id,
                with_ghost_cells=with_ghost_cells,
            )

            return field.__class__(
                grid=self[node_id],
                data=data,
                label=field.label,
                dtype=field.dtype,
                with_ghost_cells=with_ghost_cells,
            )

        if isinstance(field, FieldCollection):
            # extract data from a field collection

            # extract individual fields
            fields = [
                self.extract_subfield(f, node_id, with_ghost_cells=with_ghost_cells)
                for f in field
            ]

            # combine everything to a field collection
            return field.__class__(fields, label=field.label)

        msg = f"Field type {field.__class__.__name__} unsupported"
        raise TypeError(msg)

    def extract_boundary_conditions(self, bcs_base: BoundariesBase) -> BoundariesList:
        """Extract boundary conditions for current subgrid from global conditions.

        Args:
            bcs_base (:class:`~pde.grids.boundaries.axes.BoundariesBase`):
                The boundary conditions that will be split

        Returns:
            :class:`~pde.grids.boundaries.axes.BoundariesList`: Boundary conditions of
            the subgrid
        """
        if not isinstance(bcs_base, BoundariesList):
            msg = (
                "Simulations parallelized with MPI only work with boundary conditions "
                "based on the `BoundariesList` class, i.e., conditions need to be "
                "specified for each axes separately and cannot be set using a "
                "function globally."
            )
            raise TypeError(msg)

        bcs: list[BoundaryAxisBase] = []
        for axis in range(self.num_axes):
            bcs_axis = []
            for upper in [False, True]:
                bc = bcs_base[axis][upper]
                if self.get_neighbor(axis, upper=upper) is None:
                    # extract BC for subgrid
                    bc = bc.to_subgrid(self.current_grid)
                else:
                    # set an MPI boundary condition
                    bc = _MPIBC(self, axis, upper, rank=bc.rank)
                bcs_axis.append(bc)
            bcs.append(BoundaryPair(*bcs_axis))

        return BoundariesList(bcs)

    def split_field_data_mpi(
        self, field_data: NumericArray, *, with_ghost_cells: bool = False
    ) -> NumericArray:
        """Extract one subfield from a global field.

        Args:
            field (:class:`~pde.fields.base.DataFieldBase`):
                The field that will be split. An array with the correct shape and dtype
                also needs to be passed to the receiving nodes.
            with_ghost_cells (bool):
                Indicates whether the ghost cells are included in data. If `True`,
                `field_data` must be the full data field.

        Returns:
            :class:`~pde.fields.base.DataFieldBase`: The sub field of the current node
        """
        from ..tools.mpi import mpi_recv, mpi_send

        if len(self) != mpi.size:
            msg = f"GridMesh size differs from MPI size ({len(self)} != {mpi.size})"
            raise RuntimeError(msg)

        if mpi.is_main:
            # mpi_send fields to all client processes
            for i in range(1, len(self)):
                subfield_data = self.extract_field_data(
                    field_data, i, with_ghost_cells=with_ghost_cells
                )
                mpi_send(subfield_data, i, MPIFlags.field_split)

            # extract field for the current process
            return self.extract_field_data(
                field_data, 0, with_ghost_cells=with_ghost_cells
            )

        # receive subfield from main process
        subgrid = self.current_grid

        # determine shape of resulting data
        shape = field_data.shape[: -self.num_axes]
        shape += subgrid._shape_full if with_ghost_cells else subgrid.shape

        subfield_data = np.empty(shape, dtype=field_data.dtype)
        mpi_recv(subfield_data, 0, MPIFlags.field_split)
        return subfield_data

    def split_field_mpi(self: GridMesh, field: TField) -> TField:
        """Split a field onto the subgrids by communicating data via MPI.

        The ghost cells of the returned fields will be set according to the values of
        the original field.

        Args:
            field (:class:`~pde.fields.base.DataFieldBase`):
                The field that will be split

        Results:
            :class:`~pde.fields.base.DataFieldBase`: The field defined on the subgrid
        """
        assert field.grid == self.basegrid

        if isinstance(field, DataFieldBase):
            # split individual field
            return field.__class__(
                self.current_grid,
                data=self.split_field_data_mpi(field._data_full, with_ghost_cells=True),
                label=field.label,
                dtype=field.dtype,
                with_ghost_cells=True,
            )

        if isinstance(field, FieldCollection):
            # split field collection
            field_classes = [f.__class__ for f in field]
            data = self.split_field_data_mpi(field._data_full, with_ghost_cells=True)
            return field.__class__.from_data(  # type: ignore
                field_classes,
                self.current_grid,
                data,
                label=field.label,
                labels=field.labels,
            )

        msg = f"Field type {field.__class__.__name__} unsupported"
        raise TypeError(msg)

    def combine_field_data(
        self,
        subfields: Sequence[NumericArray],
        out: NumericArray | None = None,
        *,
        with_ghost_cells: bool = False,
    ) -> NumericArray:
        """Combine data of multiple fields defined on subgrids.

        Args:
            subfields (:class:`~numpy.ndarray`):
                The data of the fields that will be combined
            out (:class:`~numpy.ndarray`):
                Full field to which the combined data is written. If `None`, a new array
                is allocated.
            with_ghost_cells (bool):
                Indicates whether the ghost cells are included in data. If `True`,
                `fields` must be the full data fields.

        Returns:
            :class:`~numpy.ndarray`: Combined field. This is `out` if out is not `None`
        """
        assert len(subfields) == len(self)

        # allocate memory to collect the data
        field0 = subfields[0]
        shape = field0.shape[: -self.num_axes]
        shape += self.basegrid._shape_full if with_ghost_cells else self.basegrid.shape
        if out is None:
            out = np.empty(shape, dtype=field0.dtype)
        else:
            assert out.shape == shape

        # collect data from all fields
        data_idx = self._get_data_indices(with_ghost_cells=with_ghost_cells)
        for i in range(len(self)):
            idx = self._id2idx(i)
            out[(...,) + data_idx[idx]] = subfields[i]

        return out

    def combine_field_data_mpi(
        self,
        subfield: NumericArray,
        out: NumericArray | None = None,
        *,
        with_ghost_cells: bool = False,
    ) -> NumericArray | None:
        """Combine data of all subfields using MPI.

        Args:
            subfield (:class:`~numpy.ndarray`):
                The data of the field defined on the current subgrid
            out (:class:`~numpy.ndarray`):
                Full field to which the combined data is written
            with_ghost_cells (bool):
                Indicates whether the ghost cells are included in data. If `True`,
                `subfield` must be the full data fields.

        Returns:
            :class:`~numpy.ndarray`: Combined field if we are on the main node. For all
            other cases, this function returns `None`. On the main node, this array is
            `out` if `out` was supplied.
        """
        from ..tools.mpi import mpi_recv, mpi_send

        assert len(self) == mpi.size

        if mpi.is_main:
            # simply copy the subfield that is on the main node
            field_data = [subfield]
            element_shape = subfield.shape[: -self.num_axes]

            # collect all subfields from all nodes
            for i in range(1, len(self)):
                if with_ghost_cells:
                    shape = element_shape + self[i]._shape_full
                else:
                    shape = element_shape + self[i].shape
                subfield_data = np.empty(shape, dtype=subfield.dtype)
                mpi_recv(subfield_data, i, MPIFlags.field_combine)
                field_data.append(subfield_data)

            # combine the data into a single field
            return self.combine_field_data(
                field_data, out=out, with_ghost_cells=with_ghost_cells
            )

        # mpi_send our subfield to the main node
        mpi_send(subfield, 0, MPIFlags.field_combine)
        return None

    def broadcast(self, data: TData) -> TData:
        """Distribute a value from the main node to all nodes.

        Args:
            data:
                The data that will be broadcasted from the main node

        Returns:
            The same data, but on all nodes
        """
        from mpi4py.MPI import COMM_WORLD

        return COMM_WORLD.bcast(data, root=0)  # type: ignore

    def gather(self, data: TData) -> list[TData] | None:
        """Gather a value from all nodes.

        Args:
            data:
                The data that will be sent to the main node

        Returns:
            None on all nodes, except the main node, which receives an ordered list with
            the data from all nodes.
        """
        from mpi4py.MPI import COMM_WORLD

        return COMM_WORLD.gather(data, root=0)

    def allgather(self, data: TData) -> list[TData]:
        """Gather a value from reach node and sends them to all nodes.

        Args:
            data:
                The data that will be sent to the main node

        Returns:
            list: data from all nodes.
        """
        from mpi4py.MPI import COMM_WORLD

        return COMM_WORLD.allgather(data)

    @plot_on_axes()
    def plot(self, ax, **kwargs) -> None:
        r"""Visualize the grid mesh.

        Args:
            {PLOT_ARGS}
            \**kwargs:
                Extra arguments are passed on the to the matplotlib plotting routines,
                e.g., to set the color of the lines
        """
        if self.num_axes not in {1, 2}:
            msg = f"Cannot plot data of dimension {self.basegrid.dim}"
            raise NotImplementedError(msg)

        kwargs.setdefault("color", "k")
        for x in np.arange(self.shape[0] + 1) - 0.5:
            ax.axvline(x, **kwargs)
        ax.set_xlim(-0.5, self.shape[0] - 0.5)
        ax.set_xlabel(self.subgrids.flat[0].axes[0])

        if self.num_axes == 2:
            for y in np.arange(self.shape[1] + 1) - 0.5:
                ax.axhline(y, **kwargs)
            ax.set_ylim(-0.5, self.shape[1] - 0.5)
            ax.set_ylabel(self.subgrids.flat[0].axes[1])

            ax.set_aspect(1)

        for num, idx in enumerate(np.ndindex(self.shape)):
            pos = (idx, 0) if self.num_axes == 1 else idx
            ax.text(
                *pos, str(num), horizontalalignment="center", verticalalignment="center"
            )
