"""1D linear transport equation with TimeDiscreteDGscheme, periodic BC.

    du/dt + a du/dx = 0   on  [0, 1],  periodic
    u(x, 0) = shu_ic(x)

``shu_ic`` is the classic test problem for a numerical advection scheme's
dispersion/dissipation error (Shu & Osher, "Efficient Implementation of
Essentially Non-Oscillatory Shock-Capturing Schemes, II", JCP 83 (1989);
also Jiang & Shu, JCP 126 (1996), sec. 5): a combination of four features on
one period -- a smooth triple-Gaussian bump, a flat-top square wave, a
triangle wave with a genuine corner, and a smooth half-ellipse -- exercising
a scheme's smooth-region accuracy, its handling of a jump discontinuity, a
slope discontinuity, and a curved-but-not-smooth-derivative feature, all in
one initial condition. Rescaled here from the original [-1, 1] domain / [0, 1]
range to this folder's [0, 1] domain / [1, 2] range: same shapes, same order
along the interval, only the coordinate frame differs.

Advecting it at speed a = 1 for one full period (T = 1) on a periodic [0, 1]
domain returns *exactly* the initial condition -- so ``shu_ic`` doubles as
the exact reference at t = T, unlike Burgers' (dataset ``burgers`` of the
benchmark ``benchmarks/benchmarks_jax/dg_hyperbolic``), which needs a
method-of-characteristics solver because its solution changes shape (and
eventually shocks). It reuses exactly the same library volume weak form --
:class:`~scimba_jax.physical_models.classical_weakform.
conservation_law_weak_form.ConservationLawWeakForm1D` -- and
:class:`~scimba_jax.linear_approximation.galerkin.dg.flux.
LocalLaxFriedrichsFlux` pairing as Burgers' (with a *linear* flux f(u) = a u
this time, built by :func:`~scimba_jax.physical_models.classical_weakform.
conservation_law_fluxes.make_linear_advection_flux_fns` -- Burgers' own
``burgers_flux`` sits right next to it in that module --
LocalLaxFriedrichsFlux's Rusanov formula then collapses to the exact upwind
flux, since alpha = |f'(u)| = |a| is constant) and the same
:class:`~scimba_jax.linear_approximation.galerkin.dg.
artificial_viscosity.ArtificialViscosity` limiter, demonstrating all three
are general DG infrastructure, not Burgers-specific. The same benchmark
(dataset ``transport_1d``) measures this case's error against ``shu_ic``
across degrees, with and without the limiter.

**Periodic boundary condition.** Neither ``Mesh`` nor ``EllipticDGscheme``
has any notion of periodicity: a boundary condition there is always a
Dirichlet ghost value ``x -> g``, which cannot express "whatever the
solution's own trace at the other end currently is". So this problem uses
:class:`~scimba_jax.linear_approximation.galerkin.dg.periodic_dg_scheme.
PeriodicEllipticDGscheme` (``scheme_cls=`` below) instead, which adds the
missing coupling between the mesh's two ends as one more interior-style face
-- see that module's docstring for why periodicity is not a boundary
condition to plug into the existing ``dirichlet=`` slot at all. Register no
Dirichlet condition (``dirichlet=None``): the two mesh-boundary faces
contribute nothing on their own once wrapped this way, by design.
"""

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

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.artificial_viscosity import (
    ArtificialViscosity,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import LocalLaxFriedrichsFlux
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.variables.variables_dg import VariablesDG
from scimba_jax.physical_models.classical_weakform.conservation_law_fluxes import (
    make_linear_advection_flux_fns,
)
from scimba_jax.physical_models.classical_weakform.conservation_law_weak_form import (
    ConservationLawWeakForm1D,
)
from scimba_jax.time_discrete.butcher_tableau import build_rk4_tableau

DIM = 1
A_SPEED = 1.0
N_CELLS = 100
T_FINAL = 1.0  # one full period at speed A_SPEED on a length-1 domain

# f(u) = a u, f'(u) = a: linear, but the same ConservationLawWeakForm1D /
# LocalLaxFriedrichsFlux pairing as Burgers'.
transport_flux, dtransport_flux = make_linear_advection_flux_fns(A_SPEED)


# ── Initial condition: Shu & Osher's four-feature test ──────────────────────


def _gaussian(x, beta, z):
    return jnp.exp(-beta * (x - z) ** 2)


def _ellipse(x, alpha, a):
    return jnp.sqrt(jnp.maximum(1.0 - alpha**2 * (x - a) ** 2, 0.0))


def shu_ic(x):
    """Shu & Osher's four-feature test, rescaled from [-1, 1]/[0, 1] to this
    folder's [0, 1]/[1, 2] convention: x0 = 2x - 1 recovers the classical
    coordinate, and +1 lifts the classical [0, 1] range to [1, 2].
    """
    x0 = 2.0 * x - 1.0
    a, z, delta, alpha = 0.5, -0.7, 0.005, 10.0
    beta = jnp.log(2.0) / (36.0 * delta**2)

    gaussian_bump = (
        _gaussian(x0, beta, z - delta)
        + _gaussian(x0, beta, z + delta)
        + 4.0 * _gaussian(x0, beta, z)
    ) / 6.0
    ellipse_bump = (
        _ellipse(x0, alpha, a - delta)
        + _ellipse(x0, alpha, a + delta)
        + 4.0 * _ellipse(x0, alpha, a)
    ) / 6.0
    triangle = 1.0 - jnp.abs(10.0 * (x0 - 0.1))

    val = jnp.where(
        (x0 >= -0.8) & (x0 <= -0.6),
        gaussian_bump,
        jnp.where(
            (x0 >= -0.4) & (x0 <= -0.2),
            1.0,
            jnp.where(
                (x0 >= 0.0) & (x0 <= 0.2),
                triangle,
                jnp.where((x0 >= 0.4) & (x0 <= 0.6), ellipse_bump, 0.0),
            ),
        ),
    )
    return val + 1.0


# ── DG setup and time-stepping ───────────────────────────────────────────────


def make_variables(n_cells, order, quad_order):
    mesh = cartesian_mesh(n_cells=[n_cells], quad_order=quad_order)
    basis = AnalyticBasis(
        nb_basis=order + 1,
        out_dim=1,
        mesh=mesh,
        local_basis=lambda y, i, m, k=order: local_taylor_basis(
            y, i, m, order=k, out_dim=1
        ),
        basis_type="scalar",
    )
    return VariablesDG(basis=basis, nb_variables=1)


def solve(butcher_tableau, order, dt, nt, limiter=False):
    quad_order = order + 3
    variables = make_variables(N_CELLS, order, quad_order)

    artificial_viscosity = (
        ArtificialViscosity(wave_speed=dtransport_flux) if limiter else None
    )

    scheme = TimeDiscreteDGscheme(
        spatial_weak_form_factory=ConservationLawWeakForm1D(
            dim=DIM, flux_fn=transport_flux
        ),
        variables=variables,
        flux=LocalLaxFriedrichsFlux(transport_flux, dtransport_flux),
        butcher_tableau=butcher_tableau,
        dt=dt,
        dirichlet=None,
        scheme_cls=PeriodicEllipticDGscheme,
        artificial_viscosity=artificial_viscosity,
    )
    dofsl_init = scheme.initialize(shu_ic)
    dofsl_final, _ = scheme.solve(dofsl_init, t0=0.0, nt=nt)
    return variables, dofsl_final


# ── Plotting ─────────────────────────────────────────────────────────────────


def plot_panel(ax, dofsl_unlimited, dofsl_limited, variables, title):
    x_plot = jnp.linspace(0.0, 1.0, 400)[:, None]
    u_exact = shu_ic(x_plot[:, 0])

    variables.dofsl = dofsl_unlimited
    u_unlimited = variables.evaluate(x_plot)[:, 0]
    variables.dofsl = dofsl_limited
    u_limited = variables.evaluate(x_plot)[:, 0]

    ax.plot(x_plot[:, 0], u_exact, "k-", linewidth=1.5, label="exact (= u0)")
    ax.plot(x_plot[:, 0], u_unlimited, color="tab:red", label="no limiter")
    ax.plot(x_plot[:, 0], u_limited, color="tab:green", label="with limiter")
    ax.set_xlabel("x")
    ax.set_ylabel("u")
    ax.set_title(title)
    ax.legend(fontsize=8)
    ax.grid(True, alpha=0.3)


if __name__ == "__main__":
    h = 1.0 / N_CELLS
    ORDERS = [2]

    # See the module docstring: a much tighter CFL than Burgers' needed here.
    dt_explicit = {p: 0.05 * h / (2 * p + 1) for p in ORDERS}

    configs = {"RK4": (build_rk4_tableau(), dt_explicit)}

    fig, axes = plt.subplots(
        len(ORDERS),
        len(configs),
        figsize=(5 * len(ORDERS), 5 * len(configs)),
        squeeze=False,
    )

    for col, (name, (tableau, dt_by_order)) in enumerate(configs.items()):
        for row, order in enumerate(ORDERS):
            dt = dt_by_order[order]
            nt = round(T_FINAL / dt)
            print(f"{name}, p={order}: dt={dt:.3e}, nt={nt}, no limiter")
            variables_u, dofsl_unlimited = solve(tableau, order, dt, nt, limiter=False)
            print(f"{name}, p={order}: dt={dt:.3e}, nt={nt}, with limiter")
            _, dofsl_limited = solve(tableau, order, dt, nt, limiter=True)
            plot_panel(
                axes[row][col],
                dofsl_unlimited,
                dofsl_limited,
                variables_u,
                title=f"{name}, p={order} (n_cells={N_CELLS})",
            )

    fig.suptitle(
        "Linear transport, periodic, u(x,0)=shu_ic(x), a=1, one full period\n"
        "DG + LocalLaxFriedrichsFlux (upwind), with/without "
        "Persson-Peraire artificial viscosity"
    )
    fig.tight_layout()
    plt.show()