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(\pi x_1) \sin(\pi x_2)` and :math:`g = 0`.

Boundary conditions are enforced weakly.

The neural network used is a simple MLP (Multilayer Perceptron), and the optimization is done using either Adam, Natural Gradient Descent or Anagram. This example illustrates the use of the `LinearOrder2PDE` framework to design a Poisson problem.
"""

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 (
    AnagramPinnsElliptic,
    NaturalGradientPinnsElliptic,
    PinnsElliptic,
)
from scimba_torch.optimizers.optimizers_data import OptimizerData
from scimba_torch.physical_models.elliptic_pde.laplacians import (
    Laplacian2DDirichletStrongForm,
)
from scimba_torch.physical_models.elliptic_pde.linear_order_2 import LinearOrder2PDE
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces
from scimba_torch.utils.scimba_tensors import LabelTensor

torch.manual_seed(0)

bc_weight = 40.0


def f_rhs(x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    mu1 = mu.get_components()
    return 2 * mu1 * 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


# sol exacte : u(x) = mu*sin(2*pi*x1)*sin(2*pi*x2)
def exact_sol(x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    # mu1 = mu.get_components()
    return torch.sin(torch.pi * x1) * torch.sin(torch.pi * x2)


domain_x = Square2D([(0.0, 1), (0.0, 1)], is_main_domain=True)
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler([(1.0, 1.0 + 1e-5)])]
)

space = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32],
)
pde = Laplacian2DDirichletStrongForm(space, f=f_rhs, g=f_bc)
opt_1 = {
    "name": "adam",
    "optimizer_args": {"lr": 1.8e-2, "betas": (0.9, 0.999)},
}
pinns = PinnsElliptic(
    pde, bc_type="weak", optimizers=OptimizerData(opt_1), bc_weight=bc_weight
)

resume_solve = False
if resume_solve or not pinns.load(__file__, "pinns"):
    pinns.solve(epochs=1000, n_collocation=3000, verbose=True)
    pinns.save(__file__, "pinns")

pinns.space.load_from_best_approx()

space2 = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32],
)
pde2 = LinearOrder2PDE(space2, 2, f=f_rhs, g=f_bc)

pinns2 = NaturalGradientPinnsElliptic(pde2, bc_type="weak", bc_weight=bc_weight)

resume_solve = True
if resume_solve or not pinns2.load(__file__, "ENG"):
    # load from pinns.space
    # pinns2.space.load_from_dict(pinns.space.dict_for_save())
    # pinns2.space.load_from_best_approx()
    pinns2.solve(epochs=200, n_collocation=900, verbose=True, n_bc_collocation=120)
    pinns2.save(__file__, "ENG")

pinns2.space.load_from_best_approx()

space3 = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32],
)
pde3 = LinearOrder2PDE(space3, 2, f=f_rhs, g=f_bc)

pinns3 = AnagramPinnsElliptic(
    pde3, bc_type="weak", bc_weight=bc_weight, svd_threshold=5e-4
)

resume_solve = True
if resume_solve or not pinns3.load(__file__, "Anagram"):
    # load from pinns.space
    # pinns3.space.load_from_dict(pinns.space.dict_for_save())
    # pinns3.space.load_from_best_approx()
    pinns3.solve(epochs=200, n_collocation=900, verbose=True, n_bc_collocation=120)
    pinns3.save(__file__, "Anagram")

pinns3.space.load_from_best_approx()


plot_abstract_approx_spaces(
    (
        pinns.space,
        pinns2.space,
        pinns3.space,
    ),  # an Iterable of AbstractSpace
    domain_x,  # either a VolumetricDomain, or an Iterable of VolumetricDomain of length 1 or len(first argument)
    ((1.0, 1.0 + 1e-5),),
    loss=(
        pinns.losses,
        pinns2.losses,
        pinns3.losses,
    ),  # same as previously; if only one is given, it will be used for all spaces
    residual=(
        pinns.pde,
        pinns2.pde,
        pinns3.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,
    parameters_values="random",
    title="2D Laplacian with Dirichlet boundary conditions",
    titles=("no preconditioner", "ENG preconditioner", "Anagram preconditioner"),
)

plt.show()