"""Defines the FDTD grid."""
from __future__ import annotations
from typing import Tuple, List, Union

import numpy as np
import pydantic.v1 as pd

from ..base import Tidy3dBaseModel
from ..data.data_array import DataArray, SpatialDataArray, ScalarFieldDataArray
from ..data.dataset import UnstructuredGridDatasetType, UnstructuredGridDataset
from ..types import ArrayFloat1D, Axis, InterpMethod, Literal
from ..geometry.base import Box

from ...exceptions import SetupError

# data type of one dimensional coordinate array.
Coords1D = ArrayFloat1D



[docs]
class Coords(Tidy3dBaseModel):
    """Holds data about a set of x,y,z positions on a grid.

    Example
    -------
    >>> x = np.linspace(-1, 1, 10)
    >>> y = np.linspace(-1, 1, 11)
    >>> z = np.linspace(-1, 1, 12)
    >>> coords = Coords(x=x, y=y, z=z)
    """

    x: Coords1D = pd.Field(
        ..., title="X Coordinates", description="1-dimensional array of x coordinates."
    )

    y: Coords1D = pd.Field(
        ..., title="Y Coordinates", description="1-dimensional array of y coordinates."
    )

    z: Coords1D = pd.Field(
        ..., title="Z Coordinates", description="1-dimensional array of z coordinates."
    )

    @property
    def to_dict(self):
        """Return a dict of the three Coord1D objects as numpy arrays."""
        return {key: self.dict()[key] for key in "xyz"}

    @property
    def to_list(self):
        """Return a list of the three Coord1D objects as numpy arrays."""
        return list(self.to_dict.values())

    def _interp_from_xarray(
        self,
        array: Union[SpatialDataArray, ScalarFieldDataArray],
        interp_method: InterpMethod,
        fill_value: Union[Literal["extrapolate"], float] = "extrapolate",
    ) -> Union[SpatialDataArray, ScalarFieldDataArray]:
        """
        Similar to ``xarrray.DataArray.interp`` with 2 enhancements:

            1) Check if the coordinate of the supplied data are in monotonically increasing order.
            If they are, apply the faster ``assume_sorted=True``.

            2) For axes of single entry, instead of error, apply ``isel()`` along the axis.

        Parameters
        ----------
        array : Union[:class:`.SpatialDataArray`, :class:`.ScalarFieldDataArray`]
            Supplied scalar dataset
        interp_method : :class:`.InterpMethod`
            Interpolation method.
        fill_value : Union[Literal['extrapolate'], float] = "extrapolate"
            Value used to fill in for points outside the data range. If set to 'extrapolate',
            values will be extrapolated into those regions using the "nearest" method.

        Returns
        -------
        Union[:class:`.SpatialDataArray`, :class:`.ScalarFieldDataArray`]
            The interpolated spatial dataset.

        Note
        ----
        This method is called from a :class:`Coords` instance with the array to be interpolated as
        an argument, not the other way around.
        """
        # Check which axes need interpolation or selection
        interp_ax = []
        isel_ax = []
        for ax in "xyz":
            if array.sizes[ax] == 1:
                isel_ax.append(ax)
            else:
                interp_ax.append(ax)

        # apply iselection for the axis containing single entry
        if len(isel_ax) > 0:
            array = array.isel({ax: [0] * len(self.to_dict[ax]) for ax in isel_ax})
            array = array.assign_coords({ax: self.to_dict[ax] for ax in isel_ax})
            if len(interp_ax) == 0:
                return array

        # Apply interp for the rest
        is_sorted = all(np.all(np.diff(array.coords[f]) > 0) for f in interp_ax)
        interp_param = {
            "method": interp_method,
            "assume_sorted": is_sorted,
            "kwargs": {
                "bounds_error": False,
                "fill_value": fill_value,
            },
        }

        # Mark extrapolated points with nan's to fill in later
        if fill_value == "extrapolate" and interp_method != "nearest":
            interp_param["kwargs"]["fill_value"] = np.nan

        # interpolation
        interp_array = array.interp({ax: self.to_dict[ax] for ax in interp_ax}, **interp_param)

        # Fill in nan's with nearest values
        if fill_value == "extrapolate" and interp_method != "nearest":
            interp_param["method"] = "nearest"
            interp_param["kwargs"]["fill_value"] = "extrapolate"
            nearest_array = array.interp({ax: self.to_dict[ax] for ax in interp_ax}, **interp_param)
            interp_array.values[:] = np.where(
                np.isnan(interp_array.values), nearest_array.values, interp_array.values
            )

        return interp_array

    def _interp_from_unstructured(
        self,
        array: UnstructuredGridDatasetType,
        interp_method: InterpMethod,
        fill_value: Union[Literal["extrapolate"], float] = "extrapolate",
    ) -> SpatialDataArray:
        """
        Interpolate from untructured grid onto a Cartesian one.

        Parameters
        ----------
        array : Union[class:`.TriangularGridDataset`, class:`.TetrahedralGridDataset`]
            Supplied scalar dataset
        interp_method : :class:`.InterpMethod`
            Interpolation method.
        fill_value : Union[Literal['extrapolate'], float] = "extrapolate"
            Value used to fill in for points outside the data range. If set to 'extrapolate',
            values will be extrapolated into those regions using the "nearest" method.

        Returns
        -------
        :class:`.SpatialDataArray`
            The interpolated spatial dataset.

        Note
        ----
        This method is called from a :class:`Coords` instance with the array to be interpolated as
        an argument, not the other way around.
        """
        interp_array = array.interp(
            **{ax: self.to_dict[ax] for ax in "xyz"}, method=interp_method, fill_value=fill_value
        )

        return interp_array


