"""1D Sod shock tube for the compressible Euler equations, DG + HLLC.

    d/dt (rho, rho*u, E) + d/dx F(rho, rho*u, E) = 0   on [0, 1]

    F(rho, rho*u, E) = (rho*u, rho*u^2 + p, u*(E + p)),
    p = (gamma - 1) * (E - 1/2 rho u^2),  gamma = 1.4

    Riemann initial data at x0 = 0.5:
        (rho, u, p)_L = (1, 0, 1)          x < x0
        (rho, u, p)_R = (0.125, 0, 0.1)    x > x0

Toro's classic test problem (Toro, *Riemann Solvers and Numerical Methods for
Fluid Dynamics*, 3rd ed., ch. 4). Its exact solution at t = 0.2 has all three
elementary waves of a 1D Riemann problem at once: a left rarefaction fan
(x in [0.263, 0.486]), a contact discontinuity (x = 0.685), and a right shock
(x = 0.850) -- the standard benchmark for a shock-capturing scheme. It is
drawn dashed, from the library's exact Riemann solver
(:func:`~scimba_jax.physical_models.classical_weakform.euler_weak_form.
euler_exact_riemann`).

A conservative flux, integrated by parts once, paired with a face-term
numerical flux and :class:`~scimba_jax.linear_approximation.galerkin.dg.
artificial_viscosity.ArtificialViscosity` for shock capturing, on the
3-component Euler *system* (``basis_type="vec"``, ``out_dim=3``). The volume
weak form (:class:`~scimba_jax.physical_models.classical_weakform.
euler_weak_form.EulerWeakForm1D`) and the flux/wave-speed pair
(:func:`~scimba_jax.physical_models.classical_weakform.euler_weak_form.
make_euler_flux_fns`) live in the library's own ``euler_weak_form.py``. The
face term is the HLLC three-wave solver of ``finite_volume.hyperbolic``
(the one ``examples_jax/vf/hyperbolic/euler_sod_1d.py`` runs on cell
averages), applied to the DG traces by :class:`~scimba_jax.
linear_approximation.galerkin.dg.flux.HyperbolicFluxDG` -- a face asks the
same question of both schemes, only where the two states come from differs.
``ArtificialViscosity``'s own weak-form/flux pair
(``WithArtificialViscosity``/``SIPGFlux``) handles a vector state directly.

The comparisons this example used to draw -- LLF, HLL, HLLC and Roe, and a
non-uniform mesh three times finer right of ``x0`` -- are rows of the
benchmark ``benchmarks/benchmarks_jax/dg_hyperbolic`` (dataset ``sod``),
measured against the same exact solution.
"""

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.finite_volume.hyperbolic import EulerHLLCFlux
from scimba_jax.linear_approximation.galerkin.dg.artificial_viscosity import (
    ArtificialViscosity,
)
from scimba_jax.linear_approximation.galerkin.dg.flux import HyperbolicFluxDG
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.euler_weak_form import (
    EulerPDE,
    EulerWeakForm1D,
    euler_exact_riemann,
    make_euler_flux_fns,
)
from scimba_jax.time_discrete.butcher_tableau import build_ssprk3_tableau

DIM = 1
OUT_DIM = 3  # (rho, rho*u, E)
GAMMA = 1.4
X0 = 0.5  # initial discontinuity

RHO_L, U_L, P_L = 1.0, 0.0, 1.0
RHO_R, U_R, P_R = 0.125, 0.0, 0.1

N_CELLS = 64
ORDER = 2
QUAD_ORDER = 3
T_FINAL = 0.2


# ── Physics: conservative flux, pressure, wave-speed bound ──────────────────

euler_flux, euler_wave_speed = make_euler_flux_fns(GAMMA)
# The physical flux the FV solvers read, F(w) as the conservative PDE builds it.
euler_pde = EulerPDE(dim=1, gamma=GAMMA)


# ── Initial / boundary data ──────────────────────────────────────────────────


def sod_state(x):
    """Riemann data: left state for x < x0, right state for x > x0.

    Used both as the initial condition and, unchanged, as the Dirichlet
    datum -- the waves never reach x=0 or x=1 by t=T_FINAL, so the domain
    ends are exact ghost states throughout the run.
    """
    rho = jnp.where(x[0] < X0, RHO_L, RHO_R)
    p = jnp.where(x[0] < X0, P_L, P_R)
    u = jnp.where(x[0] < X0, U_L, U_R)
    mom = rho * u
    energy = p / (GAMMA - 1.0) + 0.5 * rho * u * u
    return jnp.stack([rho, mom, energy])


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


