"""Defines a vectorial field over a grid.

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

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Literal, overload

import numpy as np

from ..grids.base import DimensionError, GridBase
from ..grids.cartesian import CartesianGrid
from ..tools.docstrings import fill_in_docstring
from ..tools.misc import get_common_dtype
from ..tools.plotting import PlotReference, plot_on_figure
from .datafield_base import DataFieldBase
from .scalar import ScalarField

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence

    from numpy.typing import DTypeLike

    from ..grids.boundaries.axes import BoundariesData
    from ..tools.typing import Number, NumberOrArray, NumericArray
    from .tensorial import Tensor2Field


class VectorField(DataFieldBase):
    """Vector field discretized on a grid.

    Warning:
        Components of the vector field are given in the local basis. While the local
        basis is identical to the global basis in Cartesian coordinates, the local basis
        depends on position in curvilinear coordinate systems. Moreover, the field
        always contains all components, even if the underlying grid assumes symmetries.
    """

    rank = 1

    @classmethod
    def from_scalars(
        cls,
        fields: list[ScalarField],
        *,
        label: str | None = None,
        dtype: DTypeLike | None = None,
    ) -> VectorField:
        """Create a vector field from a list of ScalarFields.

        Note that the data of the scalar fields is copied in the process

        Args:
            fields (list):
                The list of (compatible) scalar fields
            label (str, optional):
                Name of the returned field
            dtype (numpy dtype):
                The data type of the field. If omitted, it will be determined from
                `data` automatically.

        Returns:
            :class:`VectorField`: the resulting vector field
        """
        grid = fields[0].grid

        if grid.dim != len(fields):
            msg = (
                f"Grid dimension and number of scalar fields differ ({grid.dim} != "
                f"{len(fields)})"
            )
            raise DimensionError(msg)

        data = []
        for field in fields:
            if not field.grid.compatible_with(grid):
                msg = "Grids are incompatible"
                raise ValueError(msg)
            data.append(field.data)

        return cls(grid, data, label=label, dtype=dtype)

    @classmethod
    @fill_in_docstring
    def from_expression(
        cls,
        grid: GridBase,
        expressions: Sequence[str],
        *,
        user_funcs: dict[str, Callable] | None = None,
        consts: dict[str, NumberOrArray] | None = None,
        label: str | None = None,
        dtype: DTypeLike | None = None,
    ) -> VectorField:
        """Create a vector field on a grid from given expressions.

        Warning:
            {WARNING_EXEC}

        Args:
            grid (:class:`~pde.grids.base.GridBase`):
                Grid defining the space on which this field is defined
            expressions (list of str):
                A list of mathematical expression, one for each component of the vector
                field. The expressions determine the values as a function of the
                position on the grid. The expressions may contain standard mathematical
                functions and they may depend on the axes labels of the grid.
                More information can be found in the
                :ref:`expression documentation <documentation-expressions>`.
            user_funcs (dict, optional):
                A dictionary with user defined functions that can be used in the
                expression
            consts (dict, optional):
                A dictionary with user defined constants that can be used in the
                expression. The values of these constants should either be numbers or
                :class:`~numpy.ndarray`.
            label (str, optional):
                Name of the field
            dtype (numpy dtype):
                The data type of the field. If omitted, it will be determined from
                `data` automatically.
        """
        from ..tools.expressions import ScalarExpression

        if isinstance(expressions, str) or len(expressions) != grid.dim:
            axes_names = grid.axes + grid.axes_symmetric
            msg = f"Expected {grid.dim} expressions for the coordinates {axes_names}."
            raise DimensionError(msg)

        if any("cartesian" in str(expression) for expression in expressions):
            # support Cartesian coordinates via a special constant
            if consts is None:
                consts = {}
            if "cartesian" not in consts:
                coords_cart = grid.point_to_cartesian(grid.cell_coords)
                consts["cartesian"] = np.moveaxis(coords_cart, -1, 0)
            assert "cartesian" in consts

        # obtain the coordinates of the grid points
        points = [grid.cell_coords[..., i] for i in range(grid.num_axes)]

        # evaluate all vector components at all points
        data = []
        for expression in expressions:
            expr = ScalarExpression(
                expression=expression,
                signature=grid.axes,
                user_funcs=user_funcs,
                consts=consts,
                repl=grid.c._axes_alt_repl,
                allow_indexed=True,
            )
            values = np.broadcast_to(expr(*points), grid.shape)
            data.append(values)

        # create vector field from the data
        return cls(grid=grid, data=data, label=label, dtype=dtype)

    def __getitem__(self, key: int | str) -> ScalarField:
        """Extract a component of the VectorField."""
        axis = self.grid.get_axis_index(key)
        comp_name = self.grid.c.axes[axis]
        if self.label:
            label = self.label + f"_{comp_name}"
        else:
            label = f"{comp_name} component"
        return ScalarField(
            self.grid, data=self._data_full[axis], label=label, with_ghost_cells=True
        )

    def __setitem__(self, key: int | str, value: NumberOrArray | ScalarField):
        """Set a component of the VectorField."""
        idx = self.grid.get_axis_index(key)
        if isinstance(value, ScalarField):
            self.grid.assert_grid_compatible(value.grid)
            self.data[idx] = value.data
        else:
            self.data[idx] = value

    @overload
    def dot(
        self,
        other: VectorField,
        out: ScalarField | None = ...,
        *,
        conjugate: bool = ...,
        label: str = ...,
    ) -> ScalarField: ...

    @overload
    def dot(
        self,
        other: Tensor2Field,
        out: VectorField | None = ...,
        *,
        conjugate: bool = ...,
        label: str = ...,
    ) -> VectorField: ...

    def dot(
        self,
        other: VectorField | Tensor2Field,
        out: ScalarField | VectorField | None = None,
        *,
        conjugate: bool = True,
        label: str = "dot product",
    ) -> ScalarField | VectorField:
        """Calculate the dot product involving a vector field.

        This supports the dot product between two vectors fields as well as the
        product between a vector and a tensor. The resulting fields will be a
        scalar or vector, respectively.

        Args:
            other (VectorField or Tensor2Field):
                the second field
            out (ScalarField or VectorField, optional):
                Optional field to which the result is written.
            conjugate (bool):
                Whether to use the complex conjugate for the second operand
            label (str, optional):
                Name of the returned field

        Returns:
            :class:`~pde.fields.scalar.ScalarField` or
            :class:`~pde.fields.vectorial.VectorField`: result of applying the operator
        """
        from .tensorial import Tensor2Field

        # check input
        self.grid.assert_grid_compatible(other.grid)
        if isinstance(other, VectorField):
            result_type = ScalarField
        elif isinstance(other, Tensor2Field):
            result_type = VectorField  # type: ignore
        else:
            msg = "Second term must be a vector or tensor field"
            raise TypeError(msg)

        if out is None:
            out = result_type(self.grid, dtype=get_common_dtype(self, other))
        else:
            if not isinstance(out, result_type):
                msg = f"`out` must be of type `{result_type}`"
                raise TypeError(msg)
            self.grid.assert_grid_compatible(out.grid)

        # calculate the result
        other_data = other.data.conjugate() if conjugate else other.data
        np.einsum("i...,i...->...", self.data, other_data, out=out.data)
        if label is not None:
            out.label = label

        return out

    __matmul__ = dot  # support python @-syntax for matrix multiplication

    def outer_product(
        self,
        other: VectorField,
        out: Tensor2Field | None = None,
        *,
        label: str | None = None,
    ) -> Tensor2Field:
        """Calculate the outer product of this vector field with another.

        Args:
            other (:class:`~pde.fields.vectorial.VectorField`):
                The second vector field
            out (:class:`~pde.fields.tensorial.Tensor2Field`, optional):
                Optional tensorial field to which the  result is written.
            label (str, optional):
                Name of the returned field

        Returns:
            :class:`~pde.fields.tensorial.Tensor2Field`: result of the operation
        """
        from .tensorial import Tensor2Field

        self.assert_field_compatible(other)

        if out is None:
            out = Tensor2Field(self.grid)
        else:
            self.grid.assert_grid_compatible(out.grid)

        # calculate the result
        np.einsum("i...,j...->ij...", self.data, other.data, out=out.data)
        if label is not None:
            out.label = label

        return out

    def make_outer_prod_operator(
        self, backend: Literal["numpy", "numba"] = "numba"
    ) -> Callable[[NumericArray, NumericArray, NumericArray | None], NumericArray]:
        """Return operator calculating the outer product of two vector fields.

        Warning:
            This function does not check types or dimensions.

        Args:
            backend (str):
                The backend (e.g., 'numba' or 'numba') used for this operator.

        Returns:
            function that takes two instance of :class:`~numpy.ndarray`, which contain
            the discretized data of the two operands. An optional third argument can
            specify the output array to which the result is written.
        """
        from ..backends import backends

        return backends[backend].make_outer_prod_operator(self)

    @fill_in_docstring
    def divergence(
        self, bc: BoundariesData | None, out: ScalarField | None = None, **kwargs
    ) -> ScalarField:
        """Apply divergence operator and return result as a field.

        Args:
            bc:
                The boundary conditions applied to the field.
                {ARG_BOUNDARIES_OPTIONAL}
            out (ScalarField, optional):
                Optional scalar field to which the  result is written.
            label (str, optional):
                Name of the returned field
            **kwargs:
                Additional arguments affecting how the operator behaves.

        Returns:
            :class:`~pde.fields.scalar.ScalarField`: Divergence of the field
        """
        return self.apply_operator("divergence", bc=bc, out=out, **kwargs)  # type: ignore

    @fill_in_docstring
    def gradient(
        self,
        bc: BoundariesData | None,
        out: Tensor2Field | None = None,
        **kwargs,
    ) -> Tensor2Field:
        r"""Apply vector gradient operator and return result as a field.

        The vector gradient field is a tensor field :math:`t_{\alpha\beta}` that
        specifies the derivatives of the vector field :math:`v_\alpha` with respect to
        all coordinates :math:`x_\beta`.

        Args:
            bc:
                The boundary conditions applied to the field. Boundary conditions need
                to determine all components of the vector field.
                {ARG_BOUNDARIES_OPTIONAL}
            out (VectorField, optional):
                Optional vector field to which the result is written.
            label (str, optional):
                Name of the returned field
            **kwargs:
                Additional arguments affecting how the operator behaves.

        Returns:
            :class:`~pde.fields.tensorial.Tensor2Field`: Gradient of the field
        """
        return self.apply_operator("vector_gradient", bc=bc, out=out, **kwargs)  # type: ignore

    @fill_in_docstring
    def laplace(
        self, bc: BoundariesData | None, out: VectorField | None = None, **kwargs
    ) -> VectorField:
        r"""Apply vector Laplace operator and return result as a field.

        The vector Laplacian is a vector field :math:`L_\alpha` containing the second
        derivatives of the vector field :math:`v_\alpha` with respect to the coordinates
        :math:`x_\beta`:

        .. math::
            L_\alpha = \sum_\beta
                \frac{\partial^2 v_\alpha}{\partial x_\beta \partial x_\beta}

        Args:
            bc:
                The boundary conditions applied to the field.
                {ARG_BOUNDARIES_OPTIONAL}
            out (VectorField, optional):
                Optional vector field to which the  result is written.
            label (str, optional):
                Name of the returned field
            **kwargs:
                Additional arguments affecting how the operator behaves.

        Returns:
            :class:`~pde.fields.vectorial.VectorField`: Laplacian of the field
        """
        return self.apply_operator("vector_laplace", bc=bc, out=out, **kwargs)  # type: ignore

    def to_scalar(
        self,
        scalar: str | int = "auto",
        *,
        label: str | None = "scalar `{scalar}`",
    ) -> ScalarField:
        """Return scalar variant of the field.

        Args:
            scalar (str):
                Choose the method to use. Possible  choices are `norm`, `max`, `min`,
                `squared_sum`, `norm_squared`, or an integer specifying which component
                is returned (indexing starts at `0`). The default value `auto` picks the
                method automatically: The first (and only) component is returned for
                real fields on one-dimensional spaces, while the norm of the vector is
                returned otherwise.
            label (str, optional):
                Name of the returned field

        Returns:
            :class:`pde.fields.scalar.ScalarField`:
                The scalar field after applying the operation
        """
        if scalar == "auto":
            if self.grid.dim > 1 or np.iscomplexobj(self.data):
                scalar = "norm"
            else:
                scalar = 0  # return the field unchanged

        if isinstance(scalar, int):
            data = self.data[scalar]

        elif scalar == "norm":
            data = np.linalg.norm(self.data, axis=0)

        elif scalar == "max":
            data = np.max(self.data, axis=0)

        elif scalar == "min":
            data = np.min(self.data, axis=0)

        elif scalar == "squared_sum":
            data = np.sum(self.data**2, axis=0)

        elif scalar == "norm_squared":
            data = np.sum(self.data * self.data.conjugate(), axis=0)

        else:
            msg = f"Unknown method `{scalar}` for `to_scalar`"
            raise ValueError(msg)

        if label is not None:
            label = label.format(scalar=scalar)

        return ScalarField(self.grid, data, label=label)

    def get_vector_data(
        self, transpose: bool = False, max_points: int | None = None, **kwargs
    ) -> dict[str, Any]:
        r"""Return data for a vector plot of the field.

        Args:
            transpose (bool):
                Determines whether the transpose of the data should be plotted.
            max_points (int):
                The maximal number of points that is used along each axis. This
                option can be used to sub-sample the data.
            \**kwargs:
                Additional parameters forwarded to `grid.get_image_data`

        Returns:
            dict: Information useful for plotting an vector field
        """
        if self.is_complex:
            self._logger.warning("Only the real part of complex data is shown")

        # extract the image data
        data = self.grid.get_vector_data(self.data.real, **kwargs)
        data["title"] = self.label

        # transpose the data if requested
        if transpose:
            data["x"], data["y"] = data["y"], data["x"]
            data["data_x"], data["data_y"] = data["data_y"].T, data["data_x"].T
            data["label_x"], data["label_y"] = data["label_y"], data["label_x"]
            data["extent"] = data["extent"][2:] + data["extent"][:2]

        # reduce the sampling of the vector points
        if max_points is not None:
            shape = data["data_x"].shape
            for axis, size in enumerate(shape):
                if size > max_points:
                    # sub-sample the data
                    idx_f = np.linspace(0, size - 1, max_points)
                    idx_i = np.round(idx_f).astype(int)

                    data["data_x"] = np.take(data["data_x"], idx_i, axis=axis)
                    data["data_y"] = np.take(data["data_y"], idx_i, axis=axis)
                    if axis == 0:
                        data["x"] = data["x"][idx_i]
                    elif axis == 1:
                        data["y"] = data["y"][idx_i]
                    else:
                        msg = "Only supports 2d grids"
                        raise RuntimeError(msg)

        data["shape"] = data["data_x"].shape
        data["size"] = data["data_x"].size

        return data

    @fill_in_docstring
    def interpolate_to_grid(
        self: VectorField,
        grid: GridBase,
        *,
        bc: BoundariesData | None = None,
        fill: Number | None = None,
        label: str | None = None,
    ) -> VectorField:
        """Interpolate the data of this vector field to another grid.

        Args:
            grid (:class:`~pde.grids.base.GridBase`):
                The grid of the new field onto which the current field is interpolated.
            bc:
                The boundary conditions applied to the field, which affects values close
                to the boundary. If omitted, the argument `fill` is used to determine
                values outside the domain.
                {ARG_BOUNDARIES_OPTIONAL}
            fill (Number, optional):
                Determines how values out of bounds are handled. If `None`, a
                `ValueError` is raised when out-of-bounds points are requested.
                Otherwise, the given value is returned.
            label (str, optional):
                Name of the returned field

        Returns:
            Field of the same rank as the current one.
        """
        if self.grid.dim != grid.dim:
            msg = f"Incompatible grid dimensions ({self.grid.dim:d} != {grid.dim:d})"
            raise DimensionError(msg)

        # determine the points at which data needs to be calculated
        if (
            self.grid.num_axes == grid.num_axes
            and self.grid.__class__ is grid.__class__
        ) or (  # UnitGrid and CartesianGrid should be treated the same here
            isinstance(self.grid, CartesianGrid) and isinstance(grid, CartesianGrid)
        ):
            # convert within the same grid class
            points = grid.cell_coords
            # vectors are already given in the correct basis
            data = self.interpolate(points, bc=bc, fill=fill)

        elif isinstance(grid, CartesianGrid):
            # convert Cartesian coordinates to coordinates in current grid
            points = self.grid.c.pos_from_cart(grid.cell_coords)
            points_grid_sym = self.grid._coords_symmetric(points)
            # interpolate the data to the grid; this gives the vector in the grid basis
            data_grid = self.interpolate(points_grid_sym, bc=bc, fill=fill)
            # convert the vector to the cartesian basis
            data = self.grid._vector_to_cartesian(points, data_grid)

        else:
            # this type of interpolation is not supported
            grid_in = self.grid.__class__.__name__
            grid_out = grid.__class__.__name__
            msg = f"Can't interpolate from {grid_in} to {grid_out}"
            raise NotImplementedError(msg)

        return self.__class__(grid, data, label=label)

    def _update_plot_components(self, reference: list[PlotReference]) -> None:
        """Update a component plot with the current field values.

        Args:
            reference (list of :class:`PlotReference`):
                All references of the plot to update
        """
        for i in range(self.grid.dim):
            self[i]._update_plot(reference[i])

    @plot_on_figure(update_method="_update_plot_components")
    def plot_components(
        self,
        kind: str = "auto",
        fig=None,
        **kwargs,
    ) -> list[PlotReference]:
        r"""Visualize all the components of this vector field.

        Args:
            kind (str or list of str):
                Determines the kind of the visualizations. Supported values are `image`
                or `line`. Alternatively, `auto` determines the best visualization based
                on the grid.
            {PLOT_ARGS}
            \**kwargs:
                All additional keyword arguments are forwarded to the actual plotting
                function of all subplots.

        Returns:
            list of :class:`PlotReference`: Instances that contain information
            to update all the plots with new data later.
        """
        # create all the subpanels
        dim = self.grid.dim
        axs = fig.subplots(nrows=1, ncols=dim, squeeze=False)

        # plot all the elements onto the respective axes
        kwargs.setdefault("action", "none")
        kwargs["kind"] = kind
        comps = self.grid.axes + self.grid.axes_symmetric
        references = [
            self[i].plot(
                ax=axs[0][i],
                title=f"{comps[i]} Component",
                **kwargs,
            )
            for i in range(dim)
        ]
        # return the references for all subplots
        return references

    def _get_napari_layer_data(  # type: ignore
        self, max_points: int | None = None, args: dict[str, Any] | None = None
    ) -> dict[str, Any]:
        """Returns data for plotting on a single napari layer.

        Args:
            max_points (int):
                The maximal number of points that is used along each axis. This option
                can be used to subsample the data.
            args (dict):
                Additional arguments returned in the result, which affect how the layer
                is shown.

        Returns:
            dict: all the information necessary to plot this field
        """
        result = {} if args is None else args.copy()

        # extract the vector components in the format required by napari
        data = self.get_vector_data(max_points=max_points)
        vectors = np.empty((data["size"], 2, 2))
        xs, ys = np.meshgrid(data["x"], data["y"], indexing="ij")
        vectors[:, 0, 0] = xs.flat
        vectors[:, 0, 1] = ys.flat
        vectors[:, 1, 0] = data["data_x"].flat
        vectors[:, 1, 1] = data["data_y"].flat

        result["type"] = "vectors"
        result["data"] = vectors
        return result
