r"""Solves a 2D multiscale Poisson PDE using PINNs.

.. 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 = (0, 1) \times (0, 1)`, :math:`f` such that :math:`u(x_1, x_2, \mu) = \mu \sin(8\pi x_1) \sin(8\pi x_2)`, :math:`g = 0` and :math:`\mu = \frac 1 2`.

Boundary conditions are enforced strongly.

The neural network used is a either a simple MLP (Multilayer Perceptron) or a ResNet (Residual Network); the optimization is done either using Adam or using natural gradient preconditioning.

Since the solution is quite oscillatory, the use of Fourier features or multiscale Fourier features is recommended; these two strategies are compared here to a classical MLP and to a SIREN network.
"""

# %%

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.nn_space import NNxSpace
from scimba_torch.domain.meshless_domain.domain_1d import Segment1D
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 (
    FourierMLP,
    FourierResNet,
    MultiScaleFourierMLP,
    MultiScaleFourierResNet,
)
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.neural_nets.coordinates_based_nets.siren import SirenNet
from scimba_torch.numerical_solvers.elliptic_pde.pinns import (
    NaturalGradientPinnsElliptic,
    PinnsElliptic,
)
from scimba_torch.optimizers.optimizers_data import OptimizerData
from scimba_torch.physical_models.elliptic_pde.abstract_elliptic_pde import (
    StrongFormEllipticPDE,
)
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

torch.manual_seed(0)


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


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


def exact_sol(x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    mu1 = mu.get_components()
    return mu1 * torch.sin(8 * torch.pi * x1) * torch.sin(8 * torch.pi * x2)


a = 1.0


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


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


opt_1 = {
    "name": "adam",
    "optimizer_args": {"lr": 1.2e-2, "betas": (0.9, 0.999)},
}
opt_2 = {
    "name": "lbfgs",
    "switch_at_epoch_ratio": 0.7,
}

domain_x = Square2D([(0.0, a), (0.0, a)], is_main_domain=True)
domain_mu = [(0.5, 0.5 + 1e-4)]
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler(domain_mu)]
)

# %%

#### first space: classical MLP #####
space = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[40, 40, 40],
    post_processing=post_processing,
)
pde = Laplacian2DDirichletStrongForm(
    space,
    f=f_rhs,
    g=f_bc,
)
pinns = PinnsElliptic(pde, bc_type="strong", optimizers=OptimizerData(opt_1))

new_solve = False
if (new_solve) or not pinns.load(__file__, "MLP_pinns_no_precond"):
    pinns.solve(epochs=2000, n_collocation=4000, verbose=True)
    pinns.save(__file__, "MLP_pinns_no_precond")

spaceENG = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[40, 40],
    post_processing=post_processing,
)
pdeENG = Laplacian2DDirichletStrongForm(
    spaceENG,
    f=f_rhs,
    g=f_bc,
)
pinnsENG = NaturalGradientPinnsElliptic(
    pdeENG,
    bc_type="strong",
    functional_post_processing=functional_post_processing,
    matrix_regularization=1e-4,
)

new_solve = False
if (new_solve) or not pinnsENG.load(__file__, "MLP_pinns_ENG"):
    pinnsENG.solve(epochs=100, n_collocation=2000, verbose=True)
    pinnsENG.save(__file__, "MLP_pinns_ENG")

plot_abstract_approx_spaces(
    (pinns.space, pinnsENG.space),
    domain_x,
    domain_mu,
    loss=(pinns.losses, pinnsENG.losses),
    residual=(pde, pdeENG),
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
    title=r"Approximating $-\mu\Delta u(x) = 128 \mu^2 \pi^2 sin(8 \pi x)$",
    titles=("Generic MLP, no preconditioning", "Generic MLP, ENG preconditioning"),
)
plt.show()

# %%

#### second space: MLP with Fourier features #####
space2 = NNxSpace(
    1,
    1,
    FourierMLP,
    domain_x,
    sampler,
    layer_sizes=[40, 40, 40],
    nb_features=20,
    mean=0.0,
    std=6.0,
    post_processing=post_processing,
)
pde2 = Laplacian2DDirichletStrongForm(
    space2,
    f=f_rhs,
    g=f_bc,
)
pinns2 = PinnsElliptic(pde2, bc_type="strong", optimizers=OptimizerData(opt_1))

new_solve = False
if (new_solve) or not pinns2.load(__file__, "FourierMLP_pinns_no_precond"):
    pinns2.solve(epochs=2000, n_collocation=4000, verbose=True)
    pinns2.save(__file__, "FourierMLP_pinns_no_precond")

