r"""Gray-Scott reaction-diffusion system in 2D, periodic -- with TimeDiscreteDGscheme.

.. math::

    \partial_t u &= \varepsilon_1 \Delta u + b_1(1 - u) - c_1 u v^2, \\
    \partial_t v &= \varepsilon_2 \Delta v - b_2 v + c_2 u v^2,

on :math:`(x, y) \in (-1, 1)^2`, periodic boundary conditions, with initial
conditions

.. math::

    u_0(x, y) &= 1 - \exp\!\bigl(-10\,((x+0.05)^2 + (y+0.02)^2)\bigr), \\
    v_0(x, y) &= \exp\!\bigl(-10\,((x-0.05)^2 + (y-0.02)^2)\bigr).

Same test case (domain, ICs, coefficients) as
``examples_jax/pinns/time_dependent_pdes/parabolic_systems/gray_scott_2d.py``
-- kept unchanged so the two are comparable -- but solved by classical
time-marching DG instead of a multi-window PINN. Validated the same way that
file is: against an independent RK4 finite-difference solution on a fine
periodic grid (no analytical solution exists for this system). Per-task,
this is a single mesh/degree/time-integrator check, not a convergence sweep
-- see the DG rows of ``benchmarks/benchmarks_jax/heat_time_convergence``
for that.

Two things this system needs beyond the scalar heat-equation example:

1. **A coupled system on one DG space.** ``u``/``v`` live on one
   ``out_dim=2`` ("vec") space, exactly the pattern
   ``tests/test_jax/linear_approximation/test_multi_space_dg.py`` checks
   against two separate scalar spaces: ``u.components()`` splits the trial
   trace into ``u1, u2``, and the *nonlinear* reaction term (``u1*u2*u2``,
   depending on the trial functions themselves) goes into ``bilinear_form``
   alongside the diffusion, not into ``linear_form`` -- a ``linear_form`` may
   depend on ``x``/``t`` only. :class:`~scimba_jax.physical_models.
<<<<<<<< HEAD:examples/examples_jax/dg/time_dependent/2d_grey_scott_reaction_diffusion.py
   classical_weakform.grey_scott_weak_form.GreyScottWeakForm` writes exactly
   this. The per-stage solve is already a Newton solve (see CLAUDE.md: "the
   Galerkin solve is already a Newton"), so a nonlinear volume term costs
   nothing extra, exactly as the nonlinear ``burgers_flux`` of the
   ``burgers`` dataset of ``benchmarks/benchmarks_jax/dg_hyperbolic`` (fed into :class:`~scimba_jax.physical_models.
========
   classical_weakform.gray_scott_weak_form.GrayScottWeakForm` writes exactly
   this. The per-stage solve is already a Newton solve (see CLAUDE.md: "le
   solve Galerkin est déjà un Newton"), so a nonlinear volume term costs
   nothing extra, exactly as ``1d_burgers_equation.py``'s nonlinear
   ``burgers_flux`` (fed into :class:`~scimba_jax.physical_models.
>>>>>>>> develop:examples/examples_jax/dg/time_dependent/2d_gray_scott_reaction_diffusion.py
   classical_weakform.conservation_law_weak_form.ConservationLawWeakForm1D`)
   shows for the hyperbolic case.

2. **Different diffusivities per component need a custom flux.**
   :class:`~scimba_jax.linear_approximation.galerkin.dg.flux.SIPGFlux` reads
   a single ``fields["A"]`` -- one ``(dim, dim)`` tensor applied to *every*
   component identically. That is enough for a system whose components share
   one diffusion operator (see the coupled test above, or
   ``solve_2d_periodic_reaction_diffusion.py``), but ``D_u != D_v`` here
   needs a per-*component* scaling, a different axis than what ``A``
   multiplies (component axis, not spatial axis) -- confirmed by direct
   inspection: the trial Jacobian ``u.jacobian("x")`` a face passes to the
   flux has shape ``(out_dim, dim)``, and ``SIPGFlux``'s ``A @ gradu``
   contracts ``A``'s columns against ``gradu``'s *rows* (the component axis).
   That is only harmless when ``A`` is a scalar multiple of the identity
   (any of the same-diffusivity systems above), where the contraction axis
   does not matter; it silently mixes components otherwise. So
   :class:`~scimba_jax.linear_approximation.galerkin.dg.flux.
   DiagonalSIPGFlux` reimplements the same SIPG bilinear form with a
   per-component diffusivity vector instead of a shared spatial tensor --
   new physics (anisotropic-*across-components* diffusion), not a duplicate
   of ``SIPGFlux``, in the same spirit as ``LocalLaxFriedrichsFlux`` being a
   different flux from ``SIPGFlux`` rather than a variant of it. Both this
   flux and :class:`~scimba_jax.physical_models.classical_weakform.
   diagonal_reaction_diffusion_weak_form.DiagonalReactionDiffusionWeakForm`
   (the ``GrayScottWeakForm`` base above) now live in the library, not this
   example -- the diffusion part is generic to any diagonal
   reaction-diffusion system, only the reaction is Gray-Scott-specific.

Periodic BC comes from
:class:`~scimba_jax.linear_approximation.galerkin.dg.periodic_dg_scheme.PeriodicEllipticDGscheme`
passed as ``TimeDiscreteDGscheme(..., scheme_cls=...)`` -- DG's own
mechanism for gluing a Cartesian mesh into a torus, already exercised (for a
single-diffusivity problem) by
``solve_2d_periodic_reaction_diffusion.py``. No such scheme exists for FEM
(:mod:`scimba_jax.linear_approximation.galerkin.fem` has no periodic
variant), which is why the FEM counterpart of this file
(``examples_jax/fem/solve/time_dependent/2d_gray_scott_reaction_diffusion.py``)
uses natural (zero-flux Neumann) boundaries instead -- close enough for this
short a time horizon, since the two Gaussian blobs sit well inside the
domain and never reach the boundary, but flagged there explicitly.

The reaction coefficients (``c1 = c2 = 1000``) make this a genuinely stiff
system -- explicit RK schemes here would need a prohibitively small ``dt``
for stability having nothing to do with accuracy, exactly the case CLAUDE.md
flags implicit/ENG-style solves for. An L-stable, 2-stage SDIRK
(Pareschi-Russo, one of the implicit tableaux of
``benchmarks/benchmarks_jax/heat_time_convergence``) is used here, fully implicit (diffusion *and*
reaction Newton-solved together each stage -- this codebase's time-dependent
Galerkin schemes reject IMEX tableaux outright, so there is no
partially-explicit alternative for the reaction).
"""