[docs]
    def spatial_interp(
        self,
        array: Union[SpatialDataArray, ScalarFieldDataArray, UnstructuredGridDatasetType],
        interp_method: InterpMethod,
        fill_value: Union[Literal["extrapolate"], float] = "extrapolate",
    ) -> Union[SpatialDataArray, ScalarFieldDataArray]:
        """
        Similar to ``xarrray.DataArray.interp`` with 2 enhancements:

            1) (if input data is an ``xarrray.DataArray``) Check if the coordinate of the supplied
            data are in monotonically increasing order. If they are, apply the faster
            ``assume_sorted=True``.

            2) Data is assumed invariant along zero-size dimensions (if any).

        Parameters
        ----------
        array : Union[
                :class:`.SpatialDataArray`,
                :class:`.ScalarFieldDataArray`,
                :class:`.TriangularGridDataset`,
                :class:`.TetrahedralGridDataset`,
        ]
            Supplied scalar dataset
        interp_method : :class:`.InterpMethod`
            Interpolation method.
        fill_value : Union[Literal['extrapolate'], float] = "extrapolate"
            Value used to fill in for points outside the data range. If set to 'extrapolate',
            values will be extrapolated into those regions using the "nearest" method.

        Returns
        -------
        Union[:class:`.SpatialDataArray`, :class:`.ScalarFieldDataArray`]
            The interpolated spatial dataset.

        Note
        ----
        This method is called from a :class:`Coords` instance with the array to be interpolated as
        an argument, not the other way around.
        """

        # Check for empty dimensions
        result_coords = dict(self.to_dict)
        if any(len(v) == 0 for v in result_coords.values()):
            if isinstance(array, (SpatialDataArray, ScalarFieldDataArray)):
                for c in array.coords:
                    if c not in result_coords:
                        result_coords[c] = array.coords[c].values
            result_shape = tuple(len(v) for v in result_coords.values())
            result = DataArray(np.empty(result_shape, dtype=array.dtype), coords=result_coords)
            return result

        # interpolation
        if isinstance(array, UnstructuredGridDataset):
            return self._interp_from_unstructured(
                array=array, interp_method=interp_method, fill_value=fill_value
            )
        else:
            return self._interp_from_xarray(
                array=array, interp_method=interp_method, fill_value=fill_value
            )





[docs]
class FieldGrid(Tidy3dBaseModel):
    """Holds the grid data for a single field.

    Example
    -------
    >>> x = np.linspace(-1, 1, 10)
    >>> y = np.linspace(-1, 1, 11)
    >>> z = np.linspace(-1, 1, 12)
    >>> coords = Coords(x=x, y=y, z=z)
    >>> field_grid = FieldGrid(x=coords, y=coords, z=coords)
    """

    x: Coords = pd.Field(
        ...,
        title="X Positions",
        description="x,y,z coordinates of the locations of the x-component of a vector field.",
    )

    y: Coords = pd.Field(
        ...,
        title="Y Positions",
        description="x,y,z coordinates of the locations of the y-component of a vector field.",
    )

    z: Coords = pd.Field(
        ...,
        title="Z Positions",
        description="x,y,z coordinates of the locations of the z-component of a vector field.",
    )