space2ENG = NNxSpace(
    1,
    1,
    FourierMLP,
    domain_x,
    sampler,
    layer_sizes=[40],
    nb_features=10,
    mean=0.0,
    std=6.0,
    post_processing=post_processing,
)
pde2ENG = Laplacian2DDirichletStrongForm(
    space2ENG,
    f=f_rhs,
    g=f_bc,
)
pinns2ENG = NaturalGradientPinnsElliptic(
    pde2ENG,
    bc_type="strong",
    functional_post_processing=functional_post_processing,
    matrix_regularization=1e-4,
)

new_solve = False
if (new_solve) or not pinns2ENG.load(__file__, "FourierMLP_pinns_ENG"):
    pinns2ENG.solve(epochs=200, n_collocation=2000, verbose=True)
    pinns2ENG.save(__file__, "FourierMLP_pinns_ENG")

plot_abstract_approx_spaces(
    (pinns2.space, pinns2ENG.space),
    domain_x,
    domain_mu,
    loss=(pinns2.losses, pinns2ENG.losses),
    residual=(pde2, pde2ENG),
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
    title=r"Approximating $-\mu\Delta u(x) = 128 \mu^2 \pi^2 sin(8 \pi x)$",
    titles=(
        "MLP with Fourier features, no preconditioning",
        "MLP with Fourier features, ENG preconditioning",
    ),
)
plt.show()

# %%

#### second space, bis: ResNet with Fourier features #####
space2b = NNxSpace(
    1,
    1,
    FourierResNet,
    domain_x,
    sampler,
    layer_structure=[20, 5, [1, 3], [3, 5]],
    nb_features=20,
    mean=0.0,
    std=6.0,
    post_processing=post_processing,
)
pde2b = Laplacian2DDirichletStrongForm(
    space2b,
    f=f_rhs,
    g=f_bc,
)
pinns2b = PinnsElliptic(pde2b, bc_type="strong", optimizers=OptimizerData(opt_1))

new_solve = True
if (new_solve) or not pinns2b.load(__file__, "FourierResNet_pinns_no_precond"):
    pinns2b.solve(epochs=2000, n_collocation=4000, verbose=True)
    pinns2b.save(__file__, "FourierResNet_pinns_no_precond")

space2bENG = NNxSpace(
    1,
    1,
    FourierResNet,
    domain_x,
    sampler,
    layer_structure=[20, 5, [1, 3], [3, 5]],
    nb_features=10,
    mean=0.0,
    std=6.0,
    post_processing=post_processing,
)
pde2bENG = Laplacian2DDirichletStrongForm(
    space2bENG,
    f=f_rhs,
    g=f_bc,
)
pinns2bENG = NaturalGradientPinnsElliptic(
    pde2bENG,
    bc_type="strong",
    functional_post_processing=functional_post_processing,
    matrix_regularization=1e-4,
)

new_solve = True
if (new_solve) or not pinns2bENG.load(__file__, "FourierResNet_pinns_ENG"):
    pinns2bENG.solve(epochs=200, n_collocation=2000, verbose=True)
    pinns2bENG.save(__file__, "FourierResNet_pinns_ENG")

plot_abstract_approx_spaces(
    (pinns2b.space, pinns2bENG.space),
    domain_x,
    domain_mu,
    loss=(pinns2b.losses, pinns2bENG.losses),
    residual=(pde2b, pde2bENG),
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
    title=r"Approximating $-\mu\Delta u(x) = 128 \mu^2 \pi^2 sin(8 \pi x)$",
    titles=(
        "ResNet with Fourier features, no preconditioning",
        "ResNet with Fourier features, ENG preconditioning",
    ),
)
plt.show()

# %%

#### third space: MLP with multiscale Fourier features #####
space3 = NNxSpace(
    1,
    1,
    MultiScaleFourierMLP,
    domain_x,
    sampler,
    layer_sizes=[40, 40, 40],
    nb_features=20,
    means=[0.0, 0.0],
    stds=[1.5, 10.0],
    post_processing=post_processing,
)
pde3 = Laplacian2DDirichletStrongForm(
    space3,
    f=f_rhs,
    g=f_bc,
)
pinns3 = PinnsElliptic(pde3, bc_type="strong", optimizers=OptimizerData(opt_1))

new_solve = False
if (new_solve) or not pinns3.load(__file__, "MultiScaleFourierMLP_pinns_no_precond"):
    pinns3.solve(epochs=2000, n_collocation=4000, verbose=True)
    pinns3.save(__file__, "MultiScaleFourierMLP_pinns_no_precond")

space3ENG = NNxSpace(
    1,
    1,
    MultiScaleFourierMLP,
    domain_x,
    sampler,
    layer_sizes=[40],
    nb_features=10,
    means=[0.0, 0.0],
    stds=[1.5, 10.0],
    post_processing=post_processing,
)
pde3ENG = Laplacian2DDirichletStrongForm(
    space3ENG,
    f=f_rhs,
    g=f_bc,
)
pinns3ENG = NaturalGradientPinnsElliptic(
    pde3ENG,
    bc_type="strong",
    functional_post_processing=functional_post_processing,
    matrix_regularization=1e-2,
)

