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

.. math::

    -\delta u & = f in \Omega

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

The boundary conditions are either homogeneous Dirichlet conditions enforced strongly or weakly, or periodic boundary conditions enforced strongly by using a periodic embedding within the neural network.

The neural network is a simple MLP (Multilayer Perceptron).

The optimization is done using either Adam, Natural Gradient Descent or Nystrom Natural Gradient.
"""

# %%

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


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


print("@@@ Strong boundary conditions @@@")

print("@@@@@@@@@@@@@@@ ENG @@@@@@@@@@@@@@@@@@@")

torch.manual_seed(0)

space = NNxSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32, 32],
    post_processing=post_processing,
)
pde = Laplacian2DDirichletStrongFormNoParam(space, f=f_rhs, g=f_bc)
pinns = NaturalGradientPinnsElliptic(
    pde, bc_type="strong", functional_post_processing=functional_post_processing
)

new_solve = False
if new_solve or not pinns.load(__file__, "ENG"):
    pinns.solve(epochs=100, n_collocation=1000, verbose=False)
    pinns.save(__file__, "ENG")

print(f"Training time: {pinns.training_time:.2f} seconds")
print(f"Best Loss    : {pinns.best_loss:.2e}")

print("@@@@@@@@@@@@@@@ ANaGRAM @@@@@@@@@@@@@@@@@@@")

torch.manual_seed(0)

space2 = NNxSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32, 32],
    post_processing=post_processing,
)
pde2 = Laplacian2DDirichletStrongFormNoParam(space2, f=f_rhs, g=f_bc)
pinns2 = NaturalGradientPinnsElliptic(
    pde2,
    bc_type="strong",
    ng_algo="ANaGRAM",
    functional_post_processing=functional_post_processing,
)

new_solve = True
if new_solve or not pinns2.load(__file__, "Anagram"):
    pinns2.solve(epochs=100, n_collocation=1000, verbose=False)
    pinns2.save(__file__, "Anagram")

print(f"Training time: {pinns2.training_time:.2f} seconds")
print(f"Best Loss    : {pinns2.best_loss:.2e}")

print("@@@@@@@@@@@@@@@ Sketchy NG @@@@@@@@@@@@@@@@@@@")

torch.manual_seed(0)

space3 = NNxSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32, 32],
    post_processing=post_processing,
)
pde3 = Laplacian2DDirichletStrongFormNoParam(space3, f=f_rhs, g=f_bc)
pinns3 = NaturalGradientPinnsElliptic(
    pde3,
    bc_type="strong",
    ng_algo="sketchy",
    functional_post_processing=functional_post_processing,
)

new_solve = True
if new_solve or not pinns3.load(__file__, "SNG"):
    pinns3.solve(epochs=100, n_collocation=1000, verbose=False)
    pinns3.save(__file__, "SNG")

print(f"Training time: {pinns3.training_time:.2f} seconds")
print(f"Best Loss    : {pinns3.best_loss:.2e}")

print("@@@@@@@@@@@@@@@ Nyström NG @@@@@@@@@@@@@@@@@@@")

torch.manual_seed(0)

space4 = NNxSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32, 32],
    post_processing=post_processing,
)
pde4 = Laplacian2DDirichletStrongFormNoParam(space4, f=f_rhs, g=f_bc)
pinns4 = NaturalGradientPinnsElliptic(
    pde4,
    bc_type="strong",
    ng_algo="Nyström",
    functional_post_processing=functional_post_processing,
    # matrix_free=True,
    debug=True,
)

new_solve = True
if new_solve or not pinns4.load(__file__, "NNG"):
    pinns4.solve(epochs=100, n_collocation=1000, verbose=False)
    pinns4.save(__file__, "NNG")

print(f"Training time: {pinns4.training_time:.2f} seconds")
print(f"Best Loss    : {pinns4.best_loss:.2e}")

plot_abstract_approx_spaces(
    (
        pinns.space,
        pinns2.space,
        pinns3.space,
        pinns4.space,
    ),
    domain_x,
    loss=(
        pinns.losses,
        pinns2.losses,
        pinns3.losses,
        pinns4.losses,
    ),
    residual=(
        pinns.pde,
        pinns2.pde,
        pinns3.pde,
        pinns4.pde,
    ),
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
)

plt.show()

print("@@@ Weak boundary conditions @@@")

print("@@@@@@@@@@@@@@@ ENG @@@@@@@@@@@@@@@@@@@")

torch.manual_seed(0)

space = NNxSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32, 32],
)
pde = Laplacian2DDirichletStrongFormNoParam(space, f=f_rhs, g=f_bc)
pinns = NaturalGradientPinnsElliptic(pde, bc_type="weak", bc_weight=30.0)

new_solve = True
if new_solve or not pinns.load(__file__, "weak_ENG"):
    pinns.solve(
        epochs=100,
        n_collocation=1000,
        n_bc_collocation=1000,
        verbose=False,
    )
    pinns.save(__file__, "weak_ENG")

print(f"Training time: {pinns.training_time:.2f} seconds")
print(f"Best Loss    : {pinns.best_loss:.2e}")

print("@@@@@@@@@@@@@@@ ANaGRAM @@@@@@@@@@@@@@@@@@@")

torch.manual_seed(0)

space2 = NNxSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32, 32],
)
pde2 = Laplacian2DDirichletStrongFormNoParam(space2, f=f_rhs, g=f_bc)
pinns2 = NaturalGradientPinnsElliptic(
    pde2, bc_type="weak", bc_weight=30.0, ng_algo="ANaGRAM", svd_threshold=5e-1
)

new_solve = True
if new_solve or not pinns2.load(__file__, "weak_Anagram"):
    pinns2.solve(
        epochs=100,
        n_collocation=1000,
        n_bc_collocation=1000,
        verbose=False,
    )
    pinns2.save(__file__, "weak_Anagram")

print(f"Training time: {pinns2.training_time:.2f} seconds")
print(f"Best Loss    : {pinns2.best_loss:.2e}")

print("@@@@@@@@@@@@@@@ Sketchy NG @@@@@@@@@@@@@@@@@@@")

torch.manual_seed(0)

space3 = NNxSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32, 32],
)
pde3 = Laplacian2DDirichletStrongFormNoParam(space3, f=f_rhs, g=f_bc)
pinns3 = NaturalGradientPinnsElliptic(
    pde3,
    bc_type="weak",
    bc_weight=30.0,
    ng_algo="sketchy",
    # tol=1e-10
)

new_solve = True
if new_solve or not pinns3.load(__file__, "weak_SNG"):
    pinns3.solve(
        epochs=100,
        n_collocation=1000,
        n_bc_collocation=1000,
        verbose=False,
    )
    pinns3.save(__file__, "weak_SNG")

print(f"Training time: {pinns3.training_time:.2f} seconds")
print(f"Best Loss    : {pinns3.best_loss:.2e}")

print("@@@@@@@@@@@@@@@ Nyström NG @@@@@@@@@@@@@@@@@@@")

torch.manual_seed(0)

space4 = NNxSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32, 32],
)
pde4 = Laplacian2DDirichletStrongFormNoParam(space4, f=f_rhs, g=f_bc)
pinns4 = NaturalGradientPinnsElliptic(
    pde4,
    bc_type="weak",
    bc_weight=30.0,
    ng_algo="Nyström",
    # matrix_free=True,
    debug=True,
)

new_solve = True
if new_solve or not pinns4.load(__file__, "weak_NNG"):
    pinns4.solve(
        epochs=100,
        n_collocation=1000,
        n_bc_collocation=1000,
        verbose=False,
    )
    pinns4.save(__file__, "weak_NNG")

print(f"Training time: {pinns4.training_time:.2f} seconds")
print(f"Best Loss    : {pinns4.best_loss:.2e}")

plot_abstract_approx_spaces(
    (
        pinns.space,
        pinns2.space,
        pinns3.space,
        pinns4.space,
    ),
    domain_x,
    loss=(
        pinns.losses,
        pinns2.losses,
        pinns3.losses,
        pinns4.losses,
    ),
    residual=(
        pinns.pde,
        pinns2.pde,
        pinns3.pde,
        pinns4.pde,
    ),
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
)

plt.show()

# %%