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 SteadyReactionDiffusion2DNeumannStrongForm(Laplacian2DDirichletStrongForm):
    """Implementation of a 2D steady reaction-diffusion 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,
        f: Callable = zeros,
        g: Callable = zeros,
        **kwargs,
    ):
        super().__init__(space, f, g, **kwargs)

    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 = 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 -D * (u_xx + u_yy) - 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 = mu[0]
        grad_u = torch.func.jacrev(func, 0)
        hessian_u = torch.func.jacrev(grad_u, 0, chunk_size=None)(x, mu, theta)
        return -D * (hessian_u[..., 0, 0] + hessian_u[..., 1, 1]) - u

    def bc_operator(
        self, w: MultiLabelTensor, x: LabelTensor, n: LabelTensor, mu: LabelTensor
    ) -> torch.Tensor:
        """Compute the operator for the boundary conditions.

        Args:
            w: State tensor.
            x: Boundary coordinates tensor.
            n: Normal vector tensor.
            mu: Parameter tensor.

        Returns:
            The boundary operator applied to the state.
        """
        u = w.get_components()
        u_x, u_y = self.grad(u, x)
        nx, ny = n.get_components()
        return u_x * nx + u_y * ny

    def functional_operator_bc(
        self,
        func: VarArgCallable,
        x: torch.Tensor,
        n: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        """Apply the functional operator for boundary conditions.

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

        Returns:
            The result of applying the functional operator.
        """
        grad_u = torch.func.jacrev(func, 0)(x, mu, theta)
        return grad_u @ n


class ReactionDiffusion2D(HeatEquation2DStrongForm):
    r"""Implementation of a 2D reaction-diffusion 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)
        **kwargs: Additional keyword arguments
    """

    def __init__(self, space: AbstractApproxSpace, init: Callable, **kwargs):
        super().__init__(space, init, **kwargs)
        self.space_component = SteadyReactionDiffusion2DNeumannStrongForm(space)

    def functional_operator_bc(
        self,
        func: VarArgCallable,
        # t: torch.Tensor,
        x: torch.Tensor,
        n: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        """Compute the functional operator for boundary conditions.

        Args:
            func: Callable representing the function
            t: Temporal coordinate tensor
            x: Spatial coordinate tensor
            n: Normal vector tensor
            mu: Parameter tensor
            theta: Additional parameters tensor

        Returns:
            Functional operator tensor for boundary conditions
        """
        return self.space_component.functional_operator_bc(func, x, n, mu, theta)


def f_exact(t: LabelTensor, x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    D = mu.get_components()
    t_ = t.get_components()
    cos = torch.cos(torch.pi * x1) * torch.cos(torch.pi * x2)
    exp = torch.exp((1 - 2 * D[0] * torch.pi**2) * t_)
    return cos * exp


def f_ini(x: LabelTensor, mu):
    t = LabelTensor(torch.zeros((x.shape[0], 1)), x.labels)
    return f_exact(t, x, mu)


domain_x = Square2D([(0, 1), (0, 1)], is_main_domain=True)
domain_mu = [(0.1, 0.1 + 1e-5)]
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler(domain_mu)]
)

space = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[20] * 2,
    parameters_bounds=domain_mu,
)

pde = ReactionDiffusion2D(space, init=f_ini)
projector_init = NaturalGradientProjector(
    space, rhs=pde.initial_condition, bool_weak_bc=True
)
projector = TimeDiscreteNaturalGradientProjector(
    pde, rhs=f_ini, bc_weight=100, bool_weak_bc=True
)
scheme = DiscretePINN(
    pde, projector, projector_init=projector_init, scheme="rk2", bool_weak_bc=True
)

dt = 1e-2
T = 20 * dt

scheme.initialization(epochs=250, n_collocation=3000)
scheme.solve(dt=dt, final_time=T, n_bc_collocation=3000, n_collocation=2000, epochs=15)

plot_time_discrete_scheme(
    scheme,
    solution=f_exact,
    error=f_exact,
    draw_contours=True,
)

plt.show()