r"""Solves the linearized Euler equations in 1D using a PINN.

.. math::

    \partial_t p + \partial_x u & = f_1 in \Omega \times (0, T) \\
    \partial_t u + \partial_x p & = f_2 in \Omega \times (0, T) \\
    p & = g_1 on \partial \Omega \times (0, T) \\
    u & = g_2 on \partial \Omega \times (0, T) \\
    p & = p_0 on \Omega \times {0} \\
    u & = u_0 on \Omega \times {0}

where :math:`p: \partial \Omega \times (0, T) \to \mathbb{R}` and :math:`u: \partial \Omega \times (0, T) \to \mathbb{R}` are the unknown functions, :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; 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.
"""

from typing import 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 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 (
    Laplacian1DDirichletStrongFormNoParam,
)
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, MultiLabelTensor
from scimba_torch.utils.typing_protocols import VarArgCallable


def f_rhs(
    w: MultiLabelTensor, t: LabelTensor, x: LabelTensor, mu: LabelTensor
) -> Tuple[torch.Tensor]:
    x = x.get_components()
    return torch.zeros_like(x), torch.zeros_like(x)


def f_bc(
    w: MultiLabelTensor, t: LabelTensor, x: LabelTensor, n: LabelTensor, mu: LabelTensor
) -> Tuple[torch.Tensor]:
    x = x.get_components()
    return torch.zeros_like(x), torch.zeros_like(x)


class LinearizedEulerSpaceComponent(Laplacian1DDirichletStrongFormNoParam):
    def __init__(self, space: AbstractApproxSpace, **kwargs):
        super().__init__(space, **kwargs)

    def operator(
        self, w: MultiLabelTensor, x: LabelTensor, mu: LabelTensor
    ) -> Tuple[torch.Tensor]:
        p, u = w.get_components()
        p_x = self.grad(p, x)
        u_x = self.grad(u, x)
        return u_x, p_x

    def functional_operator(
        self,
        func: VarArgCallable,
        x: torch.Tensor,
        mu: torch.Tensor,
        theta: torch.Tensor,
    ) -> torch.Tensor:
        pu_space = torch.func.jacrev(func, 0)(x, mu, theta)
        pu_space = torch.flip(pu_space, (-1,))
        return pu_space.squeeze()


class LinearizedEuler(GenericFirstOrderTemporalPDE):
    r"""Implementation of a Linearized Euler equation with Dirichlet boundary conditions."""

    def __init__(self, space, init, **kwargs):
        space_component = LinearizedEulerSpaceComponent(space)

        self.residual_size = 2
        self.bc_residual_size = 2
        self.ic_residual_size = 2

        super().__init__(space, space_component, init, f=f_rhs, g=f_bc, **kwargs)


def exact_solution(t: LabelTensor, x: LabelTensor, mu: LabelTensor) -> torch.Tensor:
    x = x.get_components()
    D = 0.02
    coeff = 1 / (4 * torch.pi * D) ** 0.5
    p_plus_u = coeff * torch.exp(-((x - t.x - 1) ** 2) / (4 * D))
    p_minus_u = coeff * torch.exp(-((x + t.x - 1) ** 2) / (4 * D))
    p = (p_plus_u + p_minus_u) / 2
    u = (p_plus_u - p_minus_u) / 2
    return torch.cat((p, u), dim=-1)


def initial_solution(x: LabelTensor, mu: LabelTensor) -> Tuple[torch.Tensor]:
    sol = exact_solution(LabelTensor(torch.zeros_like(x.x)), x, mu)
    return sol[..., 0:1], sol[..., 1:2]


domain_mu = []
domain_x = Segment1D((-1.0, 3.0), is_main_domain=True)

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

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

bc_weight = 30.0
ic_weight = 300.0

space = NNxtSpace(2, 0, GenericMLP, domain_x, sampler, layer_sizes=[64])
pde = LinearizedEuler(space, initial_solution)

