r"""Solves a 2D Poisson PDE with Dirichlet boundary conditions using PINNs.

.. math::

    -\delta u & = f in \Omega \\
    u & = g on \partial \Omega

where :math:`x = (x_1, x_2) \in \Omega = (0, 1) \times (0, 1)`, :math:`f` such that :math:`u(x_1, x_2) = \sin(2\pi x_1) \sin(2\pi x_2)`, and :math:`g = 0`.

Boundary conditions are enforced strongly.

The neural network is either a simple MLP (Multilayer Perceptron) or a ResNet (Residual Neural Network), and the optimization is done using either Adam or Natural Gradient Descent.
"""

# %%

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.neural_nets.coordinates_based_nets.res_net import GenericResNet
from scimba_torch.numerical_solvers.elliptic_pde.pinns import (
    NaturalGradientPinnsElliptic,
    PinnsElliptic,
)
from scimba_torch.optimizers.optimizers_data import OptimizerData
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(x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    return 2 * torch.pi**2 * torch.sin(torch.pi * x1) * torch.sin(torch.pi * x2)


def f_bc(x: LabelTensor, mu: LabelTensor):
    x1, _ = x.get_components()
    return x1 * 0.0


def exact_sol(x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    return torch.sin(torch.pi * x1) * torch.sin(torch.pi * x2)


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)

    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])


domain_x = Square2D([(-1, 1), (-1, 1)], is_main_domain=True)
sampler = TensorizedSampler([DomainSampler(domain_x), UniformParametricSampler([])])


def post_processing(inputs: torch.Tensor, x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    return inputs * (x1 + 1) * (1 - x1) * (x2 + 1) * (1 - x2)


def functional_post_processing(u, xy, mu, theta):
    return u(xy, mu, theta) * (xy[0] + 1) * (1 - xy[0]) * (xy[1] + 1) * (1 - xy[1])


# %%

#### first space: MLP #####

space_MLP = NNxSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[10] * 6,
    post_processing=post_processing,
)
pde_MLP = Laplacian2DDirichletStrongFormNoParam(space_MLP, f=f_rhs, g=f_bc)
dict_opt = {
    "name": "adam",
    "optimizer_args": {"lr": 1.8e-2, "betas": (0.9, 0.999)},
}
opt_MLP = OptimizerData(dict_opt)
PINN_MLP = PinnsElliptic(pde_MLP, bc_type="strong", optimizers=opt_MLP)

# train until loss <= 1e-2, for a maximum of 3000 epochs
PINN_MLP.solve(max_epochs=3000, loss_target=0.01, n_collocation=3000, verbose=True)

# %%

#### second space: ResNet #####

space_ResNet = NNxSpace(
    1,
    0,
    GenericResNet,
    domain_x,
    sampler,
    layer_structure=[10, 6, [1, 3], [4, 6]],
    post_processing=post_processing,
)
pde_ResNet = Laplacian2DDirichletStrongFormNoParam(space_ResNet, f=f_rhs, g=f_bc)
dict_opt = {
    "name": "adam",
    "optimizer_args": {"lr": 1.8e-2, "betas": (0.9, 0.999)},
}
opt_ResNet = OptimizerData(dict_opt)
PINN_ResNet = PinnsElliptic(pde_ResNet, bc_type="strong", optimizers=opt_ResNet)

# train until loss <= 1e-2, for a maximum of 3000 epochs
PINN_ResNet.solve(max_epochs=3000, loss_target=0.01, n_collocation=3000, verbose=True)

# %%

# for plotting PINN_MLP and PINN_ResNet
plot_abstract_approx_spaces(
    (
        PINN_MLP.space,
        PINN_ResNet.space,
    ),  # an Iterable of AbstractSpace
    domain_x,  # either a VolumetricDomain, or an Iterable of VolumetricDomain of length 1 or len(first argument)
    loss=(
        PINN_MLP.losses,
        PINN_ResNet.losses,
    ),  # same as previously; if only one is given, it will be used for all spaces
    residual=(
        PINN_MLP.pde,
        PINN_ResNet.pde,
    ),  # same as previously; if only one is given, it will be used for all spaces
    error=exact_sol,  # same as previously; if only one is given, it will be used for all spaces
    draw_contours=True,
    n_drawn_contours=20,
)

plt.show()

# %%

#### now, we do the same but with the natural gradient ####

#### first space: MLP #####

space_MLP_ENG = NNxSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[12] * 6,
    post_processing=post_processing,
)
pde_MLP_ENG = Laplacian2DDirichletStrongFormNoParam(space_MLP_ENG, f=f_rhs, g=f_bc)
PINN_MLP_ENG = NaturalGradientPinnsElliptic(
    pde_MLP_ENG, bc_type="strong", functional_post_processing=functional_post_processing
)

# train until loss <= 5e-9, for a maximum of 200 epochs
PINN_MLP_ENG.solve(max_epochs=200, loss_target=5e-9, n_collocation=3000, verbose=True)

# %%

#### second space: ResNet #####

space_ResNet_ENG = NNxSpace(
    1,
    0,
    GenericResNet,
    domain_x,
    sampler,
    layer_sizes=[12, 6, [1, 3], [4, 6]],
    post_processing=post_processing,
)
pde_ResNet_ENG = Laplacian2DDirichletStrongFormNoParam(
    space_ResNet_ENG, f=f_rhs, g=f_bc
)
PINN_ResNet_ENG = NaturalGradientPinnsElliptic(
    pde_ResNet_ENG,
    bc_type="strong",
    functional_post_processing=functional_post_processing,
)

# train until loss <= 5e-9, for a maximum of 200 epochs
PINN_ResNet_ENG.solve(
    max_epochs=200, loss_target=5e-9, n_collocation=3000, verbose=True
)

# %%

# for plotting PINN_MLP_ENG and PINN_ResNet_ENG
plot_abstract_approx_spaces(
    (
        PINN_MLP_ENG.space,
        PINN_ResNet_ENG.space,
    ),  # an Iterable of AbstractSpace
    domain_x,  # either a VolumetricDomain, or an Iterable of VolumetricDomain of length 1 or len(first argument)
    loss=(
        PINN_MLP_ENG.losses,
        PINN_ResNet_ENG.losses,
    ),  # same as previously; if only one is given, it will be used for all spaces
    residual=(
        PINN_MLP_ENG.pde,
        PINN_ResNet_ENG.pde,
    ),  # same as previously; if only one is given, it will be used for all spaces
    error=exact_sol,  # same as previously; if only one is given, it will be used for all spaces
    draw_contours=True,
    n_drawn_contours=20,
)

plt.show()
# %%