def local_basis(y, i, mesh):
    """The Taylor ``Q_ORDER`` basis of a 3-component state."""
    return local_taylor_basis(y, i, mesh, order=ORDER, out_dim=OUT_DIM)


def solve():
    """Run the Sod problem to T_FINAL on the uniform Cartesian mesh."""
    h = 1.0 / N_CELLS
    mesh = cartesian_mesh(n_cells=[N_CELLS], quad_order=QUAD_ORDER)
    basis = AnalyticBasis(
        nb_basis=ORDER + 1,
        out_dim=OUT_DIM,
        mesh=mesh,
        local_basis=local_basis,
        basis_type="vec",
    )
    variables = VariablesDG(basis=basis, nb_variables=OUT_DIM)
    # nu0 tuned down from this class' nu0=1 default (itself well below the
    # nu0=2 this file used to carry): at nu0=2 the added diffusion dominates
    # the contact/shock so completely that LLF, HLL and HLLC land on
    # indistinguishable curves. nu0=0.15 was picked from a sweep (nu0 in
    # [2, 0.02], order=2, N_CELLS=64): it roughly halves the contact's
    # numerical width versus nu0=2 (~0.035 vs ~0.07 in x) with no visible
    # undershoot (min density within 0.001 of the exact 0.125) and a wide
    # safety margin above the one real failure mode found while sweeping:
    # below nu0~0.05, HLL (and, at order 1, LLF too) develops a spurious
    # stationary spike that never advects off x0=0.5, a known HLL-type
    # pathology for a strong shock sitting exactly on a cell interface (as
    # Sod's does here -- X0=0.5 lands on a cell boundary for N_CELLS=64).
    # HLLC's own construction is robust to it.
    artificial_viscosity = ArtificialViscosity(
        wave_speed=euler_wave_speed, kappa=0.5, nu0=0.15
    )

    # CFL, explicit SSPRK3: dt ~ CFL * h / ((2p+1) * wave_speed). The
    # Riemann problem's fastest wave (the shock, speed ~1.75) is a bit faster
    # than either initial state's own sound speed, so a speed_scale margin
    # above both (~1.18, ~1.06) is used rather than the initial condition's
    # own max.
    speed_scale = 3
    dt = 0.15 * h / ((2 * ORDER + 1) * speed_scale)
    nt = round(T_FINAL / dt)
    print(f"Sod, DG Q{ORDER} + HLLC: n_cells={N_CELLS}, dt={dt:.3e}, nt={nt}")

    scheme = TimeDiscreteDGscheme(
        spatial_weak_form_factory=EulerWeakForm1D(dim=DIM, gamma=GAMMA),
        variables=variables,
        flux=HyperbolicFluxDG(EulerHLLCFlux(GAMMA), euler_pde),
        butcher_tableau=build_ssprk3_tableau(),
        dt=dt,
        dirichlet=sod_state,
        artificial_viscosity=artificial_viscosity,
    )
    dofsl_init = scheme.initialize(sod_state)
    dofsl_final, _ = scheme.solve(dofsl_init, t0=0.0, nt=nt, keep_history=False)
    return variables, dofsl_final


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


def plot(variables, dofsl):
    """Density, velocity and pressure against the exact Riemann solution."""
    x_plot = jnp.linspace(0.0, 1.0, 600)[:, None]
    variables.dofsl = dofsl
    w = variables.evaluate(x_plot)
    rho, mom, energy = w[:, 0], w[:, 1], w[:, 2]
    vel = mom / rho
    p = (GAMMA - 1.0) * (energy - 0.5 * mom * vel)

    exact = euler_exact_riemann((RHO_L, U_L, P_L), (RHO_R, U_R, P_R), GAMMA)
    x_exact = np.linspace(0.0, 1.0, 2000)
    fields_exact = exact(x_exact - X0, T_FINAL)

    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    for ax, field, field_exact, label in zip(
        axes, (rho, vel, p), fields_exact, (r"$\rho$", "u", "p")
    ):
        ax.plot(x_plot[:, 0], field, linewidth=1.8, label=f"DG Q{ORDER} + HLLC")
        ax.plot(x_exact, field_exact, "k--", linewidth=1.2, label="exact")
        ax.set_xlabel("x")
        ax.set_ylabel(label)
        ax.grid(alpha=0.3)
        ax.legend(fontsize=9)
    axes[0].set_title("density")
    axes[1].set_title("velocity")
    axes[2].set_title("pressure")
    fig.suptitle(
        f"Sod shock tube, t={T_FINAL} -- Cartesian mesh, n_cells={N_CELLS}, "
        "HLLC + Persson-Peraire artificial viscosity"
    )
    fig.tight_layout()


if __name__ == "__main__":
    variables, dofsl = solve()
    plot(variables, dofsl)
    plt.show()