[docs]
class YeeGrid(Tidy3dBaseModel):
    """Holds the yee grid coordinates for each of the E and H positions.

    Example
    -------
    >>> x = np.linspace(-1, 1, 10)
    >>> y = np.linspace(-1, 1, 11)
    >>> z = np.linspace(-1, 1, 12)
    >>> coords = Coords(x=x, y=y, z=z)
    >>> field_grid = FieldGrid(x=coords, y=coords, z=coords)
    >>> yee_grid = YeeGrid(E=field_grid, H=field_grid)
    >>> Ex_coords = yee_grid.E.x
    """

    E: FieldGrid = pd.Field(
        ...,
        title="Electric Field Grid",
        description="Coordinates of the locations of all three components of the electric field.",
    )

    H: FieldGrid = pd.Field(
        ...,
        title="Electric Field Grid",
        description="Coordinates of the locations of all three components of the magnetic field.",
    )

    @property
    def grid_dict(self):
        """The Yee grid coordinates associated to various field components as a dictionary."""
        return {
            "Ex": self.E.x,
            "Ey": self.E.y,
            "Ez": self.E.z,
            "Hx": self.H.x,
            "Hy": self.H.y,
            "Hz": self.H.z,
        }




[docs]
class Grid(Tidy3dBaseModel):
    """Contains all information about the spatial positions of the FDTD grid.

    Example
    -------
    >>> x = np.linspace(-1, 1, 10)
    >>> y = np.linspace(-1, 1, 11)
    >>> z = np.linspace(-1, 1, 12)
    >>> coords = Coords(x=x, y=y, z=z)
    >>> grid = Grid(boundaries=coords)
    >>> centers = grid.centers
    >>> sizes = grid.sizes
    >>> yee_grid = grid.yee
    """

    boundaries: Coords = pd.Field(
        ...,
        title="Boundary Coordinates",
        description="x,y,z coordinates of the boundaries between cells, defining the FDTD grid.",
    )

    @staticmethod
    def _avg(coords1d: Coords1D):
        """Return average positions of an array of 1D coordinates."""
        return (coords1d[1:] + coords1d[:-1]) / 2.0

    @staticmethod
    def _min(coords1d: Coords1D):
        """Return minus positions of 1D coordinates."""
        return coords1d[:-1]

    @property
    def centers(self) -> Coords:
        """Return centers of the cells in the :class:`Grid`.

        Returns
        -------
        :class:`Coords`
            centers of the FDTD cells in x,y,z stored as :class:`Coords` object.

        Example
        -------
        >>> x = np.linspace(-1, 1, 10)
        >>> y = np.linspace(-1, 1, 11)
        >>> z = np.linspace(-1, 1, 12)
        >>> coords = Coords(x=x, y=y, z=z)
        >>> grid = Grid(boundaries=coords)
        >>> centers = grid.centers
        """
        return Coords(**{key: self._avg(val) for key, val in self.boundaries.to_dict.items()})

    @property
    def sizes(self) -> Coords:
        """Return sizes of the cells in the :class:`Grid`.

        Returns
        -------
        :class:`Coords`
            Sizes of the FDTD cells in x,y,z stored as :class:`Coords` object.

        Example
        -------
        >>> x = np.linspace(-1, 1, 10)
        >>> y = np.linspace(-1, 1, 11)
        >>> z = np.linspace(-1, 1, 12)
        >>> coords = Coords(x=x, y=y, z=z)
        >>> grid = Grid(boundaries=coords)
        >>> sizes = grid.sizes
        """
        return Coords(**{key: np.diff(val) for key, val in self.boundaries.to_dict.items()})

    @property
    def num_cells(self) -> Tuple[int, int, int]:
        """Return sizes of the cells in the :class:`Grid`.

        Returns
        -------
        tuple[int, int, int]
            Number of cells in the grid in the x, y, z direction.

        Example
        -------
        >>> x = np.linspace(-1, 1, 10)
        >>> y = np.linspace(-1, 1, 11)
        >>> z = np.linspace(-1, 1, 12)
        >>> coords = Coords(x=x, y=y, z=z)
        >>> grid = Grid(boundaries=coords)
        >>> Nx, Ny, Nz = grid.num_cells
        """
        return [len(self.boundaries.dict()[dim]) - 1 for dim in "xyz"]

    @property
    def _primal_steps(self) -> Coords:
        """Return primal steps of the cells in the :class:`Grid`.

        Returns
        -------
        :class:`Coords`
            Distances between each of the cell boundaries along each dimension.
        """
        return self.sizes

    @property
    def _dual_steps(self) -> Coords:
        """Return dual steps of the cells in the :class:`Grid`.

        Returns
        -------
        :class:`Coords`
            Distances between each of the cell centers along each dimension, with periodicity
            applied.
        """

        primal_steps = {dim: self._primal_steps.dict()[dim] for dim in "xyz"}
        dsteps = {key: (psteps + np.roll(psteps, 1)) / 2 for (key, psteps) in primal_steps.items()}

        return Coords(**dsteps)

    @property
    def yee(self) -> YeeGrid:
        """Return the :class:`YeeGrid` defining the yee cell locations for this :class:`Grid`.


        Returns
        -------
        :class:`YeeGrid`
            Stores coordinates of all of the components on the yee lattice.

        Example
        -------
        >>> x = np.linspace(-1, 1, 10)
        >>> y = np.linspace(-1, 1, 11)
        >>> z = np.linspace(-1, 1, 12)
        >>> coords = Coords(x=x, y=y, z=z)
        >>> grid = Grid(boundaries=coords)
        >>> yee_cells = grid.yee
        >>> Ex_positions = yee_cells.E.x
        """
        yee_e_kwargs = {key: self._yee_e(axis=axis) for axis, key in enumerate("xyz")}
        yee_h_kwargs = {key: self._yee_h(axis=axis) for axis, key in enumerate("xyz")}

        yee_e = FieldGrid(**yee_e_kwargs)
        yee_h = FieldGrid(**yee_h_kwargs)
        return YeeGrid(E=yee_e, H=yee_h)