import time

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

from scimba_jax.linear_approximation.basis.analytic_bases import local_taylor_basis
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.galerkin.dg.flux import DiagonalSIPGFlux
from scimba_jax.linear_approximation.galerkin.dg.periodic_dg_scheme import (
    PeriodicEllipticDGscheme,
)
from scimba_jax.linear_approximation.galerkin.dg.time_discrete_dg_scheme import (
    TimeDiscreteDGscheme,
)
from scimba_jax.linear_approximation.meshes.cartesian_mesh import cartesian_mesh
from scimba_jax.linear_approximation.solvers.preconditioners import (
    BlockJacobiPreconditioner,
)
from scimba_jax.linear_approximation.variables.variables_dg import VariablesDG
from scimba_jax.physical_models.classical_weakform.gray_scott_weak_form import (
    GrayScottWeakForm,
    gray_scott_fd_grid,
    gray_scott_fd_reference,
)
from scimba_jax.time_discrete.butcher_tableau import build_pareschi_russo_tableau

jax.config.update("jax_enable_x64", True)

DIM = 2

# ── Problem parameters (identical to the PINN reference test case) ───────────

EPS1 = 0.2  # D_u
EPS2 = 0.1  # D_v
B1 = 40.0  # feed rate F
B2 = 100.0  # F + k
C1 = 1000.0  # reaction coefficient for u
C2 = 1000.0  # reaction coefficient for v

A_LO, A_HI = -1.0, 1.0  # domain (-1, 1)^2
BOX_BOUNDS = ((A_LO, A_HI), (A_LO, A_HI))


def u0v0(xy):
    """Initial condition, pointwise: (x, y) -> (u0, v0)."""
    x, y = xy[0], xy[1]
    u0 = 1.0 - jnp.exp(-10.0 * ((x + 0.05) ** 2 + (y + 0.02) ** 2))
    v0 = jnp.exp(-10.0 * ((x - 0.05) ** 2 + (y - 0.02) ** 2))
    return jnp.stack([u0, v0])


# ── DG space and solve ────────────────────────────────────────────────────────


def make_variables(n_cells: int, order: int, quad_order: int) -> VariablesDG:
    mesh = cartesian_mesh(
        n_cells=(n_cells, n_cells), quad_order=quad_order, bounds=BOX_BOUNDS
    )
    basis = AnalyticBasis(
        nb_basis=(order + 1) ** DIM,
        out_dim=2,
        mesh=mesh,
        local_basis=lambda y, i, m, k=order: local_taylor_basis(
            y, i, m, order=k, out_dim=2
        ),
        basis_type="vec",
    )
    return VariablesDG(basis=basis, nb_variables=2)


