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

.. math::

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

The equation is solved on a segment domain; the optimizer uses natural gradient preconditioning.
"""

# %%

import copy

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.collocation_projector import (
    NaturalGradientProjector,
)
from scimba_torch.numerical_solvers.temporal_pde.discrete_pinn import DiscretePINN
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 torch.ones_like(x1)


def ff_bc(x, n, mu):
    x1 = x.get_components()
    return torch.ones_like(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()
    # J = 1000
    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=[32])
pde = HeatEquation1DStrongForm(space, init=f_ini, f=f_rhs, g=f_bc)
projector = NaturalGradientProjector(space, rhs=f_ini, bc_rhs=ff_bc)
scheme = DiscretePINN(pde, projector, scheme="rk2", initial_time=1e-3)

resume_init = False
if resume_init or not scheme.projector.load(__file__, "init"):
    scheme.initialization(
        epochs=100, verbose=True, n_collocation=900, n_bc_collocation=120
    )
    scheme.projector.save(__file__, "init")

scheme.projector.space.load_from_best_approx()
scheme.initial_space = copy.deepcopy(scheme.projector.space)

T = 3e-3

resume_solve = True
if resume_solve or not scheme.projector.load(__file__, "solved"):
    scheme.solve(
        dt=1e-3,
        final_time=T,
        epochs=100,
        n_collocation=900,
        n_bc_collocation=120,
        verbose=True,
        nb_keep=1,
    )
    scheme.projector.save(__file__, "solved")

scheme.projector.space.load_from_best_approx()

plot_time_discrete_scheme(
    scheme, solution=exact, error=exact, derivatives=["ux"], title="test"
)

plt.show()