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

.. math::::
    \partial_t u + a \partial_x u + b \partial_y u = 0

The initial condition is :math:`u(x, y, 0, \alpha) = 1`, and a parametric time-dependent Dirichlet boundary condition is applied: :math:`u(0, y, t, \alpha) = 1 - \alpha y t^2` on the left boundary and :math:`u(x, 0, t, \alpha) = 1 - \alpha x t^2` on the bottom boundary.

The equation is solved using the neural semi-Lagrangian scheme. The space domain is :math:`[0, 1]^2`, 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_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_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):
    x1, x2 = x.get_components()
    return exact(LabelTensor(torch.zeros_like(x1)), x, mu)


def exact(t: LabelTensor, x: LabelTensor, mu: LabelTensor) -> torch.Tensor:
    """Exact solution of the 2D transport equation:
        u_t + a u_x + b u_y = 0
    with BC:
        u(0,y,t) = 1 - y*t^2
        u(x,0,t) = 1 - x*t^2
    and IC:
        u(x,y,0) = 1.
    """
    x1, x2 = x.get_components()
    t_ = t.get_components()
    alpha = mu.get_components()

    a = 1
    b = 1

    s_x = t_ - x1
    s_y = t_ - x2

    u = torch.ones_like(x1)

    # Case 1: characteristics hit x=0 first
    mask_x = s_x >= torch.maximum(torch.zeros_like(s_x), s_y)
    u = torch.where(mask_x, 1 - alpha * (x2 - x1) ** 2 * (t_ - x1 / a) ** 2, u)

    # Case 2: characteristics hit y=0 first
    mask_y = s_y >= torch.maximum(torch.zeros_like(s_y), s_x)
    u = torch.where(mask_y, 1 - alpha * (x1 - x2) ** 2 * (t_ - x2 / b) ** 2, u)

    return u


# %%


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

    domain_x = Square2D(((0, 1), (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] * 3)

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

    def left_bc(x: LabelTensor, mu: LabelTensor, t: float):
        x1, x2 = x.get_components()
        mu = mu.get_components()
        return 1 - x2**2 * mu * t**2

    right_bc = None

    def bottom_bc(x: LabelTensor, mu: LabelTensor, t: float):
        x1, x2 = x.get_components()
        mu = mu.get_components()
        return 1 - x1**2 * mu * t**2

    top_bc = None

    dirichlet = [[left_bc, right_bc], [bottom_bc, top_bc]]

    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.75, 0.025)
    end = time.perf_counter()
    print(f"Solved in {end - start:.2f} seconds.")
    plot_time_discrete_scheme(
        scheme,
        solution=exact,
        error=exact,
        cuts=[([0.0, 0.25], [0.0, 1.0])],
    )
    plt.show()

# %%