[docs]
    def __getitem__(self, coord_key: str) -> Coords:
        """quickly get the grid element by grid[key]."""

        coord_dict = {
            "centers": self.centers,
            "sizes": self.sizes,
            "boundaries": self.boundaries,
            "Ex": self.yee.E.x,
            "Ey": self.yee.E.y,
            "Ez": self.yee.E.z,
            "Hx": self.yee.H.x,
            "Hy": self.yee.H.y,
            "Hz": self.yee.H.z,
        }
        if coord_key not in coord_dict:
            raise SetupError(f"key {coord_key} not found in grid with {list(coord_dict.keys())} ")

        return coord_dict.get(coord_key)


    def _yee_e(self, axis: Axis):
        """E field yee lattice sites for axis."""

        boundary_coords = self.boundaries.to_dict

        # initially set all to the minus bounds
        yee_coords = {key: self._min(val) for key, val in boundary_coords.items()}

        # average the axis index between the cell boundaries
        key = "xyz"[axis]
        yee_coords[key] = self._avg(boundary_coords[key])

        return Coords(**yee_coords)

    def _yee_h(self, axis: Axis):
        """H field yee lattice sites for axis."""

        boundary_coords = self.boundaries.to_dict

        # initially set all to centers
        yee_coords = {key: self._avg(val) for key, val in boundary_coords.items()}

        # set the axis index to the minus bounds
        key = "xyz"[axis]
        yee_coords[key] = self._min(boundary_coords[key])

        return Coords(**yee_coords)


