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

.. math::

    -\mu \delta u & = f in \Omega \times M \\
    u & = 0 on \partial \Omega \times M

where :math:`x = (x_1, x_2) \in \Omega` (with :math:`\Omega` a square rotated by :math:`\pi/3` and with a disk-shaped hole), and :math:`\mu \in M = [1, 2]`.

Boundary conditions are enforced weakly.

The neural network used is a simple MLP (Multilayer Perceptron), and the optimization is done using either Adam + L-BFGS 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 Disk2D, 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,
    PinnsElliptic,
)
from scimba_torch.optimizers.losses import GenericLosses
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_space
from scimba_torch.utils.mapping import Mapping
from scimba_torch.utils.scimba_tensors import LabelTensor


def f_rhs(x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    mu1 = mu.get_components()
    return (
        mu1
        * 8.0
        * torch.pi
        * torch.pi
        * torch.sin(2.0 * torch.pi * x1)
        * torch.sin(2.0 * torch.pi * x2)
    )


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


domain_x = Square2D([(0.0, 1), (0.0, 1)], is_main_domain=True)
# add a hole
domain_x.add_hole(Disk2D([0.5, 0.5], 0.2))
domain_x.set_mapping(
    Mapping.rot_2d(
        torch.tensor(torch.pi / 3),
        center=torch.tensor([0.5, 0.5], dtype=torch.get_default_dtype()),
    ),
    [[-0.2, 1.2], [-0.2, 1.2]],
)
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler([(1.0, 2.0)])]
)


space = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[20] * 5,
)
pde = Laplacian2DDirichletStrongForm(space, f=f_rhs, g=f_bc)

use_natural_gradient = True

if use_natural_gradient:
    pinns = NaturalGradientPinnsElliptic(
        pde, bc_type="weak", bc_weight=40, default_lr=1e-3
    )
else:
    losses = GenericLosses(
        [
            ("residual", torch.nn.MSELoss(), 1.0),
            ("bc", torch.nn.MSELoss(), 40.0),
        ],
    )
    opt_1 = {
        "name": "adam",
        "optimizer_args": {"lr": 2.5e-2, "betas": (0.9, 0.999)},
    }
    opt_2 = {
        "name": "lbfgs",
        "switch_at_epoch_ratio": 0.5,
    }
    opt = OptimizerData(opt_1, opt_2)
    pinns = PinnsElliptic(pde, bc_type="weak", optimizers=opt, losses=losses)

resume_solve = True
if resume_solve or not pinns.load(__file__, "weak"):
    pinns.solve(epochs=100, n_collocation=3000, n_bc_collocation=1500, verbose=True)
    pinns.save(__file__, "weak")

pinns.space.load_from_best_approx()

plot_abstract_approx_space(
    pinns.space,
    domain_x,
    [[1.0, 2.0]],
    loss=pinns.losses,
    residual=pde,
    cuts=[
        ([0.0, 0.0], [-0.5, 0.5]),
        # ([0.0, 0.5], [-0.5, 0.5]),
        ([0.0, 0.5], [-0.5, -0.5]),
    ],
    draw_contours=True,
    n_drawn_contours=20,
    parameters_values="random",
)

plt.show()