r"""Solves the scalar linear advection equation in 1D using a PINN.

.. math::

    \partial_t u + \mu \partial_x u & = f in \Omega \times (0, T) \times M \\
    u & = g on \partial \Omega \times (0, T) \times M \\
    u & = u_0 on \Omega \times {0} \times M

where :math:`u: \partial \Omega \times (0, T) \times M \to \mathbb{R}` is the unknown function, :math:`\Omega \subset \mathbb{R}` is the spatial domain, :math:`(0, T) \subset \mathbb{R}` is the time domain, and :math:`M` is the parameter domain. Dirichlet boundary conditions are prescribed.

The equation is solved on a segment domain; weak boundary and initial conditions are used.
Three training strategies are compared: standard PINNs, PINNs with energy natural gradient preconditioning and PINNs with Anagram preconditioning.
"""

# %%

import matplotlib.pyplot as plt
import torch

from scimba_torch.approximation_space.nn_space import NNxtSpace
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.integration.monte_carlo_time import UniformTimeSampler
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.numerical_solvers.temporal_pde.pinns import (
    AnagramTemporalPinns,
    NaturalGradientTemporalPinns,
    TemporalPinns,
)
from scimba_torch.optimizers.optimizers_data import OptimizerData
from scimba_torch.physical_models.temporal_pde.transport_equation import (
    Transport1D,
)
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(w, t, x, mu):
    x1 = x.get_components()
    return 0 * x1


def f_bc(w, t, x, n, mu):
    x1 = x.get_components()
    return 0 * x1


def functional_a(t: torch.Tensor, x: torch.Tensor, mu: torch.Tensor) -> torch.Tensor:
    return mu


def a(t: LabelTensor, x: LabelTensor, mu: LabelTensor) -> torch.Tensor:
    return functional_a(t.get_components(), x.get_components(), mu.get_components())


def functional_exact(
    t: torch.Tensor, x: torch.Tensor, mu: torch.Tensor
) -> torch.Tensor:
    a = functional_a(t, x, mu)
    t0 = 0.015
    den = 4.0 * mu * t0
    return (1.0 / torch.sqrt(torch.pi * den)) * torch.exp(
        -((x - 1.0 - a * t) ** 2) / den
    )


def exact(t: LabelTensor, x: LabelTensor, mu: LabelTensor) -> torch.Tensor:
    return functional_exact(t.get_components(), x.get_components(), mu.get_components())


def f_ini(x, mu):
    t = LabelTensor(torch.zeros_like(x.x))
    return exact(t, x, mu)


xmax = 2.0

# def post_processing( inputs: torch.Tensor, t: LabelTensor, x: LabelTensor, mu: LabelTensor) -> torch.Tensor:
#     x1 = x.get_components()
#     t1 = t.get_components()
#     return f_ini(x, mu) + inputs * t1 * x1 * (xmax - x1)
#
#
# def functional_post_processing(func, t: torch.Tensor, x: torch.Tensor, mu: torch.Tensor, theta: torch.Tensor) -> torch.Tensor:
#     tz = torch.zeros_like(x)
#     ini = functional_exact(tz, x, mu)
#     return ini + func(t, x, mu, theta) * t[0] * x[0] * (xmax - x[0])

bc_weight = 50.0
ic_weight = 500.0

domain_x = Segment1D((0.0, xmax), is_main_domain=True)
domain_mu = [(0.8, 1.2)]

t_min, t_max = 0.0, 0.4
domain_t = (t_min, t_max)

sampler = TensorizedSampler(
    [
        UniformTimeSampler(domain_t),
        DomainSampler(domain_x),
        UniformParametricSampler(domain_mu),
    ]
)

space = NNxtSpace(1, 1, GenericMLP, domain_x, sampler, layer_sizes=[20, 40, 20])
pde = Transport1D(space, init=f_ini, f=f_rhs, g=f_bc, a=a)

opt_1 = {
    "name": "adam",
    "optimizer_args": {"lr": 1.8e-2, "betas": (0.9, 0.999)},
}
pinn = TemporalPinns(
    pde,
    bc_type="weak",
    ic_type="weak",
    bc_weight=bc_weight,
    ic_weight=ic_weight,
    optimizers=OptimizerData(opt_1),
)

# %%

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

pinn.space.load_from_best_approx()

# plot_AbstractApproxSpaces(
#     pinn.space,
#     domain_x,
#     domain_mu,
#     domain_t,
#     time_values=[0.0, 0.2, t_max],
#     loss=pinn.losses,
#     residual=pde,
#     solution=exact,
#     error=exact,
#     title="learning sol of 1D transport equation with TemporalPinns, weak initial and boundaries conditions",
# )
#
# plt.show()

space2 = NNxtSpace(1, 1, GenericMLP, domain_x, sampler, layer_sizes=[64])
pde2 = Transport1D(space2, init=f_ini, f=f_rhs, g=f_bc, a=a, functional_a=functional_a)
pinn2 = NaturalGradientTemporalPinns(
    pde2,
    bc_type="weak",
    ic_type="weak",
    bc_weight=bc_weight,
    ic_weight=ic_weight,
    matrix_regularization=1e-6,
)

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

pinn2.space.load_from_best_approx()

space3 = NNxtSpace(1, 1, GenericMLP, domain_x, sampler, layer_sizes=[64])
pde3 = Transport1D(space3, init=f_ini, f=f_rhs, g=f_bc, a=a, functional_a=functional_a)
pinn3 = AnagramTemporalPinns(
    pde3,
    bc_type="weak",
    ic_type="weak",
    bc_weight=bc_weight,
    ic_weight=ic_weight,
    svd_threshold=5e-3,
)

resume_solve = False
if resume_solve or not pinn3.load(__file__, "anagram"):
    pinn3.solve(
        epochs=200,
        n_collocation=2000,
        verbose=True,
    )
    pinn3.save(__file__, "anagram")

pinn3.space.load_from_best_approx()

plot_abstract_approx_spaces(
    (pinn.space, pinn2.space, pinn3.space),
    domain_x,
    domain_mu,
    domain_t,
    loss=(pinn.losses, pinn2.losses, pinn3.losses),
    residual=(pde, pde2, pde3),
    solution=exact,
    error=exact,
    title="learning sol of 1D transport equation with TemporalPinns, weak initial and boundaries conditions",
    titles=("no preconditioner", "ENG preconditioner", "Anagram preconditioner"),
)

plt.show()