new_solve = False
if (new_solve) or not pinns3ENG.load(__file__, "MultiScaleFourierMLP_pinns_ENG"):
    pinns3ENG.solve(epochs=100, n_collocation=2000, verbose=True)
    pinns3ENG.save(__file__, "MultiScaleFourierMLP_pinns_ENG")

plot_abstract_approx_spaces(
    (pinns3.space, pinns3ENG.space),
    domain_x,
    domain_mu,
    loss=(pinns3.losses, pinns3ENG.losses),
    residual=(pde3, pde3ENG),
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
    title=r"Approximating $-\mu\Delta u(x) = 128 \mu^2 \pi^2 sin(8 \pi x)$",
    titles=(
        "MLP with multi scale Fourier features, no preconditioning",
        "MLP with multi scale Fourier features, ENG preconditioning",
    ),
)
plt.show()

# %%

#### third space, bis: ResNet with multiscale Fourier features #####
space3b = NNxSpace(
    1,
    1,
    MultiScaleFourierResNet,
    domain_x,
    sampler,
    layer_structure=[20, 5, [1, 3], [3, 5]],
    nb_features=20,
    means=[0.0, 0.0],
    stds=[1.5, 10.0],
    post_processing=post_processing,
)
pde3b = Laplacian2DDirichletStrongForm(
    space3b,
    f=f_rhs,
    g=f_bc,
)
pinns3b = PinnsElliptic(pde3b, bc_type="strong", optimizers=OptimizerData(opt_1))

new_solve = False
if (new_solve) or not pinns3b.load(
    __file__, "MultiScaleFourierResNet_pinns_no_precond"
):
    pinns3b.solve(epochs=2000, n_collocation=4000, verbose=True)
    pinns3b.save(__file__, "MultiScaleFourierResNet_pinns_no_precond")

space3bENG = NNxSpace(
    1,
    1,
    MultiScaleFourierResNet,
    domain_x,
    sampler,
    layer_structure=[20, 5, [1, 3], [3, 5]],
    nb_features=10,
    means=[0.0, 0.0],
    stds=[1.5, 10.0],
    post_processing=post_processing,
)
pde3bENG = Laplacian2DDirichletStrongForm(
    space3bENG,
    f=f_rhs,
    g=f_bc,
)
pinns3bENG = NaturalGradientPinnsElliptic(
    pde3bENG,
    bc_type="strong",
    functional_post_processing=functional_post_processing,
    matrix_regularization=1e-2,
)

new_solve = False
if (new_solve) or not pinns3bENG.load(__file__, "MultiScaleFourierResNet_pinns_ENG"):
    pinns3bENG.solve(epochs=100, n_collocation=2000, verbose=True)
    pinns3bENG.save(__file__, "MultiScaleFourierResNet_pinns_ENG")

plot_abstract_approx_spaces(
    (pinns3b.space, pinns3bENG.space),
    domain_x,
    domain_mu,
    loss=(pinns3b.losses, pinns3bENG.losses),
    residual=(pde3b, pde3bENG),
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
    title=r"Approximating $-\mu\Delta u(x) = 128 \mu^2 \pi^2 sin(8 \pi x)$",
    titles=(
        "ResNet with multi scale Fourier features, no preconditioning",
        "ResNet with multi scale Fourier features, ENG preconditioning",
    ),
)
plt.show()

# %%

#### fourth space: SIREN #####
space4 = NNxSpace(
    1,
    1,
    SirenNet,
    domain_x,
    sampler,
    layer_sizes=[40, 40, 40],
    w=1,
    w0=30,
    post_processing=post_processing,
)
pde4 = Laplacian2DDirichletStrongForm(
    space4,
    f=f_rhs,
    g=f_bc,
)
pinns4 = PinnsElliptic(pde4, bc_type="strong", optimizers=OptimizerData(opt_1))

new_solve = False
if (new_solve) or not pinns4.load(__file__, "Siren_pinns_no_precond"):
    pinns4.solve(epochs=2000, n_collocation=4000, verbose=True)
    pinns4.save(__file__, "Siren_pinns_no_precond")

space4ENG = NNxSpace(
    1,
    1,
    SirenNet,
    domain_x,
    sampler,
    layer_sizes=[40],
    w=1,
    w0=30,
    post_processing=post_processing,
)
pde4ENG = Laplacian2DDirichletStrongForm(
    space4ENG,
    f=f_rhs,
    g=f_bc,
)
pinns4ENG = NaturalGradientPinnsElliptic(
    pde4ENG,
    bc_type="strong",
    functional_post_processing=functional_post_processing,
    matrix_regularization=1e-2,
)

