r"""Solves a 2D heat equation using neural Galerkin.

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

The equation is solved using a neural Galerkin scheme 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_2d import Square2D
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.neural_galerkin import NeuralGalerkin
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 (
    HeatEquation2DStrongForm,
)
from scimba_torch.plots.plot_time_discrete_scheme import plot_time_discrete_scheme
from scimba_torch.utils.scimba_tensors import LabelTensor


def f_ini(x: LabelTensor, mu):
    x1, x2 = x.get_components()
    f = torch.sin(torch.pi * x1) * torch.sin(torch.pi * x2)
    return f


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


def f_bc(x: LabelTensor, mu: LabelTensor):
    x1, _ = x.get_components()
    return x1 * 0.0


def f_exact(t: LabelTensor, x: LabelTensor, mu: LabelTensor):
    x1, x2 = x.get_components()
    D = torch.tensor(0.02)
    f = (
        torch.exp(-2 * D * torch.pi**2 * t.x)
        * torch.sin(torch.pi * x1)
        * torch.sin(torch.pi * x2)
    )
    return f


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)


domain_x = Square2D([(0.0, 1), (0.0, 1)], is_main_domain=True)
domain_mu = [(0.02, 0.02 + 1e-5)]
sampler = TensorizedSampler(
    [DomainSampler(domain_x), UniformParametricSampler(domain_mu)]
)

space = NNxSpace(
    1,
    1,
    GenericMLP,
    domain_x,
    sampler,
    layer_sizes=[20] * 2,
)

pde = HeatEquation2DStrongForm(space, init=f_ini, f=f_rhs, g=f_bc)
projector = TimeDiscreteCollocationProjector(pde, rhs=f_ini, optimizers=opt())
scheme = NeuralGalerkin(pde, projector, scheme="rk2")

scheme.initialization(epochs=2000, verbose=True, n_collocation=2000)
scheme.projector.save("ini_heat2D")

T = 3e-1
scheme.projector.space.load_from_best_approx()
scheme.solve(dt=1e-3, final_time=T, n_collocation=3000, epochs=500)

plot_time_discrete_scheme(
    scheme,
    solution=f_exact,
    error=f_exact,
    draw_contours=True,
)

plt.show()