r"""Solves the heat equation in 1D using a PINN.

.. math::

    \partial_t u - \partial_{xx} u & = f in \Omega \times (0, T) \\
    \partial_x 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. Neumann boundary conditions are prescribed, and the initial condition is discontinuous.

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.heat_equation import (
    HeatEquation1DStrongForm,
)
from scimba_torch.plots.plots_nd import plot_abstract_approx_spaces
from scimba_torch.utils.scimba_tensors import LabelTensor

# %%


def f_rhs(w, t, x, mu):
    mu1 = mu.get_components()
    return 0 * mu1


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


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


def exact(t, x, mu, J=1000):
    x1 = x.get_components()
    t1 = t.get_components() + 1e-3

    u_J = 5 / 4
    jpi = torch.pi * torch.arange(1, J + 1)[:, None, None]
    c_j = 2 * (torch.sin(5 * jpi / 8) - torch.sin(3 * jpi / 8)) / jpi

    u_J += torch.sum(
        c_j * torch.exp(-(jpi**2) * t1[None, ...]) * torch.cos(jpi * x1[None, ...]),
        axis=0,
    )

    return u_J


domain_x = Segment1D((0, 1), is_main_domain=True)

t_min, t_max = 0.0, 1.0

sampler = TensorizedSampler(
    [
        UniformTimeSampler((t_min, t_max)),
        DomainSampler(domain_x),
        UniformParametricSampler([(1.0, 1.0 + 1e-5)]),
    ]
)

# %%

bc_weight = 50.0
ic_weight = 500.0

# weak BC #
space = NNxtSpace(1, 1, GenericMLP, domain_x, sampler, layer_sizes=[20, 40, 20])
pde = HeatEquation1DStrongForm(space, init=f_ini, f=f_rhs, g=f_bc)

# %%

# losses = GenericLosses(
#     [
#         ("residual", torch.nn.MSELoss(), 1.0),
#         ("bc", torch.nn.MSELoss(), bc_weight),
#         ("ic", torch.nn.MSELoss(), ic_weight),
#     ],
# )
opt_1 = {
    "name": "adam",
    "optimizer_args": {"lr": 1.8e-2, "betas": (0.9, 0.999)},
}
# opt = OptimizerData(opt_1)

# %%

pinn = TemporalPinns(
    pde,
    bc_type="weak",
    ic_type="weak",
    optimizers=OptimizerData(opt_1),
    bc_weight=bc_weight,
    ic_weight=ic_weight,
)

# %%

resume_solve = False
if resume_solve or not pinn.load(__file__, "neumann_weak_no_precond"):
    pinn.solve(
        epochs=1000,
        n_collocation=2000,
        n_bc_collocation=50,
        n_ic_collocation=1500,
        verbose=True,
    )
    pinn.save(__file__, "neumann_weak_no_precond")

pinn.space.load_from_best_approx()

# plot_AbstractApproxSpaces(
#     pinn.space,
#     domain_x,
#     [(1.0, 1.0 + 1e-5)],
#     (t_min, t_max),
#     parameters_values=[t_max],
#     time_values=[0.0, 1e-3, 0.1],
#     loss=pinn.losses,
#     residual=pde,
#     solution=exact,
#     error=exact,
#     derivatives=["ux", "ut", "utx"],
#     title="learning sol of 1D heat equation with TemporalPinns",
# )
#
# plt.show()

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

resume_solve = False
if resume_solve or not pinn2.load(__file__, "neumann_weak_ENG"):
    # pinn2.space.load_from_dict(pinn.space.dict_for_save())
    # pinn2.space.load_from_best_approx()
    pinn2.solve(
        epochs=20,
        n_collocation=1200,
        n_bc_collocation=60,
        n_ic_collocation=1200,
        verbose=True,
    )
    pinn2.save(__file__, "neumann_weak_ENG")

pinn2.space.load_from_best_approx()

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

resume_solve = False
if resume_solve or not pinn3.load(__file__, "neumann_weak_anagram"):
    # pinn3.space.load_from_dict(pinn2.space.dict_for_save())
    # pinn3.space.load_from_best_approx()
    pinn3.solve(
        epochs=20,
        n_collocation=1200,
        n_bc_collocation=60,
        n_ic_collocation=1200,
        verbose=True,
    )
    pinn3.save(__file__, "neumann_weak_anagram")

pinn3.space.load_from_best_approx()

plot_abstract_approx_spaces(
    (pinn.space, pinn2.space, pinn3.space),
    domain_x,
    [(1.0, 1.0 + 1e-5)],
    (t_min, t_max),
    parameters_values=([1.0], [1.0], [1.0]),
    time_values=[t_max],
    loss=(pinn.losses, pinn2.losses, pinn3.losses),
    residual=(pde, pde2, pde3),
    solution=exact,
    error=exact,
    derivatives=["ux"],
    # title="learning sol of 1D heat equation with TemporalPinns",
)

plt.show()