r"""Solves a 2D advection-diffusion-reaction equation with time-discrete kernel collocation.

.. math::

    \partial_t u - \nabla \cdot (d \nabla u) + b \cdot \nabla u + c u & = f
    \text{ in } \Omega \times (0, T) \\
    u & = 0 \text{ on } \partial \Omega \times (0, T) \\
    u & = \sin(\pi x) \sin(\pi y) \text{ on } \Omega \times \{0\}

on :math:`\Omega = [-1, 1]^2`, with :math:`d = I` (isotropic diffusion),
:math:`b = (1, 0)` and :math:`c = 1`. The source :math:`f` is manufactured so
that the exact solution is
:math:`u(x, y, t) = e^{-\alpha t} \sin(\pi x) \sin(\pi y)`, which vanishes on
the whole boundary at every time.

The spatial operator is the library's ``GeneralEllipticResidual(A=d, b=b,
c=c)``, rebuilt at each stage time by ``spatial_residual_factory`` since the
source depends on time. The solution is expanded on a Gaussian-kernel basis
centered on a structured ``N_POINTS_PER_DIM x N_POINTS_PER_DIM`` grid, which
is also the set of interior collocation points; its boundary points are the
boundary collocation points, where Dirichlet is imposed weakly. The kernel
width is tied to the grid spacing (``sigma = SIGMA_FACTOR * spacing``):
conditioning is less forgiving than in 1D, and ``SIGMA_FACTOR`` in
``{3.0, 3.5}`` on even grid sizes was found to behave well (around 2, the fit
becomes ill-conditioned). The time integration uses the Pareschi-Russo
Runge-Kutta scheme through ``TimeDiscreteCollocationScheme``, with ``NT`` time
steps, each stage being a dense Newton solve.

The time and space convergence of the three schemes (explicit/implicit Euler,
Pareschi-Russo) on this problem are studied in
``comparative_studies_jax/kernel/
advection_reaction_diffusion_2d_square_isotropic_diffusion_convergence.py``,
and measured by the ``advection_reaction_diffusion_2d`` datasets of
``benchmarks/benchmarks_jax/kernel``.
"""

import timeit

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

from scimba_jax.domains.meshless_domains.domains_2d import Square2D
from scimba_jax.linear_approximation.basis.kernel_basis import KernelBasis
from scimba_jax.linear_approximation.basis.kernel_function import GaussianKernel
from scimba_jax.linear_approximation.collocation.time_discrete_collocation_scheme import (
    TimeDiscreteCollocationScheme,
)
from scimba_jax.linear_approximation.variables.collocation_variables import (
    CollocationVariables,
)
from scimba_jax.physical_models.elliptic_pde.general_elliptic import (
    GeneralEllipticResidual,
)
from scimba_jax.time_discrete.butcher_tableau import build_pareschi_russo_tableau

T_FINAL = 0.1
ALPHA = 1.0  # decay rate of the manufactured solution's time factor
N_POINTS_PER_DIM = 16
SIGMA_FACTOR = 3.0
NT = 8
DT = T_FINAL / NT
N_EVAL = 60


def d_fn(x: jnp.ndarray, mu=None) -> jnp.ndarray:
    """Diffusion coefficient: d(x, y) = I (isotropic)."""
    return jnp.eye(x.shape[0])


def b_fn(x: jnp.ndarray, mu=None) -> jnp.ndarray:
    """Advection velocity: b = (1, 0)."""
    return jnp.array([1.0, 0.0])


def c_fn(x: jnp.ndarray, mu=None) -> jnp.ndarray:
    """Reaction coefficient: c = 1."""
    return 1.0


def u0(xy, mu=None):
    """Initial condition u(x, y, 0) = sin(pi*x) * sin(pi*y)."""
    x, y = xy[0:1], xy[1:2]
    return jnp.sin(jnp.pi * x) * jnp.sin(jnp.pi * y)


def u_exact(t, xy):
    """Exact solution u(x, y, t) = exp(-alpha*t) * sin(pi*x) * sin(pi*y)."""
    x, y = xy[0:1], xy[1:2]
    return (jnp.exp(-ALPHA * t) * jnp.sin(jnp.pi * x) * jnp.sin(jnp.pi * y)).squeeze()


def f_source(xy: jnp.ndarray, mu, t: float) -> jnp.ndarray:
    """RHS of du/dt + A(u) = f, with A(u) = -div(d*grad(u)) + b.grad(u) + c*u.

    With u = exp(-alpha*t)*sin(pi*x)*sin(pi*y), d = I, b = (1, 0), c = 1:
        du/dt = -alpha*u
        -div(d*grad(u)) = 2*pi^2*u   (laplacian(u) = -2*pi^2*u)
        b.grad(u) = du/dx = exp(-alpha*t)*pi*cos(pi*x)*sin(pi*y)
        c*u = u

    Args:
        xy: Point, shape ``(2,)``.
        mu: Unused parameters.
        t: Time.

    Returns:
        The source at ``(xy, t)``, shape ``(1,)``.
    """
    x, y = xy[0:1], xy[1:2]
    decay = jnp.exp(-ALPHA * t)
    u = decay * jnp.sin(jnp.pi * x) * jnp.sin(jnp.pi * y)
    du_dt = -ALPHA * u
    du_dx = decay * jnp.pi * jnp.cos(jnp.pi * x) * jnp.sin(jnp.pi * y)
    diffusion_term = 2 * jnp.pi**2 * u
    advection_term = du_dx
    reaction_term = u
    return du_dt + diffusion_term + advection_term + reaction_term


