r"""Solves the 2D Helmholtz PDE with homogeneous Dirichlet BCs using PINNs.

.. math::

    \delta u + (2 \pi k)^2 u & = f in \Omega \\
    u & = 0 on \partial \Omega

where :math:`x = (x_1, x_2) \in \Omega = (-1, 1) \times (-1, 1)` and :math:`k = 2`.

The boundary conditions are enforced strongly.

The neural network used is a simple MLP (Multilayer Perceptron) with multiscale Fourier features, and the optimization is done using either Adam, Natural Gradient Descent, or Anagram.
"""

from typing import Callable, Tuple

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.abstract_space import AbstractApproxSpace
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.features import (
    MultiScaleFourierMLP,
)
from scimba_torch.numerical_solvers.elliptic_pde.pinns import (
    AnagramPinnsElliptic,
    NaturalGradientPinnsElliptic,
    PinnsElliptic,
)
from scimba_torch.optimizers.optimizers_data import OptimizerData
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces
from scimba_torch.utils.scimba_tensors import LabelTensor, MultiLabelTensor
from scimba_torch.utils.typing_protocols import VarArgCallable

torch.manual_seed(0)


class HelmholtzDirichlet:
    def __init__(
        self,
        space: AbstractApproxSpace,
        spatial_dim: int,
        f: Callable,
        g: Callable,
        **kwargs,
    ):
        self.space = space
        self.spatial_dim = spatial_dim
        self.f = f
        self.g = g

        self.linear = True
        self.residual_size = 1
        self.bc_residual_size = 1

    def grad(
        self,
        w: torch.Tensor | MultiLabelTensor,
        y: torch.Tensor | LabelTensor,
    ) -> torch.Tensor | Tuple[torch.Tensor, ...]:
        return self.space.grad(w, y)

    def rhs(self, w: MultiLabelTensor, x: LabelTensor, mu: LabelTensor) -> torch.Tensor:
        return self.f(x, mu)

    def bc_rhs(
        self, w: MultiLabelTensor, x: LabelTensor, n: LabelTensor, mu: LabelTensor
    ) -> LabelTensor:
        return self.g(x, mu)

    def operator(
        self, w: MultiLabelTensor, x: LabelTensor, mu: LabelTensor
    ) -> torch.Tensor:
        u = w.get_components()
        k = mu.get_components()

        omega = 2.0 * torch.pi * k
        omega2_u = omega**2 * u

        grad_u = torch.cat(tuple(self.grad(u, x)), dim=-1)
        div_grad_u = tuple(self.grad(grad_u[:, 0], x))[0]
        for i in range(1, self.spatial_dim):
            div_grad_u += tuple(self.grad(grad_u[:, i], x))[i]

        return omega2_u + div_grad_u

    def functional_operator(
        self,
        func: VarArgCallable,
        x: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        omega = 2.0 * torch.pi * mu[0]
        omega2_u = omega**2 * func(x, mu, theta)

        grad_u = torch.func.jacrev(func, 0)
        grad_grad_u = torch.func.jacrev(grad_u, 0)(x, mu, theta).squeeze()
        div_grad_u = torch.einsum("ii", grad_grad_u)

        return omega2_u + div_grad_u

    # Dirichlet conditions
    def bc_operator(
        self, w: MultiLabelTensor, x: LabelTensor, n: LabelTensor, mu: LabelTensor
    ) -> torch.Tensor:
        u = w.get_components()
        return u

    def functional_operator_bc(
        self,
        func: VarArgCallable,
        x: torch.Tensor,
        n: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        return func(x, mu, theta)


# bc_weight = 40.0

sigma = 0.03
center = [0.0, 0.5]


def f_rhs(x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    # mu1 = mu.get_components()
    return (
        -100
        * torch.exp(
            -((x1 - center[0]) ** 2 + (x2 - center[1]) ** 2) / (2.0 * sigma**2.0)
        )
        / ((2.0 * torch.pi) * sigma)
    )


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


a = 1.0


def post_processing(inputs: torch.Tensor, x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    return inputs * (-a - x1) * (a - x1) * (-a - x2) * (a - x2)


def functional_post_processing(func, x, mu, theta) -> torch.Tensor:
    # print("x.shape: ", x.shape)
    return func(x, mu, theta) * (-a - x[0]) * (a - x[0]) * (-a - x[1]) * (a - x[1])


# sigma := 0.03
# f(x,y) := (100/(2*Pi*sigma))*exp( -( x^2 + (y-0.9)^2 ) / (2*sigma^2) );
# omega := (Pi*4)^2;
# PDE := -( diff(diff(u(x,y),x),x) + diff(diff(u(x,y),y),y) + omega*u(x,y) ) = f(x,y);
# BC := u(-1,y)=0, u(1,y)=0, u(x,-1)=0, u(x,1)=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_mu = [(2.0, 2.0 + 1e-4)]
domain_x = Square2D([(-a, a), (-a, a)], is_main_domain=True)
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler(domain_mu)]
)

space = NNxSpace(
    1,
    1,
    MultiScaleFourierMLP,
    domain_x,
    sampler,
    # layer_sizes=[40, 60, 60, 60, 40],
    layer_sizes=[50, 50],
    means=[0.0, 0.0],
    stds=[4.0, 10.0],
    nb_features=40,
    activation_type="sine",
    post_processing=post_processing,
)

pde = HelmholtzDirichlet(space, 2, f=f_rhs, g=f_bc)
opt_1 = {
    "name": "adam",
    "optimizer_args": {"lr": 9.9e-3, "betas": (0.9, 0.999)},
}
opt_2 = {
    "name": "lbfgs",
    "switch_at_epoch_ratio": 0.7,
}
pinns = PinnsElliptic(
    pde,
    bc_type="strong",
    optimizers=OptimizerData(opt_1, opt_2),
)

solving_mode = "new"
# can be "new"    for solve from scratch,
#       "resume" for load from file and continue solving,
#    or "load"   for load from file
load_post_fix = "pinns"
# the post_fix of the file where to load the pinn
save_post_fix = "pinns"
# the post_fix of the file where to save the pinn

new_solving = solving_mode == "new"
resume_solving = solving_mode == "resume"
try_load = resume_solving or (not new_solving)
load_ok = try_load and pinns.load(__file__, load_post_fix)
solve = (not load_ok) or (resume_solving or new_solving)

if solve:
    pinns.solve(epochs=9, n_collocation=8000, verbose=True)
    pinns.save(__file__, save_post_fix)

pinns.space.load_from_best_approx()

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

space2 = NNxSpace(
    1,
    1,
    MultiScaleFourierMLP,
    domain_x,
    sampler,
    layer_sizes=[20, 20],
    means=[0.0, 0.0],
    stds=[4.0, 8.0],
    nb_features=20,
    activation_type="sine",
    post_processing=post_processing,
)

pde2 = HelmholtzDirichlet(space2, 2, f=f_rhs, g=f_bc)
#
pinns2 = NaturalGradientPinnsElliptic(
    pde2,
    bc_type="strong",
    functional_post_processing=functional_post_processing,
    matrix_regularization=2e-4,
)

solving_mode = "new"
# can be "new"    for solve from scratch,
#       "resume" for load from file and continue solving,
#    or "load"   for load from file
load_post_fix = "ENG"
# the post_fix of the file where to load the pinn
save_post_fix = "ENG"
# the post_fix of the file where to save the pinn

new_solving = solving_mode == "new"
resume_solving = solving_mode == "resume"
try_load = resume_solving or (not new_solving)
load_ok = try_load and pinns2.load(__file__, load_post_fix)
solve = (not load_ok) or (resume_solving or new_solving)

if solve:
    pinns2.solve(epochs=300, n_collocation=10000, verbose=True)
    pinns2.save(__file__, save_post_fix)
#
pinns2.space.load_from_best_approx()

space3 = NNxSpace(
    1,
    1,
    MultiScaleFourierMLP,
    domain_x,
    sampler,
    layer_sizes=[24, 24],
    means=[0.0, 0.0],
    stds=[4.0, 8.0],
    nb_features=20,
    activation_type="sine",
    post_processing=post_processing,
)

pde3 = HelmholtzDirichlet(space3, 2, f=f_rhs, g=f_bc)

pinns3 = AnagramPinnsElliptic(
    pde3,
    bc_type="strong",
    svd_threshold=1.0,
    functional_post_processing=functional_post_processing,
)

solving_mode = "new"
# can be "new"    for solve from scratch,
#       "resume" for load from file and continue solving,
#    or "load"   for load from file
load_post_fix = "Anagram"
# the post_fix of the file where to load the pinn
save_post_fix = "Anagram"
# the post_fix of the file where to save the pinn

new_solving = solving_mode == "new"
resume_solving = solving_mode == "resume"
try_load = resume_solving or (not new_solving)
load_ok = try_load and pinns3.load(__file__, load_post_fix)
solve = (not load_ok) or (resume_solving or new_solving)

if solve:
    pinns3.solve(epochs=10, n_collocation=4000, verbose=True)
    pinns3.save(__file__, save_post_fix)

pinns3.space.load_from_best_approx()


plot_abstract_approx_spaces(
    (
        pinns.space,
        pinns2.space,
        pinns3.space,
    ),
    domain_x,
    domain_mu,
    loss=(
        pinns.losses,
        pinns2.losses,
        pinns3.losses,
    ),
    residual=(
        pinns.pde,
        pinns2.pde,
        pinns3.pde,
    ),
    draw_contours=True,
    n_drawn_contours=20,
    parameters_values="mean",
    title="Helmholtz equation $\\Delta u + (2\\pi\\mu)^2 u = -\\frac{100}{2\\pi\\sigma}\\exp{-\\frac{x^2 + y^2}{2\\sigma^2}}$, $\\sigma = 0.03$",
    titles=("no preconditioner", "ENG preconditioner", "Anagram preconditioner"),
)

plt.show()