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.
"""

import torch

from scimba_torch.approximation_space.nn_space import NNxtSpace
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.integration.monte_carlo_time import UniformTimeSampler
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.numerical_solvers.temporal_pde.pinns import (
    NaturalGradientTemporalPinns,
)
from scimba_torch.physical_models.elliptic_pde.abstract_elliptic_pde import (
    StrongFormEllipticPDE,
)
from scimba_torch.physical_models.temporal_pde.abstract_temporal_pde import (
    FirstOrderTemporalPDE,
)
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces
from scimba_torch.utils.scimba_tensors import LabelTensor


def zeros_rhs(w, t, x, mu, nb_func: int = 1):
    return torch.zeros(x.shape[0], nb_func)


def zeros_bc_rhs(w, t, x, n, mu, nb_func: int = 1):
    return torch.zeros(x.shape[0], nb_func)


class ReactionDiffusionNDSpaceComponent(StrongFormEllipticPDE):
    def __init__(self, space, f, g, **kwargs):
        super().__init__(
            space, linear=True, residual_size=1, bc_residual_size=1, **kwargs
        )
        self.f = f
        self.g = g

    def rhs(self, w, x, mu):
        return self.f(x, mu)

    def operator(self, w, x, mu):
        u = w.get_components()
        D = mu.get_components()

        grad_u = torch.cat(tuple(self.grad(u, x)), dim=-1)
        space_dim = grad_u.shape[1]
        laplacian_u = [tuple(self.grad(grad_u[:, i], x))[i] for i in range(space_dim)]
        laplacian_u = torch.sum(torch.stack(laplacian_u, dim=-1), dim=-1)

        return -D * laplacian_u - u

    def bc_rhs(self, w, x, n, mu):
        return self.g(x, mu)

    def bc_operator(self, w, x, n, mu):
        u = w.get_components()
        grad_u = torch.cat(tuple(self.grad(u, x)), dim=-1)
        n_ = torch.cat(tuple(n.get_components()), dim=-1)
        return torch.sum(grad_u * n_, dim=-1, keepdim=True)


class ReactionDiffusionND(FirstOrderTemporalPDE):
    def __init__(self, space, init, f=zeros_rhs, g=zeros_bc_rhs, **kwargs):
        super().__init__(space, linear=True, **kwargs)
        self.space_component = ReactionDiffusionNDSpaceComponent(space, f, g)

        self.f = f
        self.g = g
        self.init = init

        self.ic_residual_size = 1

    def space_operator(self, w, t, x, mu):
        return self.space_component.operator(w, x, mu)

    def bc_operator(self, w, t, x, n, mu):
        return self.space_component.bc_operator(w, x, n, mu)

    def rhs(self, w, t, x, mu: LabelTensor):
        return self.f(w, t, x, mu)

    def bc_rhs(self, w, t, x, n, mu):
        return self.g(w, t, x, n, mu)

    def initial_condition(self, x, mu):
        return self.init(x, mu)

    def functional_operator(self, func, t, x, mu, theta):
        u = func(t, x, mu, theta)
        D = mu[0]

        time_op = torch.func.jacrev(func, 0)(t, x, mu, theta)

        grad_u = torch.func.jacrev(func, 1)
        grad_grad_u = torch.func.jacrev(grad_u, 1)(t, x, mu, theta).squeeze()
        laplacian_u = torch.einsum("ii", grad_grad_u)

        return time_op[0] - D * laplacian_u[None] - u

    def functional_operator_bc(self, func, t, x, n, mu, theta):
        grad_u = torch.func.jacrev(func, 1)(t, x, mu, theta)
        return grad_u @ n

    def functional_operator_ic(self, func, x, mu, theta):
        return func(torch.zeros(1), x, 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)


def solve() -> NaturalGradientTemporalPinns:
    """Solves the problem for the reaction-diffusion equation."""

    domain_x = Square2D([(0, 1), (0, 1)], is_main_domain=True)
    domain_mu = [(0.1, 0.1 + 1e-5)]
    t_min, t_max = 0.0, 1.0

    sampler = TensorizedSampler(
        [
            UniformTimeSampler((t_min, t_max)),
            DomainSampler(domain_x),
            UniformParametricSampler(domain_mu),
        ]
    )

    space = NNxtSpace(1, 1, GenericMLP, domain_x, sampler, layer_sizes=[16, 16])

    pde = ReactionDiffusionND(space, init=f_ini)

    pinn = NaturalGradientTemporalPinns(
        pde, ic_type="weak", bc_type="weak", bc_weight=10, ic_weight=10
    )

    pinn.solve(
        epochs=100, n_collocation=2_000, n_bc_collocation=2_000, n_ic_collocation=2_000
    )

    return {"x": domain_x, "mu": domain_mu, "t": (t_min, t_max)}, pinn


# %%

if __name__ == "__main__":
    domains, pinn = solve()

    plot_abstract_approx_spaces(
        pinn.space,
        domains["x"],
        domains["mu"],
        domains["t"],
        loss=pinn.losses,
        exact=f_exact,
        error=f_exact,
    )