new_solve = False
if (new_solve) or not pinns4ENG.load(__file__, "Siren_pinns_ENG"):
    pinns4ENG.solve(epochs=200, n_collocation=2000, verbose=True)
    pinns4ENG.save(__file__, "Siren_pinns_ENG")

plot_abstract_approx_spaces(
    (pinns4.space, pinns4ENG.space),
    domain_x,
    domain_mu,
    loss=(pinns4.losses, pinns4ENG.losses),
    residual=(pde4, pde4ENG),
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
    title=r"Approximating $-\mu\Delta u(x) = 128 \mu^2 \pi^2 sin(8 \pi x)$",
    titles=("Siren net, no preconditioning", "Siren net, ENG preconditioning"),
)
plt.show()

# %%


class Laplacian1DDirichletStrongForm2Params(StrongFormEllipticPDE):
    def __init__(self, space, f, g, **kwargs):
        super().__init__(space, residual_size=1, bc_residual_size=1, **kwargs)
        self.f = f
        self.g = g

    def rhs(self, w, x, mu):
        return self.f(x, mu)

    def operator(self, w, x, mu):
        u = w.get_components()
        # mu1, mu2 = mu.get_components()
        u_xx = self.grad(self.grad(u, x), x)
        return -u_xx

    def functional_operator(
        self,
        func,
        x: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        grad_u = torch.func.jacrev(func, 0)
        grad_grad_u = torch.func.jacrev(grad_u, 0)(x, mu, theta)
        return -grad_grad_u[..., 0, 0]

    def bc_rhs(self, w, x, n, mu):
        return self.g(x, mu)

    def bc_operator(self, w, x, n, mu):
        return w.get_components()

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


mu1 = 2.0
mu2 = 20.0
mu3 = 0.04


def f_rhs_1d(x: LabelTensor, mu: LabelTensor):
    x1 = x.get_components()
    return torch.pi**2 * mu1**2 * torch.sin(
        mu1 * torch.pi * x1
    ) + mu3 * torch.pi**2 * mu2**2 * torch.sin(mu2 * torch.pi * x1)


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


def exact_sol_1d(x: LabelTensor, mu: LabelTensor):
    x1 = x.get_components()
    # mu1, mu2 = mu.get_components()
    return torch.sin(mu1 * torch.pi * x1) + mu3 * torch.sin(mu2 * torch.pi * x1)


a = 1.0


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


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


domain_x = Segment1D((0.0, a), is_main_domain=True)
# domain_mu = [(2., 2. + 1e-4), (20., 20. + 1e-4)]
domain_mu = []
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler(domain_mu)]
)

# %%

### first space: multiscale MLP with Fourier features #####
space5 = NNxSpace(
    1,
    0,
    MultiScaleFourierMLP,
    domain_x,
    sampler,
    layer_sizes=[40, 40, 40, 40],
    nb_features=20,
    means=[0.0, 0.0],
    stds=[1.0, 60.0],
    post_processing=post_processing_1d,
)
pde5 = Laplacian1DDirichletStrongForm2Params(
    space5,
    f=f_rhs_1d,
    g=f_bc_1d,
)
pinns5 = PinnsElliptic(pde5, bc_type="strong", optimizers=OptimizerData(opt_1))

new_solve = False
if (new_solve) or not pinns5.load(__file__, "1D_MultiScaleFourier_pinns_no_precond"):
    pinns5.solve(epochs=2000, n_collocation=4000, verbose=True)
    pinns5.save(__file__, "1D_MultiScaleFourier_pinns_no_precond")

space5ENG = NNxSpace(
    1,
    0,
    MultiScaleFourierMLP,
    domain_x,
    sampler,
    layer_sizes=[20, 20],
    nb_features=10,
    means=[0.0, 0.0],
    stds=[1.0, 60.0],
    post_processing=post_processing_1d,
)
pde5ENG = Laplacian1DDirichletStrongForm2Params(
    space5ENG,
    f=f_rhs_1d,
    g=f_bc_1d,
)
pinns5ENG = NaturalGradientPinnsElliptic(
    pde5ENG,
    bc_type="strong",
    functional_post_processing=functional_post_processing_1d,
    matrix_regularization=1e-4,
)

new_solve = False
if (new_solve) or not pinns5ENG.load(__file__, "1D_MultiScaleFourier_pinns_ENG"):
    pinns5ENG.solve(epochs=100, n_collocation=900, verbose=True)
    pinns5ENG.save(__file__, "1D_MultiScaleFourier_pinns_ENG")

plot_abstract_approx_spaces(
    (pinns5.space, pinns5ENG.space),
    domain_x,
    domain_mu,
    loss=(pinns5.losses, pinns5ENG.losses),
    residual=(pde5, pde5ENG),
    solution=exact_sol_1d,
    error=exact_sol_1d,
    draw_contours=True,
    n_drawn_contours=20,
)
plt.show()

# %%