# ── Kernel collocation space on a structured grid ────────────────────────────
domain = Square2D([(-1.0, 1.0), (-1.0, 1.0)], is_main_domain=True)
xs_grid = jnp.linspace(-1.0, 1.0, N_POINTS_PER_DIM)
xx, yy = jnp.meshgrid(xs_grid, xs_grid, indexing="ij")
xy_centers = jnp.stack([xx.ravel(), yy.ravel()], axis=-1)
on_boundary = ((xx == -1.0) | (xx == 1.0) | (yy == -1.0) | (yy == 1.0)).ravel()
xy_centers_bc = xy_centers[on_boundary]
sigma = SIGMA_FACTOR * 2.0 / (N_POINTS_PER_DIM - 1)

basis = KernelBasis(
    dim=2,
    output_dim=1,
    kernel_function=GaussianKernel(sigma=sigma),
    centers=xy_centers,
    basis_type="scalar",
)
variables = CollocationVariables(basis=basis, nb_variables=1)

# ── Time-discrete scheme ─────────────────────────────────────────────────────
scheme = TimeDiscreteCollocationScheme(
    spatial_residual_factory=lambda t: GeneralEllipticResidual(
        domain=domain,
        model_type="x_mu",
        f_rhs=lambda x, mu: f_source(x, mu, t),
        A=d_fn,
        b=b_fn,
        c=c_fn,
    ),
    variables=variables,
    collocation_points=xy_centers,
    bc_collocation_points=xy_centers_bc,
    main_domain=domain,
    butcher_tableau=build_pareschi_russo_tableau(),
    dt=DT,
    dirichlet_factory=lambda t: (lambda x, n, mu: jnp.zeros(1)),
)

print(
    "Solving the 2D advection-diffusion-reaction equation with kernel collocation: "
    f"{N_POINTS_PER_DIM}x{N_POINTS_PER_DIM} points, sigma_factor={SIGMA_FACTOR}, "
    f"nt={NT}, dt={DT:.2e} -- Pareschi-Russo"
)
start = timeit.default_timer()
dofsl_init = scheme.initialize(u0)
dofsl_final, history = scheme.solve(dofsl_init, t0=0.0, nt=NT)
# The solve is dispatched asynchronously: block before reading the clock.
jax.block_until_ready(history)
end = timeit.default_timer()
print(f"time for {NT} time steps: {end - start:.2f}s")

# ── Error against the exact solution ─────────────────────────────────────────
xs = jnp.linspace(-1.0, 1.0, N_EVAL)
xx_eval, yy_eval = jnp.meshgrid(xs, xs, indexing="ij")
xy_eval = jnp.stack([xx_eval.ravel(), yy_eval.ravel()], axis=-1)


def u_h_at(dofsl):
    """u_h on the evaluation grid for the given DOFs, shape (N_EVAL, N_EVAL)."""
    variables.dofsl = dofsl
    return jax.vmap(variables.evaluate)(xy_eval).reshape(N_EVAL, N_EVAL)


def u_ex_at(t):
    """u_exact(t, .) on the evaluation grid, shape (N_EVAL, N_EVAL)."""
    return jax.vmap(lambda xy: u_exact(t, xy))(xy_eval).reshape(N_EVAL, N_EVAL)


def integral(values):
    """2D trapezoidal quadrature on the evaluation grid."""
    return jnp.trapezoid(jnp.trapezoid(values, xs, axis=1), xs)


u_h = u_h_at(dofsl_final)
u_ex = u_ex_at(T_FINAL)
# Collocation has no mesh/quadrature machinery: the L2 norm is estimated by
# trapezoidal quadrature on a dense evaluation grid.
rel_l2 = jnp.sqrt(integral((u_h - u_ex) ** 2) / integral(u_ex**2))
print(f"rel. L2 error at t={T_FINAL}: {rel_l2:.3e}")
print(f"max abs error at t={T_FINAL}: {jnp.max(jnp.abs(u_h - u_ex)):.3e}")

# ── Plot: solution and pointwise error at a few time steps ───────────────────
steps = (0, NT // 2, NT)
fig, axes = plt.subplots(2, len(steps), figsize=(13, 8))
for ax_u, ax_err, step in zip(axes[0], axes[1], steps):
    t = step * DT
    u_h_t = u_h_at(history[step])
    err_t = jnp.abs(u_h_t - u_ex_at(t))

    im = ax_u.pcolormesh(xx_eval, yy_eval, u_h_t, shading="auto", cmap="viridis")
    fig.colorbar(im, ax=ax_u)
    ax_u.set_title(f"collocation, t={t:.3f}")

    im = ax_err.pcolormesh(xx_eval, yy_eval, err_t, shading="auto", cmap="magma")
    fig.colorbar(im, ax=ax_err)
    ax_err.set_title(rf"$|u_h - u_\mathrm{{exact}}|$, t={t:.3f}")

for ax in axes.ravel():
    ax.set_aspect("equal")
    ax.set_xlabel("x")
    ax.set_ylabel("y")

fig.suptitle(
    "2D advection-diffusion-reaction, kernel collocation "
    f"(Pareschi-Russo, {N_POINTS_PER_DIM}x{N_POINTS_PER_DIM} points, "
    f"{NT} time steps, rel. L2 error at t={T_FINAL}: {rel_l2:.2e})"
)
fig.tight_layout()
plt.show()