"""Sod's shock tube in finite volume: Rusanov, HLL, HLLC and Roe on one plot.

    W_t + F(W)_x = 0,   W = (rho, rho u, E),   gamma = 1.4
    (rho, u, p) = (1, 0, 1) for x < 0.5,  (0.125, 0, 0.1) for x > 0.5,  t = 0.2

The classical test (Toro, *Riemann Solvers and Numerical Methods for Fluid
Dynamics*, 3rd ed., sec. 4.3.3 and ch. 6): a left rarefaction, a contact and a
right shock, against the exact solution of :func:`euler_exact_riemann`.

First order in space and time (P0 cells, explicit Euler, as ``burgers_1d.py``),
so what separates the four fluxes is their dissipation, and where. Read the
density at the contact (x ~ 0.69): Rusanov and HLL, two-wave solvers with a
single intermediate state, smear it with O(h) diffusion; HLLC and Roe carry
the contact wave and keep it sharp. The shock (x ~ 0.85) is captured by all
four -- a first-order scheme's shock width is a few cells whatever the flux.

The fluxes are the n-D objects of ``finite_volume.hyperbolic.euler``, used
here with ``n = (1,)``; ``euler_cylindrical_explosion_2d.py`` runs the same
four in 2D.
"""

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

from scimba_jax.linear_approximation.finite_volume import (
    FiniteVolumeScheme,
    TimeDiscreteFVscheme,
)
from scimba_jax.linear_approximation.finite_volume.hyperbolic import (
    EulerHLLCFlux,
    EulerRoeFlux,
    HLLFlux,
    RusanovFlux,
    euler_hll_wave_speeds,
    euler_wave_speed,
)
from scimba_jax.linear_approximation.meshes.cartesian_mesh import cartesian_mesh
from scimba_jax.linear_approximation.variables.variables_fv import VariablesFV
from scimba_jax.physical_models.abstract_physical_conservative_model import (
    AbstractPhysicalConservativeModel,
)
from scimba_jax.physical_models.classical_weakform.euler_weak_form import (
    EulerPDE,
    euler_conserved,
    euler_exact_riemann,
    euler_primitives,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet
from scimba_jax.time_discrete.butcher_tableau import build_explicit_euler_tableau

GAMMA = 1.4
N_CELLS = 200
CFL = 0.8
T_FINAL = 0.2
# Bound on |u| + c over the WHOLE run: the star state moves at u* = 0.93 with
# c* = 1.0 (1.93), faster than anything the initial data shows (c = 1.18).
SPEED_BOUND = 2.0
X_DISCONTINUITY = 0.5
LEFT = (1.0, 0.0, 1.0)  # (rho, u, p)
RIGHT = (0.125, 0.0, 0.1)

W_LEFT = euler_conserved(LEFT[0], jnp.array([LEFT[1]]), LEFT[2], GAMMA)
W_RIGHT = euler_conserved(RIGHT[0], jnp.array([RIGHT[1]]), RIGHT[2], GAMMA)


def initial_condition(x):
    return jnp.where(x[0] < X_DISCONTINUITY, W_LEFT, W_RIGHT)


def make_time_scheme(flux, dt):
    mesh = cartesian_mesh([N_CELLS], quad_order=2)
    model = AbstractPhysicalConservativeModel(EulerPDE(dim=1, gamma=GAMMA))
    # The waves never reach the ends by t = 0.2, so the initial states are
    # exact ghost states throughout the run.
    model.add_boundary_condition("west", Dirichlet(lambda _x: W_LEFT))
    model.add_boundary_condition("east", Dirichlet(lambda _x: W_RIGHT))
    spatial = FiniteVolumeScheme(
        model,
        VariablesFV(mesh, nb_variables=3),
        flux,
        assemble_second_order=False,
        assemble_reaction=False,
        assemble_source=False,
    )
    return TimeDiscreteFVscheme(spatial, build_explicit_euler_tableau(), dt=dt)


def solve(flux, dt, nt):
    scheme = make_time_scheme(flux, dt)
    initial = scheme.initialize(initial_condition)
    final = scheme.solve_final(initial, t0=0.0, nt=nt)
    return scheme, final


def primitives(w):
    rho, velocity, pressure, _ = jax.vmap(lambda s: euler_primitives(s, GAMMA))(w)
    return np.asarray(rho), np.asarray(velocity[:, 0]), np.asarray(pressure)


if __name__ == "__main__":
    h = 1.0 / N_CELLS
    dt = CFL * h / SPEED_BOUND
    nt = round(T_FINAL / dt)
    dt = T_FINAL / nt

    fluxes = {
        "Rusanov": RusanovFlux(euler_wave_speed(GAMMA)),
        "HLL": HLLFlux(euler_hll_wave_speeds(GAMMA)),
        "HLLC": EulerHLLCFlux(GAMMA),
        "Roe": EulerRoeFlux(GAMMA),
    }
    exact = euler_exact_riemann(LEFT, RIGHT, GAMMA)
    x_exact = np.linspace(0.0, 1.0, 2000)
    rho_exact, u_exact, p_exact = exact(x_exact - X_DISCONTINUITY, T_FINAL)

    figure, axes = plt.subplots(2, 2, figsize=(13, 8.5), constrained_layout=True)
    axes = axes.ravel()
    print(f"Sod, N={N_CELLS}, CFL={CFL}, {nt} steps of dt={dt:.2e}\n")
    print(f"{'flux':10s}{'L1(rho)':>12}{'L1(u)':>12}{'L1(p)':>12}{'mass':>14}")
    for name, flux in fluxes.items():
        scheme, final = solve(flux, dt, nt)
        centers = np.asarray(
            jax.vmap(scheme.variables.mesh.cell_centroid)(jnp.arange(N_CELLS))
        )[:, 0]
        rho, u, p = primitives(final)
        rho_ref, u_ref, p_ref = exact(centers - X_DISCONTINUITY, T_FINAL)
        errors = [
            np.mean(np.abs(a - b)) for a, b in ((rho, rho_ref), (u, u_ref), (p, p_ref))
        ]
        mass = float(jnp.sum(final[:, 0]) * h)
        print(
            f"{name:10s}{errors[0]:>12.3e}{errors[1]:>12.3e}{errors[2]:>12.3e}"
            f"{mass:>14.10f}"
        )
        # Full domain for the three variables, then the density again zoomed
        # on the contact and the shock -- the only place the fluxes differ.
        for axis, field in zip(axes, (rho, u, p, rho)):
            axis.step(centers, field, where="mid", lw=1.6, label=name)

    for axis, field, label in zip(
        axes, (rho_exact, u_exact, p_exact, rho_exact), ("rho", "u", "p", "rho")
    ):
        axis.plot(x_exact, field, "k--", lw=1.2, label="exact")
        axis.set(xlabel="x", ylabel=label)
        axis.grid(alpha=0.25)
    axes[0].legend()
    axes[3].set_xlim(0.55, 0.95)
    axes[3].set_ylim(0.1, 0.5)
    axes[3].set_title("density, zoomed on the contact (0.69) and the shock (0.85)")
    figure.suptitle(f"Sod shock tube, FV P0 + explicit Euler, N={N_CELLS}, t={T_FINAL}")
    print("\nRead the density at the contact: HLLC and Roe keep it sharp, the two-wave")
    print("solvers (Rusanov, HLL) smear it. All four capture the shock.")
    plt.show()