r"""Solves a 2D reaction-diffusion equation.

.. math::::
    \partial_t u - D \Delta u - u = 0

The equation is solved using an explicit discrete PINN, and a Natural Gradient optimizer.
"""

from typing import Callable

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.abstract_space import AbstractApproxSpace
from scimba_torch.approximation_space.nn_space import NNxSpace
from scimba_torch.domain.meshless_domain.domain_2d import Square2D
from scimba_torch.integration.monte_carlo import DomainSampler, TensorizedSampler
from scimba_torch.integration.monte_carlo_parameters import UniformParametricSampler
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.numerical_solvers.collocation_projector import (
    NaturalGradientProjector,
)
from scimba_torch.numerical_solvers.temporal_pde.discrete_pinn import (
    DiscretePINN,
)
from scimba_torch.numerical_solvers.temporal_pde.time_discrete import (
    TimeDiscreteNaturalGradientProjector,
)
from scimba_torch.physical_models.elliptic_pde.laplacians import (
    Laplacian2DDirichletStrongForm,
)
from scimba_torch.physical_models.temporal_pde.heat_equation import (
    HeatEquation2DStrongForm,
)
from scimba_torch.plots.plot_time_discrete_scheme import plot_time_discrete_scheme
from scimba_torch.utils.scimba_tensors import LabelTensor, MultiLabelTensor
from scimba_torch.utils.typing_protocols import VarArgCallable


def zeros(x: LabelTensor, mu: LabelTensor, nb_func: int = 1) -> torch.Tensor:
    """Function returning zeros.

    Args:
        x: Spatial coordinates tensor.
        mu: Parameter tensor.
        nb_func: Number of functions to return (default is 1).

    Returns:
        A tensor of zeros with shape (number of points, nb_func).
    """
    return torch.zeros(x.shape[0], nb_func)


class SteadyFisherKPP2DNeumannStrongForm(Laplacian2DDirichletStrongForm):
    """Implementation of a 2D steady Fisher-KPP problem with Neumann BCs in strong form.

    Args:
        space: The approximation space for the problem.
        f: Callable, represents the source term f(x, μ) (default: f=1).
        g: Callable, represents the Neumann boundary condition g(x, μ) (default: g=0).
        **kwargs: Additional keyword arguments.
    """

    def __init__(
        self,
        space: AbstractApproxSpace,
        final_time: float,
        f: Callable = zeros,
        g: Callable = zeros,
        **kwargs,
    ):
        super().__init__(space, f, g, **kwargs)
        self.final_time = final_time

    def operator(
        self, w: MultiLabelTensor, x: LabelTensor, mu: LabelTensor
    ) -> torch.Tensor:
        """Compute the differential operator of the PDE.

        Args:
            w: State tensor.
            x: Spatial coordinates tensor.
            mu: Parameter tensor.

        Returns:
            The result of applying the operator to the state.
        """
        u = w.get_components()
        D, r = mu.get_components()
        u_x, u_y = self.grad(u, x)
        u_xx, _ = self.grad(u_x, x)
        _, u_yy = self.grad(u_y, x)
        return -self.final_time * (D * (u_xx + u_yy) + r * u * (1 - u))

    def functional_operator(
        self,
        func: VarArgCallable,
        x: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        """Apply the functional operator for the differential equation.

        Args:
            func: The callable function to apply.
            x: Spatial coordinates tensor.
            mu: Parameter tensor.
            theta: Theta parameter tensor.

        Returns:
            The result of applying the functional operator.
        """
        u = func(x, mu, theta)
        D, r = mu[0], mu[1]
        grad_u = torch.func.jacrev(func, 0)
        hessian_u = torch.func.jacrev(grad_u, 0, chunk_size=None)(x, mu, theta)
        laplacian_u = hessian_u[..., 0, 0] + hessian_u[..., 1, 1]
        return -self.final_time * (D * laplacian_u + r * u * (1 - u))


class FisherKPP2D(HeatEquation2DStrongForm):
    r"""Implementation of a 2D Fisher-KPP problem with Dirichlet BCs in strong form.

    Args:
        space: The approximation space for the problem
        init: Callable for the initial condition
        f: Source term function (default is zero)
        g: Dirichlet boundary condition function (default is zero)
        final_time: The final time for the problem (default is 50.0)
        **kwargs: Additional keyword arguments
    """

    def __init__(
        self,
        space: AbstractApproxSpace,
        init: Callable,
        final_time: float = 50.0,
        **kwargs,
    ):
        super().__init__(space, init, **kwargs)
        self.space_component = SteadyFisherKPP2DNeumannStrongForm(
            space, final_time=final_time
        )


def f_exact(t: LabelTensor, x: LabelTensor, mu: LabelTensor):
    return f_ini(x, mu)


def f_ini(x: LabelTensor, mu):
    center = [10, 10]
    x1, x2 = x.get_components()
    r2 = (x1 - center[0]) ** 2 + (x2 - center[1]) ** 2
    return torch.exp(-r2 / 10.0)


def phi_neumann(x: torch.tensor) -> torch.Tensor:
    return 3 * (x / 50) ** 2 - 2 * (x / 50) ** 3


def pre_processing(x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    D, r = mu.get_components()
    r2 = torch.sum((x.x - 10) ** 2, dim=-1, keepdim=True) / 10**2
    return torch.cat([phi_neumann(x1), phi_neumann(x2), r2, D, r], dim=-1)


def functional_pre_processing(*args):
    x1, x2 = args[0].split(1, dim=-1)
    D, r = args[1].split(1, dim=-1)
    r2 = torch.sum((args[0] - 10) ** 2, dim=-1, keepdim=True) / 10**2
    return torch.cat([phi_neumann(x1), phi_neumann(x2), r2, D, r], dim=-1)


domain_x = Square2D([(0, 50), (0, 50)], is_main_domain=True)
domain_mu = [(0.05, 0.15), (0.05, 0.15)]
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler(domain_mu)]
)

space = NNxSpace(
    1,
    2,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[16, 16],
    parameters_bounds=domain_mu,
    pre_processing=pre_processing,
    activation_output="sigmoid",
    pre_processing_out_size=5,  # x1, x2, r^2, D, r
)

pde = FisherKPP2D(space, init=f_ini)
projector_init = NaturalGradientProjector(
    space,
    rhs=pde.initial_condition,
    functional_pre_processing=functional_pre_processing,
)
projector = TimeDiscreteNaturalGradientProjector(
    pde,
    rhs=f_ini,
    functional_pre_processing=functional_pre_processing,
)
scheme = DiscretePINN(pde, projector, projector_init=projector_init, scheme="rk4")

dt = 5e-2
T = 20 * dt

scheme.initialization(epochs=500, n_collocation=5_000)
scheme.solve(dt=dt, final_time=T, n_collocation=2_000, epochs=15)

plot_time_discrete_scheme(
    scheme,
    solution=f_exact,
    error=f_exact,
    derivatives=["ux", "uy"],
    draw_contours=True,
)

plt.show()