r"""Solves a 2D Poisson PDE with Dirichlet BCs and two parameters.

.. math::

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

where :math:`x = (x_1, x_2) \in \Omega = \mathcal{D}` (with \mathcal{D} being the disk with center :math:`(x_1^0, x_2^0) = (-0.5, 0.5)`), :math:`f` such that :math:`u(x_1, x_2, \mu) = \mu_1 \exp{- \mu_2 (x_1^2 + x_2^2)}`, :math:`g = 0` and :math:`\mu \in M = [1, 2] \times [0.1, 0.5]`. Boundary conditions are enforced strongly.

The neural network used is a simple MLP (Multilayer Perceptron), 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 Disk2D
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.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
from scimba_torch.utils.verbosity import set_verbosity

set_verbosity(True)


class Laplacian2DDirichletStrongForm2Params(Laplacian2DDirichletStrongForm):
    def operator(self, w, x: LabelTensor, mu: LabelTensor) -> torch.Tensor:
        """Compute the differential operator of the PDE.

        Args:
            w: State tensor.
            x: Spatial coordinates tensor.
            mu: Parameter tensor.

        Returns:
            The result of applying the operator to the state.
        """
        u = w.get_components()
        alpha, sigma = mu.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 -alpha * (u_xx + u_yy)


def exact_sol(x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    mu1, sigma = mu.get_components()
    return mu1 * torch.exp(-sigma * (x1**2 + x2**2))


boundary_value = exact_sol(
    LabelTensor(torch.tensor([[0.0, 1.0]], dtype=torch.double)),
    LabelTensor(torch.tensor([[1.0, 0.25]], dtype=torch.double)),
)


def f_rhs(x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    mu1, sigma = mu.get_components()
    # return torch.ones_like(x1)
    a = x1**2 + x2**2
    return -4 * mu1**2 * sigma * (-1.0 + sigma * a) * torch.exp(-sigma * a)


def f_bc(x: LabelTensor, mu: LabelTensor):
    x1, _ = x.get_components()
    mu1, sigma = mu.get_components()
    boundary_value = mu1 * torch.exp(-sigma)
    return x1 * 0.0 + boundary_value


domain_x = Disk2D((0.0, 0.0), 1, is_main_domain=True)
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler([[1.0, 2.0], [0.1, 0.5]])]
)


#### first space: Adam ####


def post_processing(inputs: torch.Tensor, x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    mu1, sigma = mu.get_components()
    phi = x1**2.0 + x2**2.0 - 1.0
    boundary_value = mu1 * torch.exp(-sigma)
    return inputs * phi + boundary_value


def functional_post_processing(
    func, x: torch.Tensor, mu: torch.Tensor, theta: torch.Tensor
) -> torch.Tensor:
    phi = x[0] ** 2.0 + x[1] ** 2.0 - 1.0
    boundary_value = mu[0] * torch.exp(-mu[1])
    return func(x, mu, theta) * phi + boundary_value


space = NNxSpace(
    1,
    2,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[20] * 3,
    post_processing=post_processing,
)

pde = Laplacian2DDirichletStrongForm2Params(space, f=f_rhs, g=f_bc)

losses = GenericLosses([("residual", torch.nn.MSELoss(), 1.0)])
opt = {
    "name": "adam",
    "optimizer_args": {"lr": 1.8e-2, "betas": (0.9, 0.999)},
}
pinns = PinnsElliptic(pde, bc_type="strong", optimizers=opt, losses=losses)

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


#### second space: natural gradient descent ####

space2 = NNxSpace(
    1,
    2,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[20] * 3,
    post_processing=post_processing,
)

pde2 = Laplacian2DDirichletStrongForm2Params(space2, f=f_rhs, g=f_bc)

pinns2 = NaturalGradientPinnsElliptic(
    pde2, bc_type="strong", functional_post_processing=functional_post_processing
)

new_solve = False
if new_solve or not pinns2.load(__file__, "ENG"):
    pinns2.solve(epochs=200, n_collocation=900, verbose=False)
    pinns2.save(__file__, "ENG")

title = (
    "Solving "
    + r"$-\mu\nabla^2 u = \mu(1-\sigma(x^2 + y^2))e^{-\sigma(x^2 + y^2)}$"
    + " on the unit disk"
)


def norm_of_grad(space, xy, mu):
    xy.x.requires_grad_()
    w = space.evaluate(xy, mu)
    u = w.get_components()
    ux, uy = space.grad(u, xy)
    return torch.sqrt(ux**2 + uy**2)


additional_scalar_functions = {"$\\|\\nabla u\\|_2$": norm_of_grad}

plot_abstract_approx_spaces(
    (
        pinns.space,
        pinns2.space,
    ),  # an Iterable of AbstractSpace
    (
        domain_x,
    ),  # either a VolumetricDomain, or an Iterable of VolumetricDomain of length 1 or len(first argument)
    ([[1.0, 2.0], [0.1, 0.5]],),
    loss=(
        pinns.losses,
        pinns2.losses,
    ),  # same as previously; if only one is given, it will be used for all spaces
    residual=(
        pinns.pde,
        pinns2.pde,
    ),  # same as previously; if only one is given, it will be used for all spaces
    error=exact_sol,
    # solution=exact_sol,
    additional_scalar_functions=additional_scalar_functions,
    draw_contours=True,
    n_drawn_contours=20,
    parameters_values="mean",
    title=title,
    titles=("not preconditioned PINN", "ENG preconditioned PINN"),
)
plt.show()