pinns = TemporalPinns(
    pde,
    bc_type="weak",
    ic_type="weak",
    optimizers="ssbfgs",
    bc_weight=bc_weight,
    ic_weight=ic_weight,
    one_loss_by_equation=True,
)

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=1000,
        n_collocation=3000,
        n_bc_collocation=1000,
        n_ic_collocation=1000,
    )
    pinns.save(__file__, save_post_fix)
print(f"Training time: {pinns.training_time:.2f} seconds")
print(f"Best Loss    : {pinns.best_loss:.2e}")

# plot_AbstractApproxSpaces(
#     pinns.space,
#     domain_x,
#     domain_mu,
#     domain_t,
#     time_values=[t_max,],
#     loss=pinns.losses,
#     residual=pde,
#     solution=exact_solution,
#     error=exact_solution,
#     derivatives=["ux", "ut"],
#     title="solving LinearizedEuler with TemporalPinns",
# )
#
# plt.show()

space2 = NNxtSpace(2, 0, GenericMLP, domain_x, sampler, layer_sizes=[64])
pde2 = LinearizedEuler(space2, initial_solution)

pinns2 = NaturalGradientTemporalPinns(
    pde2,
    bc_type="weak",
    ic_type="weak",
    bc_weight=bc_weight,
    ic_weight=ic_weight,
    one_loss_by_equation=True,
    matrix_regularization=1e-6,
)

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_ENG"
# the post_fix of the file where to load the pinn
save_post_fix = "pinns_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=1000,
        n_collocation=3000,
        n_bc_collocation=1000,
        n_ic_collocation=1000,
    )
    pinns2.save(__file__, save_post_fix)
print(f"Training time: {pinns2.training_time:.2f} seconds")
print(f"Best Loss    : {pinns2.best_loss:.2e}")


def norm(space, t, x, mu):
    x.x.requires_grad_()
    w = space.evaluate(t, x, mu)
    u, p = w.get_components()
    return torch.sqrt(u**2 + p**2)


additional_scalar_functions = {"$\\|(u,p)\\|_2$": norm}

plot_abstract_approx_spaces(
    (
        pinns.space,
        pinns2.space,
    ),
    domain_x,
    domain_mu,
    domain_t,
    time_values=[
        t_max,
    ],
    # loss=(
    #     pinns.losses,
    #     pinns2.losses,
    # ),
    residual=(
        pde,
        pde2,
    ),
    solution=exact_solution,
    error=exact_solution,
    derivatives=["ux", "ut"],
    additional_scalar_functions=additional_scalar_functions,
    title="solving LinearizedEuler with TemporalPinns",
    titles=("SS-BFGS", "ENG preconditioning"),
)

plt.show()

space3 = NNxtSpace(2, 0, GenericMLP, domain_x, sampler, layer_sizes=[32, 32])
pde3 = LinearizedEuler(space3, initial_solution)

pinns3 = NaturalGradientTemporalPinns(
    pde3,
    bc_type="weak",
    ic_type="weak",
    bc_weight=bc_weight,
    ic_weight=ic_weight,
    one_loss_by_equation=True,
    ng_algo="ANaGRAM",
    svd_threshold=5e-2,
)

solving_mode = "load"
# 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_ANaGRAM"
# the post_fix of the file where to load the pinn
save_post_fix = "pinns_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=100,
        n_collocation=3000,
        n_bc_collocation=1000,
        n_ic_collocation=1000,
    )
    pinns3.save(__file__, save_post_fix)

plot_abstract_approx_spaces(
    (
        pinns2.space,
        pinns3.space,
    ),
    domain_x,
    domain_mu,
    domain_t,
    time_values=[
        t_max,
    ],
    loss=(
        pinns2.losses,
        pinns3.losses,
    ),
    residual=(
        pde2,
        pde3,
    ),
    solution=exact_solution,
    error=exact_solution,
    derivatives=["ux", "ut"],
    title="solving LinearizedEuler with TemporalPinns",
    titles=("ENG preconditioning", "ANaGRAM preconditioning"),
)

plt.show()