r"""Solves the viscous Burgers advection equation in 1D using a PINN.

.. math::

    \partial_t u + \partial_x \frac {u^2}{2} - \sigma \partial_{xx} u & = f in \Omega \times (0, T) \\
    u & = g on \partial \Omega \times (0, T) \\
    u & = u_0 on \Omega \times {0}

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

The equation is solved on a segment domain; strong boundary and initial conditions are used.
Three training strategies are compared: standard PINNs with SS-BFGS optimizer, 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 (
    NaturalGradientTemporalPinns,
    TemporalPinns,
)
from scimba_torch.physical_models.elliptic_pde.laplacians import (
    Laplacian1DDirichletStrongForm,
)
from scimba_torch.physical_models.temporal_pde.abstract_temporal_pde import (
    GenericFirstOrderTemporalPDE,
)
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces
from scimba_torch.utils.scimba_tensors import LabelTensor

torch.manual_seed(0)

PI = torch.pi
DEFAULT_σ = 2e-3


class ViscousBurgers1DSpaceComponent(Laplacian1DDirichletStrongForm):
    """Implementation of a 2D steady reaction-diffusion problem with Neumann BCs in strong form.

    Args:
        space: The approximation space for the problem.
        f: Callable, represents the source term f(x, μ) (default: f=1).
        g: Callable, represents the Neumann boundary condition g(x, μ) (default: g=0).
        **kwargs: Additional keyword arguments.
    """

    def __init__(self, space, **kwargs):
        super().__init__(space, **kwargs)
        self.σ = kwargs.get("sigma", DEFAULT_σ)

    def operator(self, w, x, mu):
        u = w.get_components()
        u2_x = self.space.grad(u**2 / 2, x)
        u_xx = self.space.grad(self.space.grad(u, x), x)
        return u2_x - self.σ * u_xx

    def functional_operator(self, func, x, mu, theta):
        # space operator
        def u2(*args):
            return func(*args) ** 2 / 2

        dx_u2 = torch.func.jacrev(u2, 0)(x, mu, theta).squeeze()

        dx_u = torch.func.jacrev(func, 0)
        d2x_u = torch.func.jacrev(dx_u, 0)(x, mu, theta).squeeze()

        space_op = dx_u2 - self.σ * d2x_u
        return space_op.unsqueeze(0)


class ViscousBurgers1D(GenericFirstOrderTemporalPDE):
    r"""Implementation of a 1D viscous Burgers' equation with Dirichlet boundary conditions."""

    def __init__(self, space, init, f, **kwargs):
        space_component = ViscousBurgers1DSpaceComponent(space)
        super().__init__(space, space_component, init, f, **kwargs)


def f_rhs(w, t, x, mu, σ):
    x_ = x.get_components()
    t_ = t.get_components()

    exp_neg_t = torch.exp(-t_)
    sin_term = torch.sin(2 * PI * x_)
    cos_term = torch.cos(2 * PI * x_)

    return exp_neg_t * sin_term * (2 * PI * (cos_term * exp_neg_t + 2 * PI * σ) - 1)


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


def functional_exact(t, x, mu):
    return torch.sin(2 * PI * x) * torch.exp(-t)


def exact_solution(t, x, mu):
    return functional_exact(t.get_components(), x.get_components(), mu.get_components())


def initial_condition(x, mu):
    t = LabelTensor(torch.zeros_like(x.x))
    return exact_solution(t, x, mu)


def post_processing(inputs, t, x, mu):
    x1 = x.get_components()
    t1 = t.get_components()
    return initial_condition(x, mu) + inputs * t1 * x1 * (1 - x1)


def functional_post_processing(func, t, x, mu, theta):
    tz = torch.zeros_like(x)
    ini = functional_exact(tz, x, mu)
    return ini + func(t, x, mu, theta) * t[0] * x[0] * (1 - x[0])


domain_x = Segment1D((0, 1), is_main_domain=True)
domain_mu = []

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,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[20, 40, 20],
    post_processing=post_processing,
)
pde = ViscousBurgers1D(
    space, init=initial_condition, f=lambda *args: f_rhs(*args, DEFAULT_σ)
)

pinn = TemporalPinns(pde, bc_type="strong", ic_type="strong", optimizers="ssbfgs")

new_solve = True
if new_solve or not pinn.load(__file__, "ssbfgs"):
    pinn.solve(epochs=1000, n_collocation=3000)
    pinn.save(__file__, "ssbfgs")

plot_abstract_approx_spaces(
    pinn.space,
    domain_x,
    domain_mu,
    domain_t,
    time_values=[0.0, 0.2, t_max],
    loss=pinn.losses,
    residual=pde,
    solution=exact_solution,
    error=exact_solution,
    title="learning sol of 1D viscous Burgers' equation with TemporalPinns, strong initial and boundaries conditions",
)

plt.show()

# %%

space2 = NNxtSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32],
    post_processing=post_processing,
)
pde2 = ViscousBurgers1D(
    space2, init=initial_condition, f=lambda *args: f_rhs(*args, DEFAULT_σ)
)
pinn2 = NaturalGradientTemporalPinns(
    pde2,
    bc_type="strong",
    ic_type="strong",
    matrix_regularization=1e-6,
    functional_post_processing=functional_post_processing,
    adaptive_matrix_regularization=True,
    adaptive_matrix_regularization_decrease=0.25,
)

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

plot_abstract_approx_spaces(
    pinn2.space,
    domain_x,
    domain_mu,
    domain_t,
    time_values=[0.0, 0.2, t_max],
    loss=pinn2.losses,
    residual=pde2,
    solution=exact_solution,
    error=exact_solution,
    title="learning sol of 1D viscous Burgers' equation with TemporalPinns, strong initial and boundaries conditions",
)

plt.show()

# %%

space3 = NNxtSpace(
    1,
    0,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[32],
    post_processing=post_processing,
)
pde3 = ViscousBurgers1D(
    space3, init=initial_condition, f=lambda *args: f_rhs(*args, DEFAULT_σ)
)
pinn3 = NaturalGradientTemporalPinns(
    pde3,
    bc_type="strong",
    ic_type="strong",
    ng_algo="ANaGRAM",
    svd_threshold=5e-4,
    functional_post_processing=functional_post_processing,
)

resume_solve = True
if resume_solve or not pinn3.load(__file__, "ANaGRAM"):
    pinn3.solve(epochs=100, n_collocation=2000)
    pinn3.save(__file__, "ANaGRAM")

plot_abstract_approx_spaces(
    pinn3.space,
    domain_x,
    domain_mu,
    domain_t,
    time_values=[0.0, 0.2, t_max],
    loss=pinn3.losses,
    residual=pde3,
    solution=exact_solution,
    error=exact_solution,
    title="learning sol of 1D viscous Burgers' equation with TemporalPinns, strong initial and boundaries conditions",
)

plt.show()

# %%

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_solution,
    error=exact_solution,
    title="learning sol of 1D viscous Burgers' equation with TemporalPinns, strong initial and boundaries conditions",
    titles=("SS-BFGS", "ENG preconditioner", "ANaGRAM preconditioner"),
)

plt.show()

# %%