r"""Solves a 1D advection equation with a parametric Dirichlet boundary condition.

.. math::::
    \partial_t u + c \partial_x u = 0

The initial condition is :math:`u(x, 0, \alpha) = 1`, and a parametric time-dependent Dirichlet boundary condition is applied: :math:`u(0, t, \alpha) = 1 - \alpha t^2`. The exact solution is :math:`u(x, t, \alpha) = 1` if :math:`x \geq c t` and :math:`u(x, t, \alpha) = 1 - \alpha (t - x/c)^2` if :math:`x < c t`.

The equation is solved using the neural semi-Lagrangian scheme. The space domain is :math:`[0, 1]`, the advection speed is :math:`c = 1`, and the parameter :math:`\alpha` varies in :math:`[1, 2]`.
"""

# %%

import time

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.neural_semilagrangian import (
    Characteristic,
    NeuralSemiLagrangian,
)
from scimba_torch.numerical_solvers.temporal_pde.time_discrete import (
    TimeDiscreteNaturalGradientProjector,
)
from scimba_torch.physical_models.temporal_pde.advection_diffusion_equation import (
    AdvectionReactionDiffusionDirichletStrongForm,
)
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, mu):
    return exact(LabelTensor(torch.zeros_like(x.x)), x, mu)


def exact(t: LabelTensor, x: LabelTensor, mu: LabelTensor) -> torch.Tensor:
    x_ = x.get_components()
    t_ = t.get_components()
    alpha = mu.get_components()

    c = 1

    u = torch.ones_like(x_)

    mask = x_ < c * t_
    u[mask] = 1 - alpha[mask] * (t_[mask] - x_[mask] / c) ** 2

    return u


# %%


def solve_with_neural_sl(T: float, dt: float):
    torch.random.manual_seed(0)

    domain_x = Segment1D((0, 1), is_main_domain=True)
    domain_mu = [[1, 2]]

    sampler = TensorizedSampler(
        [DomainSampler(domain_x), UniformParametricSampler(domain_mu)]
    )

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

    pde = AdvectionReactionDiffusionDirichletStrongForm(
        space,
        u0=f_ini,
        constant_advection=True,
        zero_diffusion=True,
        exact_solution=exact,
    )

    def dirichlet(
        w: torch.Tensor, x: LabelTensor, mu: LabelTensor, t: float, dt: float
    ) -> torch.Tensor:
        x = x.get_components()
        alpha = mu.get_components()
        mask = x < 0
        w[mask] = 1 - alpha[mask] * (t + dt / 2) ** 2
        return w

    characteristic = Characteristic(pde, dirichlet=dirichlet)

    projector = TimeDiscreteNaturalGradientProjector(pde, rhs=f_ini)
    scheme = NeuralSemiLagrangian(characteristic, projector)

    scheme.initialization(epochs=25, verbose=True, n_collocation=3000)
    scheme.solve(dt=dt, final_time=T, epochs=20, n_collocation=2000, verbose=True)

    return scheme


# %%

if __name__ == "__main__":
    start = time.perf_counter()
    scheme = solve_with_neural_sl(0.8, 0.025)
    end = time.perf_counter()
    print(f"Solved in {end - start:.2f} seconds.")
    plot_time_discrete_scheme(scheme, solution=exact, error=exact)
    plt.show()

# %%