r"""Solves the Vlasov equation in 1Dx1V using the neural semi-Lagrangian method.

.. math::

    \partial_t f + v \partial_x f + \sin(x) \partial_x f & = 0 in \Omega \times (0, T) \\
    f & = f_0 on \Omega \times {0}

where :math:`f: \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. Periodic boundary conditions are prescribed.

The equation is solved on a square domain; the periodic boundary conditions are applied using a periodic embedding directly at the neural network level.
Two training strategies are usable: Adam and natural gradient preconditioning.
"""

# %%

import torch

from scimba_torch import PI
from scimba_torch.approximation_space.nn_space import NNxvSpace
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,
    UniformVelocitySamplerOnCuboid,
)
from scimba_torch.neural_nets.coordinates_based_nets.features import PeriodicMLP
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.numerical_solvers.collocation_projector import (
    CollocationProjector,
    NaturalGradientProjector,
)
from scimba_torch.numerical_solvers.temporal_pde.neural_semilagrangian import (
    Characteristic,
    NeuralSemiLagrangian,
)
from scimba_torch.optimizers.optimizers_data import OptimizerData
from scimba_torch.physical_models.kinetic_pde.vlasov import Vlasov

SQRT_2_PI = (2 * PI) ** 0.5

DOMAIN_X = Segment1D((0, 2 * PI), is_main_domain=True)
DOMAIN_V = Segment1D((-6, 6))
DOMAIN_MU = []

SAMPLER = TensorizedSampler(
    [
        DomainSampler(DOMAIN_X),
        UniformVelocitySamplerOnCuboid(DOMAIN_V),
        UniformParametricSampler(DOMAIN_MU),
    ]
)


def initial_condition(x, v, mu):
    v_ = v.get_components()
    return torch.exp(-(v_**2) / 2) / SQRT_2_PI


def electric_field(t, x, mu):
    x_ = x.get_components()
    return torch.sin(x_)


def opt():
    opt_1 = {
        "name": "adam",
        "optimizer_args": {"lr": 2.5e-2, "betas": (0.9, 0.999)},
    }
    return OptimizerData(opt_1)


# %%


def solve_with_neural_sl(
    T: float,
    dt: float,
    with_natural_gradient: bool = True,
    with_classical_projector: bool = False,
    N_c: int = 10_000,
    periodic_mlp: bool = True,
):
    torch.random.manual_seed(0)

    mlp = PeriodicMLP if periodic_mlp else GenericMLP

    if with_classical_projector:
        space = NNxvSpace(
            1,
            0,
            mlp,
            DOMAIN_X,
            DOMAIN_V,
            SAMPLER,
            layer_sizes=[60, 60, 60],
            activation_type="tanh",
        )
        pde = Vlasov(
            space,
            initial_condition=initial_condition,
            electric_field=electric_field,
        )
        characteristic = Characteristic(pde, periodic=True)

        projector = CollocationProjector(space, rhs=initial_condition, optimizers=opt())

        scheme = NeuralSemiLagrangian(characteristic, projector)

        scheme.initialization(epochs=750, n_collocation=N_c, verbose=True)
        scheme.projector.save("ini_transport_2D_SL")

        scheme.projector = CollocationProjector(
            space, rhs=initial_condition, optimizers=opt()
        )
        scheme.projector.load("ini_transport_2D_SL")
        scheme.projector.space.load_from_best_approx()
        scheme.solve(dt=dt, final_time=T, epochs=750, n_collocation=N_c, verbose=True)

    if with_natural_gradient:
        space = NNxvSpace(
            1,
            0,
            mlp,
            DOMAIN_X,
            DOMAIN_V,
            SAMPLER,
            layer_sizes=[20, 20, 20],
            activation_type="tanh",
        )
        pde = Vlasov(
            space,
            initial_condition=initial_condition,
            electric_field=electric_field,
        )
        characteristic = Characteristic(pde, periodic=True)

        projector = NaturalGradientProjector(space, rhs=initial_condition)

        scheme = NeuralSemiLagrangian(characteristic, projector)

        scheme.initialization(epochs=100, n_collocation=N_c, verbose=True)

        scheme.solve(dt=dt, final_time=T, epochs=100, n_collocation=N_c, verbose=True)

    return scheme


# %%

if __name__ == "__main__":
    scheme = solve_with_neural_sl(
        3.0,
        1.5,
        with_natural_gradient=True,
        with_classical_projector=False,
        N_c=64**2,
        periodic_mlp=False,
    )

# %%

if __name__ == "__main__":
    import matplotlib.pyplot as plt

    from scimba_torch.plots.plots_nd import plot_abstract_approx_space

    cuts = [
        ([0.0, 0.0], [-0.5, 0.5]),
        ([0.0, 0.2], [0.0, 1.0]),
    ]

    plot_abstract_approx_space(
        scheme.pde.space,
        scheme.pde.space.spatial_domain,
        velocity_domain=scheme.pde.space.velocity_domain,
        loss=scheme.projector.losses,
        aspect="auto",
        derivatives=["ux", "uv"],
        cuts=cuts,
    )

    plt.show()

# %%