[docs]
    def discretize_inds(self, box: Box, extend: bool = False) -> List[Tuple[int, int]]:
        """Start and stopping indexes for the cells that intersect with a :class:`Box`.

        Parameters
        ----------
        box : :class:`Box`
            Rectangular geometry within simulation to discretize.
        extend : bool = False
            If ``True``, ensure that the returned indexes extend sufficiently in every direction to
            be able to interpolate any field component at any point within the ``box``, for field
            components sampled on the Yee grid.

        Returns
        -------
        List[Tuple[int, int]]
            The (start, stop) indexes of the cells that intersect with ``box`` in each of the three
            dimensions.
        """

        pts_min, pts_max = box.bounds
        boundaries = self.boundaries

        inds_list = []

        # for each dimension
        for axis, (pt_min, pt_max) in enumerate(zip(pts_min, pts_max)):
            bound_coords = np.array(boundaries.to_list[axis])
            assert pt_min <= pt_max, "min point was greater than max point"

            # index of smallest coord greater than pt_max
            inds_gt_pt_max = np.where(bound_coords > pt_max)[0]
            ind_max = len(bound_coords) - 1 if len(inds_gt_pt_max) == 0 else inds_gt_pt_max[0]

            # index of largest coord less than or equal to pt_min
            inds_leq_pt_min = np.where(bound_coords <= pt_min)[0]
            ind_min = 0 if len(inds_leq_pt_min) == 0 else inds_leq_pt_min[-1]

            # handle extensions
            if ind_max > ind_min and extend:
                # Left side
                if box.bounds[0][axis] < self.centers.to_list[axis][ind_min]:
                    # Box bounds on the left side are to the left of the closest grid center
                    ind_min -= 1

                # We always need an extra pixel on the right for the tangential components
                ind_max += 1

            # store indexes
            inds_list.append((ind_min, ind_max))

        return inds_list



[docs]
    def extended_subspace(
        self,
        axis: Axis,
        ind_beg: int = 0,
        ind_end: int = 0,
        periodic: bool = True,
    ) -> Coords1D:
        """Pick a subspace of 1D boundaries within ``range(ind_beg, ind_end)``. If any indexes lie
        outside of the grid boundaries array, padding is used based on the boundary conditions.
        For periodic BCs, the zeroth and last element of the grid boundaries are identified.
        For other BCs, the zeroth and last element of the boundaries are a reflection plane.

        Parameters
        ----------
        axis : Axis
            Axis index along which to pick the subspace.
        ind_beg : int = 0
            Starting index for the subspace.
        ind_end : int = 0
            Ending index for the subspace.
        periodic : bool = True
            Whether to pad out of bounds indexes with a periodic or reflected pattern.

        Returns
        -------
        Coords1D
            The subspace of the grid along ``axis``.
        """

        coords = self.boundaries.to_list[axis]
        padded_coords = coords
        num_cells = coords.size - 1

        reverse = True
        while ind_beg < 0:
            if periodic or not reverse:
                offset = padded_coords[0] - coords[-1]
                padded_coords = np.concatenate([coords[:-1] + offset, padded_coords])
                reverse = True
            else:
                offset = padded_coords[0] + coords[0]
                padded_coords = np.concatenate([offset - coords[:0:-1], padded_coords])
                reverse = False
            ind_beg += num_cells
            ind_end += num_cells

        reverse = True
        while ind_end >= padded_coords.size:
            if periodic or not reverse:
                offset = padded_coords[-1] - coords[0]
                padded_coords = np.concatenate([padded_coords, coords[1:] + offset])
                reverse = True
            else:
                offset = padded_coords[-1] + coords[-1]
                padded_coords = np.concatenate([padded_coords, offset - coords[-2::-1]])
                reverse = False

        return padded_coords[ind_beg:ind_end]



[docs]
    def snap_to_box_zero_dim(self, box: Box):
        """Snap a grid to an exact box position for dimensions for which the box is size 0.
        If the box location is outside of the grid, an error is raised.

        Parameters
        ----------
        box : :class:`Box`
            Box to use for the zero dim check.

        Returns
        -------
        class:`Grid`
            Snapped copy of the grid.
        """

        boundary_dict = self.boundaries.to_dict.copy()
        for dim, center, size in zip("xyz", box.center, box.size):
            # Overwrite grid boundaries with box center if box is size 0 along dimension
            if size == 0:
                if boundary_dict[dim][0] > center or boundary_dict[dim][-1] < center:
                    raise ValueError("Cannot snap grid to box center outside of grid domain.")
                boundary_dict[dim] = np.array([center, center])
        return self.updated_copy(boundaries=Coords(**boundary_dict))