"""Cylindrical explosion in 2D: the same four Euler fluxes as the 1D shock tube.

    W_t + div F(W) = 0  on [-1, 1]^2,   gamma = 1.4,   t = 0.25
    (rho, u, p) = (1, 0, 1) for r < 0.4,  (0.125, 0, 0.1) outside

Toro, *Riemann Solvers and Numerical Methods for Fluid Dynamics*, 3rd ed.,
sec. 17.1: Sod's data on a disc. A circular shock and contact run outward, a
rarefaction runs inward and, after it focuses at the centre, an inward shock
forms. The waves have not reached the box by t = 0.25, so the ambient state
is an exact ghost state on the four sides.

The point of the file is that NOTHING in the fluxes is 2D-specific. They are
the objects of ``euler_sod_1d.py`` (``finite_volume.hyperbolic.euler``),
written in the frame of each face normal -- ``u_n = v . n``, tangential
velocity carried with the contact -- so the same class serves ``n = (1,)``
and ``n = (n_x, n_y)``. A flux that were not rotationally invariant would
show up here as an explosion that is not round.

Reference: there is no closed form, but the problem is one-dimensional in
``r``. The 1D radial Euler system with the geometric source

    W_t + F(W)_r = -(1/r) (rho u, rho u^2, u (E + p))

is solved on a fine mesh (HLLC, a reflecting ``Wall`` at r = 0, the ambient
state at r = 1), and the 2D cells are compared to it along the x-axis and the
diagonal: two cuts at 45 degrees of each other, which agree only if the
scheme is isotropic. First order in space and time, as in 1D.
"""

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.nonlinear_approximation.model_class.funcparam_vectorial import (
    ParamVecFunction,
)
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_primitives,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet, Wall
from scimba_jax.time_discrete.butcher_tableau import build_explicit_euler_tableau

GAMMA = 1.4
N_CELLS = 200  # per direction
N_RADIAL = 2000  # the 1D reference
CFL = 0.8
T_FINAL = 0.25
# Bound on |v . n| + c over the WHOLE run, not on the initial data: Sod's star
# state moves at u* = 0.93 with c* = 1.0, so 1.93 -- the inner state's c = 1.18
# would be wrong by a factor 1.6. The step is CFL h / (dim S): in 2D the
# Rusanov dissipation of the four faces sums to 2 alpha / h on the diagonal,
# so its explicit-Euler bound is dt <= h / (2 alpha), i.e. CFL <= 1 here.
SPEED_BOUND = 2.0
RADIUS = 0.4
INSIDE = (1.0, 0.0, 1.0)  # (rho, u, p)
OUTSIDE = (0.125, 0.0, 0.1)


def conserved(state, dim):
    rho, _, p = state
    return euler_conserved(rho, jnp.zeros(dim), p, GAMMA)


def initial_condition(x):
    """Sod's data on a disc; cut cells take the cell average of the projection."""
    dim = x.shape[0]
    inside = jnp.dot(x, x) < RADIUS * RADIUS
    return jnp.where(inside, conserved(INSIDE, dim), conserved(OUTSIDE, dim))


class RadialEulerPDE(EulerPDE):
    """1D Euler with the cylindrical geometric source, ``x = r``.

    ``div F + R = 0`` with ``R = (1/r) (rho u, rho u^2, u (E + p))``: the
    reaction slot of the conservative PDE, since it depends on the state.
    """

    def __init__(self):
        super().__init__(dim=1, gamma=GAMMA)

    def construct_R(self, u):  # noqa: N802
        gamma = self.gamma

        def geometric(scheme, x):
            w = u(scheme, x)
            rho, velocity, pressure, _ = euler_primitives(w, gamma)
            u_r = velocity[0]
            radial_flux = jnp.array(
                [rho * u_r, rho * u_r * u_r, u_r * (w[-1] + pressure)]
            )
            return radial_flux / x[0]

        return ParamVecFunction(3, u.dims, geometric, f_type=u.f_type)


def time_step(h, dim):
    """``CFL h / (dim S_max)``, see ``SPEED_BOUND``."""
    return CFL * h / (dim * SPEED_BOUND)


def make_2d_scheme(flux, dt):
    mesh = cartesian_mesh([N_CELLS, N_CELLS], quad_order=2, bounds=[(-1.0, 1.0)] * 2)
    model = AbstractPhysicalConservativeModel(EulerPDE(dim=2, gamma=GAMMA))
    for side in ("west", "east", "south", "north"):
        model.add_boundary_condition(side, Dirichlet(lambda _x: conserved(OUTSIDE, 2)))
    spatial = FiniteVolumeScheme(
        model,
        VariablesFV(mesh, nb_variables=4),
        flux,
        assemble_second_order=False,
        assemble_reaction=False,
        assemble_source=False,
    )
    return TimeDiscreteFVscheme(spatial, build_explicit_euler_tableau(), dt=dt)


