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 Natural Gradient Descent.

Four sampling strategies are used:
- first, we sample :math:`N` new points at each epoch;
- then, we use a fixed set of :math:`P \gg N` spatial points sampled before training and draw :math:`N` points from this set at each epoch;
- third, we do the same but we also pre-sample the parameters and not only the spatial points;
- finally, we use a fixed set of :math:`N` spatial points sampled before training and use all of them at each epoch.
"""

# %%

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

DOMAIN_X = Disk2D((0.0, 0.0), 1, is_main_domain=True)
DOMAIN_MU = [[1.0, 2.0], [0.1, 0.5]]

DEFAULT_SAMPLER_X = DomainSampler(DOMAIN_X)
DEFAULT_SAMPLER_MU = UniformParametricSampler(DOMAIN_MU)


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


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


def post_processing(inputs: torch.Tensor, x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    mu1, sigma = mu.get_components()
    phi = x1**2 + x2**2 - 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 + x[1] ** 2 - 1.0
    boundary_value = mu[0] * torch.exp(-mu[1])
    return func(x, mu, theta) * phi + boundary_value


def create_pinn(sampler):
    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)

    return NaturalGradientPinnsElliptic(
        pde,
        bc_type="strong",
        functional_post_processing=functional_post_processing,
    )


# %%  1/ RE-SAMPLE POINTS AT EACH EPOCH

sampler_re_sampled = TensorizedSampler([DEFAULT_SAMPLER_X, DEFAULT_SAMPLER_MU])

pinn_re_sampled = create_pinn(sampler_re_sampled)
pinn_re_sampled.solve(epochs=50, n_collocation=768, verbose=False)

# %%  2/ DEFINE SOME PRE-SAMPLED POINTS AND SAMPLE FROM THEM AT EACH EPOCH

sampler_x = DomainSampler(DOMAIN_X, pre_sampling=True, n_pre_sampled_points=10_000)
sampler_pre_sampled = TensorizedSampler([sampler_x, DEFAULT_SAMPLER_MU])

pinn_pre_sampled = create_pinn(sampler_pre_sampled)
pinn_pre_sampled.solve(epochs=50, n_collocation=768, verbose=False)

# %%  3/ DEFINE SOME PRE-SAMPLED POINTS AND PARAMETERS AND SAMPLE FROM THEM AT EACH EPOCH

sampler_x = DomainSampler(DOMAIN_X, pre_sampling=True, n_pre_sampled_points=10_000)
sampler_mu = UniformParametricSampler(
    DOMAIN_MU, pre_sampling=True, n_pre_sampled_points=5_000
)
sampler_pre_sampled_params = TensorizedSampler([sampler_x, sampler_mu])

pinn_pre_sampled_params = create_pinn(sampler_pre_sampled_params)
pinn_pre_sampled_params.solve(epochs=50, n_collocation=768, verbose=False)

# %%  4/ DEFINE SOME PRE-SAMPLED POINTS AND USE ONLY THESE POINTS AT EACH EPOCH

sampler_x = DomainSampler(DOMAIN_X, pre_sampling=True, n_pre_sampled_points=768)
sampler_pre_sampled_fixed = TensorizedSampler([sampler_x, DEFAULT_SAMPLER_MU])

pinn_pre_sampled_fixed = create_pinn(sampler_pre_sampled_fixed)
pinn_pre_sampled_fixed.solve(epochs=50, n_collocation=768, verbose=False)

# %%  5/ PLOT THE RESULTS

all_pinns = [
    pinn_re_sampled,
    pinn_pre_sampled,
    pinn_pre_sampled_params,
    pinn_pre_sampled_fixed,
]

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


plot_abstract_approx_spaces(
    [pinn.space for pinn in all_pinns],
    (DOMAIN_X,),
    (DOMAIN_MU,),
    loss=[pinn.losses for pinn in all_pinns],
    residual=[pinn.pde for pinn in all_pinns],
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
    parameters_values="mean",
    title=title,
    titles=(
        "PINN with re-sampling",
        "PINN with pre-sampled points",
        "PINN with pre-sampled points and parameters",
        "PINN with pre-sampled points (fixed)",
    ),
)
plt.show()

# %%