r"""Solves a 2D Reaction-Diffusion PDE with Dirichlet boundary conditions.

.. math::

    -\delta u + u & = f in \Omega \times M

where :math:`x = (x_1, x_2, \alpha) \in \Omega \times M = (0, 1) \times (0, 1) \times [1, 3]` and :math:`f` such that :math:`u(x_1, x_2, \alpha) = \sin(\pi x_1) \sin(\pi x_2) \sin(\pi \alpha x_1 x_2)`.

The homogeneous Dirichlet boundary conditions are enforced strongly, the neural network is a simple MLP (Multilayer Perceptron), and the optimization is done using Natural Gradient Descent.

This example mainly illustrates how to update the parameter bounds of a PINN, and how to use this feature to solve the PDE for increasing values of the parameter α. Indeed, trying to directly solve the PDE for a large range of α will fail: a good strategy is to start with a small range of α, and increase it progressively, using the solution obtained for the smaller range as a starting point for the next training.
"""

# %%

import matplotlib.pyplot as plt
import torch

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.elliptic_pde.pinns import (
    NaturalGradientPinnsElliptic,
)
from scimba_torch.physical_models.elliptic_pde.laplacians import (
    Laplacian2DDirichletStrongForm,
)
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces
from scimba_torch.utils.scimba_tensors import LabelTensor

torch.manual_seed(0)


def f_rhs(xy: LabelTensor, mu: LabelTensor, k: float = 1):
    x, y = xy.get_components()
    α = mu.get_components()

    pi = torch.pi

    sin_x = torch.sin(pi * x)
    cos_x = torch.cos(pi * x)

    sin_y = torch.sin(pi * y)
    cos_y = torch.cos(pi * y)

    sin_k = torch.sin(pi * k * x * y * α)
    cos_k = torch.cos(pi * k * x * y * α)

    term1 = pi**2 * k**2 * α**2 * (x**2 + y**2) + (1 + 2 * pi**2)

    term2 = 2 * pi**2 * k * α * (cos_x * y * sin_y + x * sin_x * cos_y)

    return term1 * sin_x * sin_y * sin_k - term2 * cos_k


def exact_sol(xy: LabelTensor, mu: LabelTensor, k: float = 1):
    x, y = xy.get_components()
    α = mu.get_components()

    pi = torch.pi

    return (
        torch.sin(torch.pi * x)
        * torch.sin(torch.pi * y)
        * torch.sin(pi * k * α * x * y)
    )


class Laplacian2DDirichletStrongFormNoParam(Laplacian2DDirichletStrongForm):
    def operator(self, w, x, mu):
        u = w.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 -(u_xx + u_yy) + u

    def functional_operator(self, func, x, mu, theta):
        grad_u = torch.func.jacrev(func, 0)
        hessian_u = torch.func.jacrev(grad_u, 0, chunk_size=None)(x, mu, theta)
        return -(hessian_u[..., 0, 0] + hessian_u[..., 1, 1]) + func(x, mu, theta)


def post_processing(inputs: torch.Tensor, x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    phi = torch.sin(torch.pi * x1) * torch.sin(torch.pi * x2)
    return inputs * phi


def functional_post_processing(u, xy, mu, theta):
    return u(xy, mu, theta) * torch.sin(torch.pi * xy[0]) * torch.sin(torch.pi * xy[1])


def create_pde(domain_x, sampler):
    space = NNxSpace(
        1,
        1,
        GenericMLP,
        domain_x,
        sampler,
        layer_sizes=[16, 16],
        post_processing=post_processing,
    )
    return Laplacian2DDirichletStrongFormNoParam(space, f=f_rhs)


domain_x = Square2D([(0, 1), (0, 1)], is_main_domain=True)
sampler_x = DomainSampler(domain_x)

default_param_domain = [[1.0, 1.5]]
default_param_sampler = UniformParametricSampler(default_param_domain)

sampler = TensorizedSampler([sampler_x, default_param_sampler])

pde = create_pde(domain_x, sampler)
pinn = NaturalGradientPinnsElliptic(
    pde, bc_type="strong", functional_post_processing=functional_post_processing
)

# %%

for α_max in [1.0, 1.5, 2.0, 2.5, 3.0]:
    # update the parameter bounds to [[1.0, α_max]]
    pinn.update_parameter_bounds([[1.0, α_max]])
    # train until loss <= α_max * 1e-5, for a maximum of 500 epochs
    pinn.solve(max_epochs=500, loss_target=α_max * 1e-5, n_collocation=3000)

# %%

plot_abstract_approx_spaces(
    (pinn.space,),
    domain_x,
    [[α_max, α_max]],
    loss=(pinn.losses,),
    solution=exact_sol,
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
)

plt.show()

# %%