# tests/test_riemann.py
import numpy as np
from riemann import hll_flux


GAMMA = 1.4


def _to_conservative(rho, u, P, gamma):
    """(rho, u, P) -> [rho, rho*u, rho*E] where E = e + u^2/2."""
    return np.array([
        rho,
        rho * u,
        P / (gamma - 1) + 0.5 * rho * u ** 2,
    ])


def _physical_flux(W, gamma):
    """F(W) = [rho*u, rho*u^2 + P, u*(rho*E + P)]."""
    rho = W[0]
    u = W[1] / rho
    E = W[2]
    P = (gamma - 1) * (E - 0.5 * rho * u ** 2)
    return np.array([rho * u, rho * u ** 2 + P, u * (E + P)])


def test_hll_identical_states_returns_physical_flux():
    """
    When W_L == W_R, HLL must return the exact physical flux F(W)
    with zero numerical dissipation (the (W_R - W_L) term vanishes).
    """
    W = _to_conservative(rho=1.2, u=50.0, P=2.5e5, gamma=GAMMA)
    F = hll_flux(W, W, GAMMA)
    expected = _physical_flux(W, GAMMA)
    assert np.allclose(F, expected, rtol=1e-12), f"F={F}, expected={expected}"


def test_hll_equal_pressure_equal_energy_gives_exact_pressure_flux():
    """
    Two stationary states (u=0) with same P but different rho:
    - Both have the same energy density E = P/(gamma-1), so the HLL
      (W_R - W_L)[2] term vanishes -> exact zero energy flux.
    - F_L[1] = F_R[1] = P, and W_L[1] = W_R[1] = 0, so the momentum
      flux is exactly P.
    - The mass flux is NOT exactly zero for HLL (the density jump
      triggers the (W_R - W_L)[0] dissipation term) — this is a known
      HLL limitation for stationary contact discontinuities. We do
      not assert on F[0] here.
    """
    W_L = _to_conservative(rho=10.0, u=0.0, P=1e5, gamma=GAMMA)
    W_R = _to_conservative(rho=1.0,  u=0.0, P=1e5, gamma=GAMMA)
    F = hll_flux(W_L, W_R, GAMMA)
    assert abs(F[1] - 1e5) < 1e-6, f"momentum flux should equal P=1e5, got {F[1]}"
    assert abs(F[2]) < 1e-8, f"energy flux should be exactly 0, got {F[2]}"


def test_hll_sod_shock_tube_directional_fluxes_all_positive():
    """
    Classical Sod initial values:
      left:  (rho, u, P) = (1.0,   0, 1.0)
      right: (rho, u, P) = (0.125, 0, 0.1)
    The pressure gradient drives flow from left to right, so the
    HLL flux at the interface should have all three components
    strictly positive:
      F[0] > 0 : mass flux rightward
      F[1] > 0 : momentum flux rightward
      F[2] > 0 : energy flux rightward
    """
    W_L = _to_conservative(rho=1.0,   u=0.0, P=1.0, gamma=GAMMA)
    W_R = _to_conservative(rho=0.125, u=0.0, P=0.1, gamma=GAMMA)
    F = hll_flux(W_L, W_R, GAMMA)
    assert F[0] > 0, f"expected positive mass flux, got {F[0]}"
    assert F[1] > 0, f"expected positive momentum flux, got {F[1]}"
    assert F[2] > 0, f"expected positive energy flux, got {F[2]}"