r"""Solves a 1D heat equation using a discrete PINN.

..math::
    \partial_t u - \partial_{xx} u = 0

The initial condition is a bump function.

The equation is solved using a discrete PINN and an Adam optimizer followed by L-BFGS.
"""

# %%

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.temporal_pde.discrete_pinn import DiscretePINN
from scimba_torch.numerical_solvers.temporal_pde.time_discrete import (
    TimeDiscreteCollocationProjector,
)
from scimba_torch.optimizers.optimizers_data import OptimizerData
from scimba_torch.physical_models.temporal_pde.heat_equation import (
    HeatEquation1DStrongForm,
)
from scimba_torch.plots.plot_time_discrete_scheme import plot_time_discrete_scheme
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):
    # start from a regularized initial condition
    t = LabelTensor(1e-3 * torch.ones_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()

    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


def opt():
    opt_1 = {
        "name": "adam",
        "optimizer_args": {"lr": 2.5e-2, "betas": (0.9, 0.999)},
    }
    opt_2 = {
        "name": "lbfgs",
        "switch_at_epoch_ratio": 0.9,
        "optimizer_args": {
            "history_size": 50,
            "max_iter": 50,
            "tolerance_grad": 1e-11,
            "tolerance_change": 1e-9,
        },
    }
    return OptimizerData(opt_1, opt_2)


# %%

torch.random.manual_seed(0)

domain_x = Segment1D((0, 1), is_main_domain=True)
domain_mu = [[1.0, 1.0 + 1e-5]]
sampler = TensorizedSampler(
    [
        DomainSampler(domain_x),
        UniformParametricSampler(domain_mu),
    ]
)

space = NNxSpace(1, 1, GenericMLP, domain_x, sampler, layer_sizes=[20] * 4)
pde = HeatEquation1DStrongForm(space, init=f_ini, f=f_rhs, g=f_bc)
projector = TimeDiscreteCollocationProjector(pde, rhs=f_ini)
scheme = DiscretePINN(pde, projector, scheme="rk2", initial_time=1e-3)

scheme.initialization(epochs=500, verbose=True, n_collocation=3000)
scheme.projector.save("ini_heat")

T = 5e-3
# scheme.projector = CollocationProjector(space, rhs=f_ini)
# scheme.projector.load("ini_heat")
scheme.projector.space.load_from_best_approx()
# scheme.solve(dt=1e-3, T=T, verbose=True)
scheme.solve(dt=1e-3, final_time=T, epochs=250, n_collocation=2000, verbose=True)


plot_time_discrete_scheme(scheme, solution=exact, error=exact)

plt.show()