def make_radial_scheme(dt):
    mesh = cartesian_mesh([N_RADIAL], quad_order=2)  # r in [0, 1]
    model = AbstractPhysicalConservativeModel(RadialEulerPDE())
    model.add_boundary_condition("west", Wall())  # the axis r = 0
    model.add_boundary_condition("east", Dirichlet(lambda _x: conserved(OUTSIDE, 1)))
    spatial = FiniteVolumeScheme(
        model,
        VariablesFV(mesh, nb_variables=3),
        EulerHLLCFlux(GAMMA),
        assemble_second_order=False,
        assemble_reaction=True,
        assemble_source=False,
    )
    return TimeDiscreteFVscheme(spatial, build_explicit_euler_tableau(), dt=dt)


def run(scheme, dt):
    nt = round(T_FINAL / dt)
    initial = scheme.initialize(initial_condition)
    final = scheme.solve_final(initial, t0=0.0, nt=nt)
    centers = np.asarray(
        jax.vmap(scheme.variables.mesh.cell_centroid)(
            jnp.arange(scheme.variables.mesh.n_cells_total)
        )
    )
    return centers, np.asarray(final)


def primitives(w):
    """``(rho, |v|, p)`` cell by cell; ``|v|`` is the radial speed on a round flow."""
    rho, velocity, pressure, _ = jax.vmap(lambda s: euler_primitives(s, GAMMA))(
        jnp.asarray(w)
    )
    return (
        np.asarray(rho),
        np.asarray(jnp.linalg.norm(velocity, axis=-1)),
        np.asarray(pressure),
    )


if __name__ == "__main__":
    h = 2.0 / N_CELLS
    dt = T_FINAL / round(T_FINAL / time_step(h, 2))
    fluxes = {
        "Rusanov": RusanovFlux(euler_wave_speed(GAMMA)),
        "HLL": HLLFlux(euler_hll_wave_speeds(GAMMA)),
        "HLLC": EulerHLLCFlux(GAMMA),
        "Roe": EulerRoeFlux(GAMMA),
    }

    # The radial reference, on a mesh 10x finer than the 2D one.
    dt_radial = T_FINAL / round(T_FINAL / time_step(1.0 / N_RADIAL, 1))
    r_ref, w_ref = run(make_radial_scheme(dt_radial), dt_radial)
    rho_ref, u_ref, p_ref = primitives(w_ref)
    r_ref = r_ref[:, 0]

    results = {}
    print(
        f"2D: {N_CELLS}x{N_CELLS} cells, {round(T_FINAL / dt)} steps;  reference: {N_RADIAL} radial cells\n"
    )
    print(f"{'flux':10s}{'mass':>14}{'rho min':>10}{'rho max':>10}")
    for name, flux in fluxes.items():
        centers, final = run(make_2d_scheme(flux, dt), dt)
        results[name] = (centers, final)
        mass = float(np.sum(final[:, 0]) * h * h)
        print(
            f"{name:10s}{mass:>14.8f}{final[:, 0].min():>10.4f}{final[:, 0].max():>10.4f}"
        )

    # ── One image per flux ────────────────────────────────────────────────
    figure, axes = plt.subplots(2, 2, figsize=(10, 9.5), constrained_layout=True)
    for axis, (name, (centers, final)) in zip(axes.ravel(), results.items()):
        # Cells are numbered `unravel_index(idx, (n, n))`: x index first.
        image = final[:, 0].reshape(N_CELLS, N_CELLS)
        axis.imshow(
            image.T,
            origin="lower",
            extent=(-1, 1, -1, 1),
            vmin=0.1,
            vmax=1.0,
            cmap="viridis",
        )
        axis.set(title=f"{name}: density at t = {T_FINAL}", xlabel="x", ylabel="y")
    figure.suptitle(
        f"Cylindrical explosion, FV P0 + explicit Euler, {N_CELLS}x{N_CELLS}"
    )

    # ── The cuts: x-axis and diagonal, against the radial reference ───────
    figure, axes = plt.subplots(2, 3, figsize=(15, 8), constrained_layout=True)
    cuts = (
        (
            "x-axis cut (cells with |y| < h/2, x > 0)",
            lambda c: (np.abs(c[:, 1]) < h / 2) & (c[:, 0] > 0),
        ),
        (
            "diagonal cut (cells with x = y > 0)",
            lambda c: (np.abs(c[:, 0] - c[:, 1]) < h / 2) & (c[:, 0] > 0),
        ),
    )
    for row, (title, select) in zip(axes, cuts):
        for name, (centers, final) in results.items():
            mask = select(centers)
            r = np.linalg.norm(centers[mask], axis=1)
            order = np.argsort(r)
            rho, speed, p = primitives(final[mask])
            for axis, field in zip(row, (rho, speed, p)):
                axis.plot(r[order], field[order], lw=1.4, label=name)
        for axis, field, label in zip(
            row, (rho_ref, u_ref, p_ref), ("rho", "|v|", "p")
        ):
            axis.plot(r_ref, field, "k--", lw=1.2, label="1D radial reference")
            axis.set(xlabel="r", ylabel=label, xlim=(0, 1))
            axis.grid(alpha=0.25)
        row[0].set_title(title)
    axes[0, 0].legend()
    figure.suptitle(
        "Cuts of the 2D solution against the 1D radial reference "
        "(shock ~ 0.8, contact ~ 0.6, focused rarefaction near r = 0)"
    )
    print(
        "\nThe two cuts must coincide (isotropy) and follow the radial reference;\n"
        "the contact at r ~ 0.6 is where the two-wave solvers smear more."
    )
    plt.show()