r"""Solves a 1D elliptic PDE with Dirichlet BCs using the Deep Ritz method and PINNs.

.. math::

    -\mu (Au')' & = f in \Omega \times M \\
    u & = g on \partial \Omega \times M

where :math:`x \in \Omega = (0, 1)`, :math:`A = 1`, :math:`f(x, \mu) = (2 \mu \pi)^2 \sin(2\pi x)`, :math:`g = 0` and :math:`\mu \in M = [0.5, 1]`. Boundary conditions are enforced weakly.
"""

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.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.deep_ritz import (
    DeepRitzElliptic,
    NaturalGradientDeepRitzElliptic,
)
from scimba_torch.numerical_solvers.elliptic_pde.pinns import (
    NaturalGradientPinnsElliptic,
)
from scimba_torch.physical_models.elliptic_pde.linear_order_2 import (
    DivAGradUPDE,
)
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces
from scimba_torch.utils.scimba_tensors import LabelTensor

# maple code:
# u(x):= mu*sin(2*Pi*x);
# f(x) := simplify(-mu*( diff(diff(u(x),x),x) ));
# PDE := -mu*( diff(diff(v(x),x),x) ) = f(x);
# test:= v(x)= u(x);
# pdetest(test, PDE);


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


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


def f_bc(x: LabelTensor, mu: LabelTensor):
    x1 = x.get_components()
    # mu1 = mu.get_components()
    return torch.zeros_like(x1)


def A(x: torch.Tensor) -> torch.Tensor:  # noqa: N802
    return torch.ones(
        x.shape[0],
        1,
        1,
        dtype=torch.get_default_dtype(),
        device=torch.get_default_device(),
    )


domain_x = Segment1D((0, 1), is_main_domain=True)
# domain_mu= []
domain_mu = [[0.5, 1.0]]
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler(domain_mu)]
)

space = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32],
)

pde = DivAGradUPDE(space, 1, f=f_rhs, g=f_bc, A=A)

pinns = NaturalGradientPinnsElliptic(
    pde,
    bc_type="weak",
)

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

space2 = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[16, 16],
)

pde2 = DivAGradUPDE(space2, 1, f=f_rhs, g=f_bc, A=A)

bc_weight = 100.0

ritz = DeepRitzElliptic(pde2, bc_type="weak", bc_weight=bc_weight)

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

space3 = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[16, 16],
)

pde3 = DivAGradUPDE(space3, 1, f=f_rhs, g=f_bc, A=A)

ritz2 = NaturalGradientDeepRitzElliptic(pde3, bc_type="weak", bc_weight=bc_weight)

new_solve = True
if new_solve or not ritz2.load(__file__, "ritz_weak_ENG"):
    ritz2.solve(epochs=200, n_collocation=4000, n_bc_collocation=2000, verbose=True)
    ritz2.save(__file__, "ritz_weak_ENG")

plot_abstract_approx_spaces(
    (
        pinns.space,
        ritz.space,
        ritz2.space,
    ),
    (domain_x,),
    (domain_mu,),
    loss=(
        pinns.losses,
        ritz.losses,
        ritz2.losses,
    ),
    error=exact_sol,
    draw_contours=True,
    n_drawn_contours=20,
    parameters_values="mean",
    title=r"Solving $-\mu \\Delta u = (2 \mu \pi)^2 \sin(2\pi x) $",
    titles=(
        "PINN with ENG preconditioning",
        "RITZ with no preconditioning",
        "RITZ with ENG preconditioning",
    ),
)
plt.show()