"""The Swift-Hohenberg equation.

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

from __future__ import annotations

from typing import TYPE_CHECKING

from ..fields import ScalarField
from ..grids.boundaries import set_default_bc
from ..tools.docstrings import fill_in_docstring
from .base import PDEBase, expr_prod

if TYPE_CHECKING:
    from collections.abc import Callable

    import torch

    from ..grids.boundaries.axes import BoundariesData
    from ..tools.typing import NumericArray


class SwiftHohenbergPDE(PDEBase):
    r"""The Swift-Hohenberg equation.

    The mathematical definition is

    .. math::
        \partial_t c =
            \left[\epsilon - \left(k_c^2 + \nabla^2\right)^2\right] c
            + \delta \, c^2 - c^3

    where :math:`c` is a scalar field and :math:`\epsilon`, :math:`k_c^2`, and
    :math:`\delta` are parameters of the equation.
    """

    explicit_time_dependence = False
    default_bc = "auto_periodic_neumann"
    """Default boundary condition used when no specific conditions are chosen."""

    @fill_in_docstring
    def __init__(
        self,
        rate: float = 0.1,
        kc2: float = 1.0,
        delta: float = 1.0,
        *,
        bc: BoundariesData | None = None,
        bc_lap: BoundariesData | None = None,
    ):
        r"""
        Args:
            rate (float):
                The bifurcation parameter :math:`\epsilon`
            kc2 (float):
                Squared wave vector :math:`k_c^2` of the linear instability
            delta (float):
                Parameter :math:`\delta` of the non-linearity
            bc:
                The boundary conditions applied to the field.
                {ARG_BOUNDARIES}
            bc_lap:
                The boundary conditions applied to the second derivative of the scalar
                field :math:`c`. If `None`, the same boundary condition as `bc` is
                chosen. Otherwise, this supports the same options as `bc`.
        """
        super().__init__()

        self.rate = rate
        self.kc2 = kc2
        self.delta = delta
        self.bc = set_default_bc(bc, self.default_bc)
        self.bc_lap = self.bc if bc_lap is None else bc_lap

    @property
    def expression(self) -> str:
        """str: the expression of the right hand side of this PDE"""
        return (
            f"{expr_prod(self.rate - self.kc2**2, 'c')} - c³"
            f" + {expr_prod(self.delta, 'c²')}"
            f" - ∇²({expr_prod(2 * self.kc2, 'c')} + ∇²c)"
        )

    def evolution_rate(  # type: ignore
        self,
        state: ScalarField,
        t: float = 0,
    ) -> ScalarField:
        """Evaluate the right hand side of the PDE.

        Args:
            state (:class:`~pde.fields.ScalarField`):
                The scalar field describing the concentration distribution
            t (float):
                The current time point

        Returns:
            :class:`~pde.fields.ScalarField`:
            Scalar field describing the evolution rate of the PDE
        """
        if not isinstance(state, ScalarField):
            msg = "`state` must be ScalarField"
            raise TypeError(msg)
        state_laplace = state.laplace(bc=self.bc, args={"t": t})
        state_laplace2 = state_laplace.laplace(bc=self.bc_lap, args={"t": t})

        result = (
            (self.rate - self.kc2**2) * state
            - 2 * self.kc2 * state_laplace
            - state_laplace2
            + self.delta * state**2
            - state**3
        )
        result.label = "evolution rate"
        return result  # type: ignore

    def make_pde_rhs_numba(
        self, state: ScalarField
    ) -> Callable[[NumericArray, float], NumericArray]:
        """Create a compiled function evaluating the right hand side of the PDE.

        Args:
            state (:class:`~pde.fields.ScalarField`):
                An example for the state defining the grid and data types

        Returns:
            A function with signature `(state_data, t)`, which can be called
            with an instance of :class:`~numpy.ndarray` of the state data and
            the time to obtained an instance of :class:`~numpy.ndarray` giving
            the evolution rate.
        """
        rate = self.rate
        kc2 = self.kc2
        delta = self.delta

        laplace = state.grid.make_operator("laplace", bc=self.bc, backend="numba")
        laplace2 = state.grid.make_operator("laplace", bc=self.bc_lap, backend="numba")

        def pde_rhs(state_data: NumericArray, t: float = 0) -> NumericArray:
            """Compiled helper function evaluating right hand side."""
            state_laplace = laplace(state_data, args={"t": t})
            state_laplace2 = laplace2(state_laplace, args={"t": t})

            return (
                (rate - kc2**2) * state_data
                - 2 * kc2 * state_laplace
                - state_laplace2
                + delta * state_data**2
                - state_data**3
            )

        return pde_rhs

    def make_pde_rhs_torch(
        self, state: ScalarField
    ) -> Callable[[torch.Tensor, float], torch.Tensor]:
        """Create a compiled function evaluating the right hand side of the PDE.

        Args:
            state (:class:`~pde.fields.ScalarField`):
                An example for the state defining the grid and data types

        Returns:
            A function with signature `(state_data, t)`, which can be called
            with an instance of :class:`torch.Tensor` of the state data and
            the time to obtained an instance of :class:`torch.Tensor` giving
            the evolution rate.
        """
        rate = self.rate
        kc2 = self.kc2
        delta = self.delta

        laplace = state.grid.make_operator(
            operator="laplace",
            bc=self.bc,
            backend="torch",
            native=True,
            dtype=state.dtype,
        )
        laplace2 = state.grid.make_operator(
            operator="laplace",
            bc=self.bc_lap,
            backend="torch",
            native=True,
            dtype=state.dtype,
        )

        def pde_rhs(state_data: torch.Tensor, t: float = 0) -> torch.Tensor:
            """Compiled helper function evaluating right hand side."""
            state_laplace = laplace(state_data, args={"t": t})
            state_laplace2 = laplace2(state_laplace, args={"t": t})

            return (
                (rate - kc2**2) * state_data
                - 2 * kc2 * state_laplace
                - state_laplace2
                + delta * state_data**2
                - state_data**3
            )

        return pde_rhs
