r"""Solves the advection equation in 1D using Neural Semi-Lagrangian (NSL).

.. math::

    \partial_t u + c \partial_{x} u & = 0 \quad \text{in } \Omega \times (0, T) \\
    u & = u_0 \quad \text{on } \Omega \times \{0\} \\
    u(0, t) & = u(1, t) \quad \text{periodic BC}

where :math:`u: \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.

The equation is solved on a segment domain with periodic boundary conditions
using the Neural Semi-Lagrangian method. The advection velocity is constant.

This example demonstrates:
- Using periodic boundary conditions via modulo operation on characteristics
- Solving the advection equation with a simple constant velocity field
- Comparing the NSL solution to the exact solution
"""

import time

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt

from scimba_jax.domains.meshless_domains.domains_1d import Segment1D
from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import (
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
    DomainSampler,
    TensorizedSampler,
)
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.numerical_solvers.neural_semi_lagrangian import (
    NeuralSemiLagrangian,
)
from scimba_jax.plots.plots_nd import plot_abstract_approx_spaces

N_COLLOC = 1000
N_EPOCHS_INIT = 100
N_EPOCHS = 10

DOMAIN_BOUNDS = ((0.0,), (1.0,))
X_MIN, X_MAX = DOMAIN_BOUNDS[0][0], DOMAIN_BOUNDS[1][0]
DOM_X = Segment1D((X_MIN, X_MAX), is_main_domain=True)
DOM_T = (0.0, 0.5)
SAMPLER = TensorizedSampler([DomainSampler(DOM_X)], model_type="x")

VELOCITY = 0.5  # Constant advection velocity
NT = 10  # Number of time steps
DT = (DOM_T[1] - DOM_T[0]) / NT  # Time step size


def exact_sol(t: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
    """Exact solution of the advection equation with periodic BC.

    The exact solution is a Gaussian profile advected at velocity c:
    u(x, t) = exp(-150 * (x - (0.25 + c*t))^2)

    With periodic boundary conditions, the profile wraps around when it
    reaches the domain boundaries.
    """
    # Apply periodic boundary conditions using modulo operation
    x_ = (x - X_MIN) % (X_MAX - X_MIN) + X_MIN

    # Compute the advected position
    x_advected = x_ - (0.25 + VELOCITY * t)

    # Compute the exact solution as a Gaussian profile
    return jnp.exp(-((x_advected) ** 2) * 150)


def f_init(x: jnp.ndarray) -> jnp.ndarray:
    """Initial condition function."""
    t = jnp.zeros_like(x)
    return exact_sol(t, x)


def advection_field(t: jnp.ndarray, x: jnp.ndarray) -> jnp.ndarray:
    """Advection field a(t, x, mu) = c (constant velocity)."""
    return jnp.ones_like(x) * VELOCITY


def exact_characteristic_foot(t: jnp.ndarray, x: jnp.ndarray, dt: float) -> jnp.ndarray:
    """Exact characteristic foot computation.

    For the advection equation u_t + c u_x = 0, the characteristics are:
    X(s) = x - c * (s - t)

    We want the foot at time t given x at time t+dt:
    X(t) = x - c * dt
    """
    a = VELOCITY
    return x - a * dt


if __name__ == "__main__":
    # Initialize the key and neural network
    key = jax.random.PRNGKey(0)
    in_size = 1
    out_size = 1
    nn = MLP(in_size=in_size, out_size=out_size, hidden_sizes=[12] * 2, key=key)
    space = ApproximationSpace({"x": 1}, [(nn, "scalar", None)], model_type="x")

    neural_sl = NeuralSemiLagrangian(
        main_domain=DOM_X,
        time_domain=DOM_T,
        sampler=SAMPLER,
        dt=DT,
        advection_field=advection_field,
        exact_characteristic_foot=exact_characteristic_foot,
        periodic=True,
        domain_bounds=DOMAIN_BOUNDS,
        out_size=out_size,
        model_type="x",
        exact_solution=exact_sol,
    )

    print("Initializing the Neural Semi-Lagrangian solver...")
    start_init = time.perf_counter()
    key, neural_sl = neural_sl.initialize(key, space, f_init, N_EPOCHS_INIT, N_COLLOC)
    space = neural_sl.space
    end_init = time.perf_counter()
    print(
        f"Initializing the Neural Semi-Lagrangian solver... "
        f"Done in {end_init - start_init:.2f} seconds\n"
    )

    # Plot initial condition
    plot_abstract_approx_spaces(
        [space],
        DOM_X,
        solution=f_init,
        error=f_init,
        title="initial condition",
    )

    print("Solving with Neural Semi-Lagrangian method...")
    start_solve = time.perf_counter()
    key, neural_sl = neural_sl.solve(key, space, N_EPOCHS, N_COLLOC)
    space = neural_sl.space
    end_solve = time.perf_counter()
    print(
        f"Solving with Neural Semi-Lagrangian method... "
        f"Done in {end_solve - start_solve:.2f} seconds\n"
    )

    # Plot final solution
    plot_abstract_approx_spaces(
        [space],
        DOM_X,
        solution=lambda x: exact_sol(jnp.ones_like(x) * DOM_T[-1], x),
        error=lambda x: exact_sol(jnp.ones_like(x) * DOM_T[-1], x),
        title=f"solution at final time t={DOM_T[-1]}",
    )

    # Print final errors
    if hasattr(neural_sl, "errors_over_time"):
        errors = neural_sl.errors_over_time
        print(f"\nFinal relative L2 error: {errors[-1, 0]:.2e}")
        print(f"Final relative Linf error: {errors[-1, 1]:.2e}")

    plt.show()