def solve(n_cells: int, order: int, dt: float, nt: int):
    quad_order = order + 3
    h = (A_HI - A_LO) / n_cells
    sigma = (order + 1) * (order + 2)
    variables = make_variables(n_cells, order, quad_order)
    flux = DiagonalSIPGFlux(sigma=sigma, h=h)
    weak_form = GrayScottWeakForm(
        dim=DIM, D_u=EPS1, D_v=EPS2, b1=B1, b2=B2, c1=C1, c2=C2
    )
    scheme = TimeDiscreteDGscheme(
        spatial_weak_form_factory=weak_form,
        variables=variables,
        flux=flux,
        butcher_tableau=build_pareschi_russo_tableau(),
        dt=dt,
        scheme_cls=PeriodicEllipticDGscheme,
        max_iter=30,
        tol=1e-10,
        cg_solver="bicgstab",
        preconditioner=BlockJacobiPreconditioner(),
    )
    dofsl_init = scheme.initialize(u0v0)
    dofsl_final, history = scheme.solve(dofsl_init, t0=0.0, nt=nt)
    return variables, dofsl_final, history


# ── Reference: RK4 finite-difference solve on a fine periodic grid ───────────
# Independent of the DG discretization (and of the PINN example's own
# reference, which validates itself the same way): see
# gray_scott_fd_reference in the library for why it lives there.

N_FD = 160
_x_fd, _xy_fd = gray_scott_fd_grid((A_LO, A_HI), N_FD)


def fd_reference(t_final: float, dt_fd: float = 2e-5):
    return gray_scott_fd_reference(
        u0v0, (A_LO, A_HI), EPS1, EPS2, B1, B2, C1, C2, t_final, N_FD, dt_fd
    )


if __name__ == "__main__":
    N_CELLS = 24
    ORDER = 2
    DT = 1e-2
    NT = 40
    T_FINAL = DT * NT

    print(
        f"Solving Gray-Scott with DG: n_cells={N_CELLS}, order={ORDER}, "
        f"dt={DT:.1e}, nt={NT} (T_final={T_FINAL:.4f}) -- Pareschi-Russo, periodic"
    )
    t0 = time.time()
    variables, dofsl_final, _history = solve(N_CELLS, ORDER, DT, NT)
    print(f"DG solve done in {time.time() - t0:.1f}s")

    print(f"Solving independent RK4/FD reference on a {N_FD}x{N_FD} periodic grid...")
    t0 = time.time()
    u_ref, v_ref = fd_reference(T_FINAL)
    print(f"FD reference done in {time.time() - t0:.1f}s")

    variables.dofsl = dofsl_final
    uv_dg = variables.evaluate(_xy_fd)
    u_dg = np.array(uv_dg[:, 0]).reshape(N_FD, N_FD)
    v_dg = np.array(uv_dg[:, 1]).reshape(N_FD, N_FD)

    err_u = np.abs(u_dg - u_ref)
    err_v = np.abs(v_dg - v_ref)
    rel_l2_u = np.linalg.norm(err_u) / np.linalg.norm(u_ref)
    rel_l2_v = np.linalg.norm(err_v) / np.linalg.norm(v_ref)
    print(
        f"u: rel. L2 error vs. FD reference = {rel_l2_u:.3e}, max abs = {err_u.max():.3e}"
    )
    print(
        f"v: rel. L2 error vs. FD reference = {rel_l2_v:.3e}, max abs = {err_v.max():.3e}"
    )

    # ── Plot: DG vs. FD reference vs. |difference| ────────────────────────────
    X, Y = np.array(_x_fd), np.array(_x_fd)
    fig, axes = plt.subplots(2, 3, figsize=(15, 9))
    panels = [
        (u_dg, "$u$ DG"),
        (u_ref, "$u$ FD reference"),
        (err_u, r"$|u_\mathrm{DG} - u_\mathrm{FD}|$"),
        (v_dg, "$v$ DG"),
        (v_ref, "$v$ FD reference"),
        (err_v, r"$|v_\mathrm{DG} - v_\mathrm{FD}|$"),
    ]
    for ax, (grid, title) in zip(axes.flat, panels):
        im = ax.pcolormesh(X, Y, grid, cmap="turbo", shading="auto")
        fig.colorbar(im, ax=ax, pad=0.02)
        ax.set_title(title)
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_aspect("equal")
    fig.suptitle(
        rf"Gray-Scott 2D, DG (Pareschi-Russo, p={ORDER}, {N_CELLS}x{N_CELLS} cells, "
        rf"periodic) vs. FD reference at $t={T_FINAL:.4f}$"
    )
    fig.tight_layout()
    plt.show()