"""Upwind P0 finite-volume transport of a Gaussian over one quarter turn.

The angular velocity is ``2 pi`` and ``T=0.25``: the initial bump therefore
moves from ``(0.45, 0)`` to ``(0, 0.45)``.  The comparison is deliberately
between two different meshes of the disk: a freely recombined Gmsh quadrilateral
mesh (one genuine unstructured mesh, no patches) and the five-patch O-grid.
Both use cell averages and explicit Euler, hence the scheme is formally first
order.  The error is always measured against the exact rotated Gaussian, not
against the initial state.
"""

import math
import time

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.cm import ScalarMappable
from matplotlib.collections import PolyCollection
from matplotlib.colors import Normalize

from scimba_jax.linear_approximation.finite_volume import (
    BlockFiniteVolumeScheme,
    FiniteVolumeScheme,
    TimeDiscreteFVscheme,
)
from scimba_jax.linear_approximation.finite_volume.hyperbolic import UpwindFlux
from scimba_jax.linear_approximation.meshes.block_structured_mesh import (
    BlockStructuredMesh,
)
from scimba_jax.linear_approximation.meshes.unstructured_mesh import (
    UnstructuredMesh,
)
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.variables.variables_fv import VariablesFV
from scimba_jax.mapping.macro_mesh import (
    macro_mesh_from_curve,
    macro_mesh_ogrid_disk,
)
from scimba_jax.nonlinear_approximation.model_class.funcparam_vectorial import (
    ParamVecFunction,
)
from scimba_jax.physical_models.abstract_conservative_pde import (
    AbstractConservativePDE,
)
from scimba_jax.physical_models.abstract_physical_conservative_model import (
    AbstractPhysicalConservativeModel,
)
from scimba_jax.time_discrete.butcher_tableau import build_explicit_euler_tableau

T_FINAL = 0.25
ANGULAR_SPEED = 2.0 * math.pi
# For forward Euler/upwind, an unnecessarily tiny CFL *increases* numerical
# viscosity.  ``time_step_count`` below enforces this as a true local FV CFL,
# not through a proxy based on the smallest cell extent.
CFL = 0.80
# A block mesh refines the cells but keeps its macro geometry.  A high-order
# O-grid is therefore essential here: with quadratic arcs it converges to a
# fixed approximate disk rather than to the radius-one disk of the PDE.  The
# same degree is requested from Gmsh, so the boundary error is well below the
# P0 discretisation error on the displayed levels.
GEOMETRY_ORDER = 8
QUADRATURE_ORDER = 10
# A quarter turn keeps the transport visible while reducing the accumulated
# P0-upwind diffusion.  Render level 32 rather than the deliberately coarse
# first level: a P0 field on cells of diameter about 0.1 is not representative
# of the converged solution.
LEVELS = (16, 32, 64, 72, 96)
PLOT_LEVEL = 32
# ⚠ Largeur du pic, et ce n'est PAS un paramètre cosmétique : c'est elle qui
# décide si les niveaux affichés sont dans le régime asymptotique. Le taux d'un
# upwind P0 ne suit pas le maillage seul mais le nombre de MAILLES PAR
# ÉCART-TYPE. Mesuré sur les blocs, mêmes maillages, mêmes pas de temps, en ne
# changeant que cette valeur ::
#
#     mailles/sigma      2.8    5.5    11     16.5
#     sigma^2 = 0.02      --    0.55   0.70   0.79
#     sigma^2 = 0.08     0.78   0.87   0.92     --
#
# À 0.02 le niveau le plus fin n'a que 8.3 mailles par sigma et plafonne à 0.80 ;
# à 0.08 les MÊMES maillages rendent 0.92 et montent encore. Élargir coûte moins
# cher que raffiner, et c'est la seule des deux façons d'atteindre l'ordre 1 ici
# sans allonger la table.
PULSE_VARIANCE = 0.08
PULSE_CENTRE = 0.45


class RotationTransport(AbstractConservativePDE):
    """``u_t + div(b u) = 0`` with ``b=2 pi (-y, x)``."""

    def __init__(self):
        super().__init__(dim=2)

    def construct_F(self, u):  # noqa: N802
        return ParamVecFunction(
            2,
            u.dims,
            lambda scheme, x: ANGULAR_SPEED
            * jnp.array([-x[1], x[0]])
            * u(scheme, x)[0],
            f_type=u.f_type,
        )

    def construct_A(self, u):  # noqa: N802
        return None

    def construct_R(self, u):  # noqa: N802
        return None


def gaussian(x):
    """A bump initially on the right side of the disk.

    ⚠ Elle ne s'annule pas sur le cercle, et ça reste exact : les
    caractéristiques de la rotation SONT les cercles, donc ``b.n = 0`` partout
    au bord et rien n'y entre ni n'en sort quel que soit le profil. La solution
    exacte est la gaussienne tournée, y compris là où elle est encore non nulle
    au bord.

    Args:
        x: Un point physique.

    Returns:
        La valeur du pic, ``(1,)``.
    """
    return jnp.array(
        [jnp.exp(-((x[0] - PULSE_CENTRE) ** 2 + x[1] ** 2) / PULSE_VARIANCE)]
    )


def disk_boundary(t):
    """Unit circle parameterisation passed to the Gmsh mesh generator."""
    angle = 2.0 * np.pi * t
    return np.cos(angle), np.sin(angle)


def exact_solution(x, time):
    """Gaussian carried by the backward characteristic of the rotation."""
    angle = ANGULAR_SPEED * time
    cosine, sine = jnp.cos(angle), jnp.sin(angle)
    foot = jnp.array([cosine * x[0] + sine * x[1], -sine * x[0] + cosine * x[1]])
    return gaussian(foot)


def normal_speed(_model, context):
    """Normal characteristic speed used by the linear upwind flux."""
    x = context.point
    return jnp.dot(ANGULAR_SPEED * jnp.array([-x[1], x[0]]), context.normal)


def patch_scheme(mesh, flux):
    return FiniteVolumeScheme(
        AbstractPhysicalConservativeModel(RotationTransport()),
        VariablesFV(mesh, projection="average"),
        flux,
        assemble_second_order=False,
        assemble_reaction=False,
        assemble_source=False,
    )


@jax.jit
def local_outflow_rates(mesh):
    """FV outflow rates ``sum_f int (b.n)^+ / |K|`` on one mesh.

    This is the monotonicity limit for scalar linear upwind/Euler.  It uses the
    same physical face quadrature, curved normals, and cell volumes as the
    residual; in particular it is not a guessed ``omega * h_min`` estimate.
    Patch-boundary faces are intentionally included: on a block interface
    they bound the outgoing contribution for that patch cell, and on the disk
    boundary the exact value is zero because the rotation is tangent.
    """
    faces = jnp.arange(mesh.n_faces)

    def one_face(face):
        # ⚠ UNE passe. La version precedente appelait
        # `_local_weights_points_interface` dans DEUX `vmap` separes -- un pour
        # les poids, un pour les points -- donc toute la geometrie de face
        # etait calculee deux fois.
        weights, points = mesh._local_weights_points_interface(face)
        normals = mesh._local_unit_normals_interface(face)
        velocity = ANGULAR_SPEED * jnp.stack([-points[..., 1], points[..., 0]], axis=-1)
        normal_velocity = jnp.sum(velocity * normals, axis=-1)
        return (
            jnp.sum(weights * jnp.maximum(normal_velocity, 0.0)),
            jnp.sum(weights * jnp.maximum(-normal_velocity, 0.0)),
        )

    outflow_left, outflow_right = jax.vmap(one_face)(faces)
    left, right = mesh._face_fidx_to_neighbors_fidx(faces)
    rates = jnp.zeros(mesh.n_cells_total)
    # ⚠ `where` et non un masque booleen : un masque exige des valeurs
    # concretes, donc interdit le `jit` -- et c'est le `jit` qui fait passer ce
    # calcul de 9.45 s a 1.91 s sur le niveau le plus fin.
    rates = rates.at[jnp.where(left >= 0, left, 0)].add(
        jnp.where(left >= 0, outflow_left, 0.0)
    )
    rates = rates.at[jnp.where(right >= 0, right, 0)].add(
        jnp.where(right >= 0, outflow_right, 0.0)
    )
    return rates / mesh.cell_measures()


def time_step_count(meshes):
    """Number of Euler steps for the requested *physical* local CFL."""
    rates = np.concatenate([np.asarray(local_outflow_rates(mesh)) for mesh in meshes])
    max_rate = float(np.max(rates))
    nt = math.ceil(T_FINAL * max_rate / CFL)
    local_cfl = T_FINAL * rates / nt
    return nt, max_rate, tuple(np.quantile(local_cfl, [0.5, 0.9, 1.0]))


def characteristic_mesh_size(mesh):
    """``h`` MOYEN : ``(|Omega| / n_mailles)^(1/dim)``.

    ⚠ Et non le diamètre MAXIMAL, qui est ce que ce fichier utilisait. L'erreur
    mesurée est une intégrale L1 sur TOUTES les mailles, donc l'échelle qui la
    gouverne est la taille moyenne ; le max est une statistique d'aberrant, que
    Gmsh fait bouger d'un niveau à l'autre sans rapport avec le raffinement
    réel. Constaté sur la table : entre les niveaux 64 et 72 le max ne bougeait
    que de 9 % alors que le nombre de mailles faisait x1.36, et le taux affiché
    sortait à **1.4954** -- impossible pour un schéma d'ordre 1. Avec la moyenne
    il vaut 0.88.

    Le maximum reste imprimé à côté, parce que l'écart entre les deux DIT la
    non-uniformité du maillage (facteur ~2 sur ces maillages Gmsh) et que la
    perdre serait perdre une information.

    Args:
        mesh: Un maillage.

    Returns:
        La taille moyenne de maille.
    """
    volume = float(jnp.sum(mesh.cell_measures()))
    return (volume / int(mesh.n_cells_total)) ** (1.0 / mesh.dim)


def largest_cell_diameter(mesh):
    """Le diamètre de la plus grosse maille, imprimé comme diagnostic.

    Args:
        mesh: Un maillage.

    Returns:
        Le plus grand diamètre de cellule.
    """
    extents = jax.vmap(mesh.cell_extent)(jnp.arange(mesh.n_cells_total))
    return float(jnp.max(jnp.linalg.norm(extents, axis=1)))


def cell_average_centres(mesh):
    """Physical barycentres, consistent with the P0 cell-average DOFs."""

    def one(cell):
        weights, points = mesh._local_weights_points(cell)
        return jnp.sum(weights[:, None] * points, axis=0) / jnp.sum(weights)

    return jax.vmap(one)(jnp.arange(mesh.n_cells_total))


def flatten_state(meshes, state, *, with_centres=True):
    """Cell centres, values and volumes from either one mesh or block patches."""
    if not isinstance(state, tuple):
        meshes, state = (meshes[0],), (state,)
    centres = []
    values = []
    volumes = []
    for mesh, dofsl in zip(meshes, state):
        if with_centres:
            centres.append(cell_average_centres(mesh))
        values.append(dofsl[:, 0])
        volumes.append(mesh.cell_measures())
    return (
        jnp.concatenate(centres) if with_centres else None,
        jnp.concatenate(values),
        jnp.concatenate(volumes),
    )


def state_tuple(state):
    """Normalise a one-mesh state and a block state to patch tuples."""
    return state if isinstance(state, tuple) else (state,)


def project_exact(meshes, time):
    """Exact cell averages at ``time`` in the representation used by FV."""
    exact = jax.vmap(lambda x: exact_solution(x, time))
    return tuple(
        VariablesFV(mesh, projection="average").projector(exact) for mesh in meshes
    )


def run(spatial_scheme, meshes):
    # ⚠ La moyenne sur TOUS les patchs, pas le max des moyennes par patch : le
    # domaine est un, l'erreur L1 y est une seule intégrale, et sur l'O-grid le
    # patch central et les pétales n'ont pas la même finesse.
    total_volume = sum(float(jnp.sum(mesh.cell_measures())) for mesh in meshes)
    total_cells = sum(int(mesh.n_cells_total) for mesh in meshes)
    h = (total_volume / total_cells) ** (1.0 / meshes[0].dim)
    h_max = max(largest_cell_diameter(mesh) for mesh in meshes)
    # ⚠ Le chrono démarre ICI, donc HORS construction de maillage : celle-ci a
    # eu lieu chez l'appelant (Gmsh pour l'un, les patchs pour l'autre) et elle
    # n'a rien à voir avec ce que coûte le schéma. Ce qui est mesuré est le pas
    # de temps et sa préparation.
    setup_start = time.perf_counter()
    nt, max_outflow, cfl_quantiles = time_step_count(meshes)
    time_scheme = TimeDiscreteFVscheme(
        spatial_scheme,
        build_explicit_euler_tableau(),
        dt=T_FINAL / nt,
    )
    initial = jax.block_until_ready(time_scheme.initialize(gaussian))
    setup_seconds = time.perf_counter() - setup_start
    # ⚠ `block_until_ready`, et ce n'est pas une formalité : JAX dispatche de
    # façon asynchrone, donc un chrono arrêté ici mesurerait la MISE EN FILE et
    # le travail atterrirait dans ce qu'on mesure ensuite.
    solve_start = time.perf_counter()
    final = jax.block_until_ready(time_scheme.solve_final(initial, t0=0.0, nt=nt))
    solve_seconds = time.perf_counter() - solve_start
    centres, initial_values, volumes = flatten_state(meshes, initial)
    _, final_values, _ = flatten_state(meshes, final, with_centres=False)
    exact_state = project_exact(meshes, T_FINAL)
    _, exact_values, _ = flatten_state(meshes, exact_state, with_centres=False)
    l1_error = float(jnp.sum(jnp.abs(final_values - exact_values) * volumes))
    mass_initial = jnp.sum(initial_values * volumes)
    mass_final = jnp.sum(final_values * volumes)
    radius_squared = jnp.sum(centres**2, axis=1)
    radial_second_initial = jnp.sum(radius_squared * initial_values * volumes)
    radial_second_final = jnp.sum(radius_squared * final_values * volumes)
    velocity = ANGULAR_SPEED * jnp.stack([-centres[:, 1], centres[:, 0]], axis=1)
    return {
        "h": h,
        "h_max": h_max,
        "n_cells": total_cells,
        "nt": nt,
        "setup_seconds": setup_seconds,
        # ⚠ Compilation COMPRISE : chaque niveau a des formes nouvelles, donc
        # une trace neuve, et la mesurer à part demanderait de refaire le solve
        # -- ce qui doublerait le coût de toute la table pour une colonne.
        "solve_seconds": solve_seconds,
        "local_cfl_quantiles": cfl_quantiles,
        "max_local_cfl": T_FINAL * max_outflow / nt,
        "meshes": tuple(meshes),
        "initial_state": state_tuple(initial),
        "final_state": state_tuple(final),
        "exact_state": exact_state,
        "centres": centres,
        "initial": initial_values,
        "final": final_values,
        "exact": exact_values,
        "error": l1_error,
        "area": float(jnp.sum(volumes)),
        "relative_mass_change": float((mass_final - mass_initial) / mass_initial),
        "mean_radius_squared_initial": float(radial_second_initial / mass_initial),
        "mean_radius_squared_final": float(radial_second_final / mass_final),
        "max_radial_velocity": float(
            jnp.max(jnp.abs(jnp.sum(velocity * centres, axis=1)))
        ),
    }


def unstructured_run(level):
    """Freely recombined Gmsh mesh: no macro-patch/O-grid topology."""
    macro = macro_mesh_from_curve(
        disk_boundary,
        order=GEOMETRY_ORDER,
        mesh_size=1.0 / level,
    )
    mesh = UnstructuredMesh.from_macro_mesh(
        macro, UnitSquareTensorized(dim=2, order=QUADRATURE_ORDER)
    )
    return run(patch_scheme(mesh, UpwindFlux(normal_speed)), [mesh])


def block_run(level):
    block_mesh = BlockStructuredMesh(
        macro_mesh_ogrid_disk(n=1, order=GEOMETRY_ORDER),
        n_cells=(level, level),
        ref_quad=UnitSquareTensorized(dim=2, order=QUADRATURE_ORDER),
    )
    flux = UpwindFlux(normal_speed)
    patches = [patch_scheme(mesh, flux) for mesh in block_mesh.meshes]
    # ⚠ Pas de `flux=` : le conteneur prend celui des patchs. Le repasser
    # ici mettrait le même objet deux fois dans le pytree.
    return run(BlockFiniteVolumeScheme(block_mesh, patches), block_mesh.meshes)


def physical_cell_outlines(mesh, samples=5):
    """Mapped, closed outlines of every cell, evaluated in one JAX batch.

    P0 needs one colour per cell, not a dense grid of coloured subcells.  The
    former renderer made one ``pcolormesh`` (36 artists) per FV cell, which
    becomes unusable on the Gmsh refinement levels.  Sampling only the curved
    boundary lets one :class:`PolyCollection` render the whole mesh.
    """
    coordinate = jnp.linspace(0.0, 1.0, samples)
    bottom = jnp.stack([coordinate, jnp.zeros_like(coordinate)], axis=-1)
    right = jnp.stack([jnp.ones_like(coordinate[1:]), coordinate[1:]], axis=-1)
    top = jnp.stack([coordinate[-2::-1], jnp.ones_like(coordinate[:-1])], axis=-1)
    left = jnp.stack(
        [jnp.zeros_like(coordinate[-2:0:-1]), coordinate[-2:0:-1]], axis=-1
    )
    reference = jnp.concatenate([bottom, right, top, left], axis=0)
    cells = jnp.arange(mesh.n_cells_total)
    if isinstance(mesh, UnstructuredMesh):
        outlines = jax.vmap(lambda cell: mesh._unit_hypercube_to_cell(cell, reference))(
            cells
        )
    else:
        logical = jax.vmap(lambda cell: mesh._unit_hypercube_to_cell(cell, reference))(
            cells
        )
        outlines = mesh.mapping.local_mapping(logical.reshape(-1, 2)).reshape(
            logical.shape
        )
    return np.asarray(outlines)


def plot_solution(axis, meshes, state, title, norm, cmap="turbo", contours=8):
    """Le champ P0 SUR ses mailles : cellules remplies, arêtes par-dessus.

    ⚠ Un champ P0 EST une valeur par maille, pas un nuage de points : le
    représenter par un marqueur au barycentre demande de deviner une taille de
    marqueur, qui ne suit ni la forme des cellules ni leur étirement sous le
    mapping -- sur l'O-grid, où les cellules du centre et celles des pétales
    n'ont ni la même taille ni la même forme, les carrés se chevauchaient au
    centre et laissaient du fond visible au bord.

    Une :class:`PolyCollection` dit la chose telle qu'elle est, et le même
    artiste porte la valeur (``array``) et le maillage (``edgecolors``) -- d'où
    la disparition de la figure de maillage séparée : elle était le même dessin,
    à moitié.

    Args:
        axis: L'axe matplotlib.
        meshes: Les maillages, un par patch.
        state: Les ddl, un tableau par patch.
        title: Titre du panneau.
        norm: Normalisation des couleurs, partagée entre les panneaux.
        cmap: Palette. ``turbo`` par défaut : elle sépare mieux les faibles
            valeurs que ``viridis``, et c'est là que se lit la diffusion
            numérique -- la traînée derrière le pic vaut quelques pour cent du
            maximum. ⚠ Elle n'est pas perceptuellement uniforme, donc ses bandes
            peuvent suggérer des marches qui n'existent pas ; les lignes de
            niveau sont là pour trancher.
        contours: Nombre de lignes de niveau, ``0`` pour aucune.

    Returns:
        La collection ajoutée.
    """
    polygons = np.concatenate([physical_cell_outlines(mesh) for mesh in meshes])
    values = np.concatenate([np.asarray(dofsl[:, 0]) for dofsl in state])
    # ⚠ L'épaisseur des arêtes suit le RAFFINEMENT : à 0.25 partout, un maillage
    # à 46 000 cellules n'est plus un maillage mais un aplat gris qui cache le
    # champ. En dessous de 0.05 matplotlib ne dessine plus rien de lisible, donc
    # les arêtes s'effacent d'elles-mêmes quand elles cesseraient d'informer.
    width = float(np.clip(400.0 / len(values), 0.0, 0.4))
    collection = PolyCollection(
        list(polygons),
        array=values,
        cmap=cmap,
        norm=norm,
        edgecolors="0.25" if width > 0.05 else "none",
        linewidths=width,
        rasterized=True,
    )
    axis.add_collection(collection)
    if contours:
        # ⚠ Sur les BARYCENTRES, donc sur une interpolation du champ P0 et non
        # sur le champ lui-même, qui est constant par maille et dont les vraies
        # lignes de niveau seraient les arêtes. C'est un outil de LECTURE : il
        # dit où passe le front et s'il est resté rond, ce qu'un aplat de
        # couleur ne permet pas de juger à l'oeil.
        centres = np.concatenate(
            [np.asarray(cell_average_centres(mesh)) for mesh in meshes]
        )
        axis.tricontour(
            centres[:, 0],
            centres[:, 1],
            values,
            levels=np.linspace(norm.vmin, norm.vmax, contours + 2)[1:-1],
            colors="black",
            linewidths=0.6,
            alpha=0.55,
        )
    axis.add_patch(
        plt.Circle((0.0, 0.0), 1.0, fill=False, color="black", ls="--", lw=1.0)
    )
    axis.set(
        title=title,
        xlabel="x",
        ylabel="y",
        xlim=(-1.05, 1.05),
        ylim=(-1.05, 1.05),
    )
    axis.set_aspect("equal")
    return collection


def observed_rates(results):
    """Consecutive L1 rates based on the actual characteristic cell sizes."""
    hs = jnp.asarray([result["h"] for result in results])
    errors = jnp.asarray([result["error"] for result in results])
    return jnp.log(errors[:-1] / errors[1:]) / jnp.log(hs[:-1] / hs[1:])


HEADER = (
    f"{'representation':18s} {'niveau':>6s} {'cellules':>9s} {'h_moy':>10s} "
    f"{'h_max':>10s} {'pas':>5s} {'CFL 50/90/max':>17s} {'erreur L1':>11s} "
    f"{'taux':>5s} {'setup s':>8s} {'solve s':>8s}"
)


def report_row(name, level, result, previous):
    """Une ligne, imprimée dès que le run est fini.

    Args:
        name: La représentation.
        level: Le niveau de raffinement.
        result: Ce que :func:`run` a rendu.
        previous: Le résultat précédent, ou ``None`` pour la première ligne.

    Returns:
        Rien ; imprime.
    """
    rate = ""
    if previous is not None:
        rate = "%.2f" % (
            math.log(previous["error"] / result["error"])
            / math.log(previous["h"] / result["h"])
        )
    cfl = (
        f"{result['local_cfl_quantiles'][0]:.2f}/"
        f"{result['local_cfl_quantiles'][1]:.2f}/"
        f"{result['max_local_cfl']:.2f}"
    )
    print(
        f"{name:18s} {level:6d} {result['n_cells']:9d} {result['h']:10.4e} "
        f"{result['h_max']:10.4e} {result['nt']:5d} {cfl:>17s} "
        f"{result['error']:11.4e} {rate:>5s} "
        f"{result['setup_seconds']:8.2f} {result['solve_seconds']:8.2f}",
        flush=True,
    )


if __name__ == "__main__":
    # ⚠ Imprimé AU FIL DE L'EAU, pas à la fin : la table complète est une
    # dizaine de solves dont le dernier est le plus long, et tout accumuler
    # laissait l'écran vide pendant toute la campagne -- au point qu'un run qui
    # divergeait ressemblait à un run qui tournait.
    print("\nTransport rotationnel 2D — FV P0 + Euler explicite, T=0.25")
    print(
        f"pic : sigma^2={PULSE_VARIANCE}, centre ({PULSE_CENTRE}, 0)"
        "   |   taux calcules sur h_moy   |   solve = compilation comprise,"
        " hors construction de maillage"
    )
    print(HEADER)
    print("-" * len(HEADER))
    families = {}
    for name, runner in (
        ("non structuré", unstructured_run),
        ("blocs structurés", block_run),
    ):
        results = []
        for level in LEVELS:
            result = runner(level)
            report_row(name, level, result, results[-1] if results else None)
            results.append(result)
        families[name] = results
    unstructured = families["non structuré"]
    block = families["blocs structurés"]
    for name, result in (
        ("non structuré", unstructured[0]),
        ("blocs structurés", block[0]),
    ):
        print(
            f"{name:18s} aire={result['area']:.12f}  "
            f"Δmasse/m={result['relative_mass_change']:+.2e}  "
            f"<r²>: {result['mean_radius_squared_initial']:.6f} -> "
            f"{result['mean_radius_squared_final']:.6f}  "
            f"max |b·x|={result['max_radial_velocity']:.1e}"
        )

    plotted_unstructured = unstructured[LEVELS.index(PLOT_LEVEL)]
    plotted_block = block[LEVELS.index(PLOT_LEVEL)]
    # ⚠ Plus de figure de maillage a part : chaque panneau ci-dessous porte le
    # maillage sous le champ, ce qui est la seule facon de voir OU la diffusion
    # numerique agit -- elle suit les cellules, pas la geometrie.
    figure, axes = plt.subplots(2, 3, figsize=(13, 8), constrained_layout=True)
    all_values = jnp.concatenate(
        [
            result[key]
            for result in (plotted_unstructured, plotted_block)
            for key in ("initial", "exact", "final")
        ]
    )
    value_norm = Normalize(
        vmin=float(jnp.min(all_values)), vmax=float(jnp.max(all_values))
    )
    all_errors = jnp.concatenate(
        [
            jnp.abs(result["final"] - result["exact"])
            for result in (plotted_unstructured, plotted_block)
        ]
    )
    error_norm = Normalize(vmin=0.0, vmax=float(jnp.max(all_errors)))
    for row, (name, result) in enumerate(
        (
            ("maillage non structuré", plotted_unstructured),
            ("blocs structurés", plotted_block),
        )
    ):
        plot_solution(
            axes[row, 0],
            result["meshes"],
            result["initial_state"],
            f"{name} : t=0",
            value_norm,
        )
        plot_solution(
            axes[row, 1],
            result["meshes"],
            result["final_state"],
            f"{name} : T={T_FINAL}",
            value_norm,
        )
        plot_solution(
            axes[row, 2],
            result["meshes"],
            tuple(
                jnp.abs(final - exact)
                for exact, final in zip(result["exact_state"], result["final_state"])
            ),
            f"{name} : |u_h(T)-u_exact(T)|",
            error_norm,
            cmap="magma",
        )
    figure.colorbar(
        ScalarMappable(norm=value_norm, cmap="turbo"),
        ax=axes[:, :2],
        shrink=0.82,
        label="u",
    )
    figure.colorbar(
        ScalarMappable(norm=error_norm, cmap="magma"),
        ax=axes[:, 2],
        shrink=0.82,
        label=r"$|u_h(T)-u_{exact}(T)|$",
    )

    convergence, axis = plt.subplots(figsize=(6.5, 4.5), constrained_layout=True)
    for name, results, marker in (
        ("non structuré", unstructured, "o"),
        ("blocs structurés", block, "s"),
    ):
        axis.loglog(
            [result["h"] for result in results],
            [result["error"] for result in results],
            f"-{marker}",
            lw=2,
            label=name,
        )
        rates = np.asarray(observed_rates(results))
        for h, error, rate in zip(
            [result["h"] for result in results[1:]],
            [result["error"] for result in results[1:]],
            rates,
        ):
            axis.annotate(
                f"p={rate:.2f}", (h, error), xytext=(4, 5), textcoords="offset points"
            )
    axis.set(
        xlabel="h moyen = (|Ω| / n_mailles)^(1/2)",
        ylabel="erreur L1 à T=0.25",
        title="Raffinement : taux L1 mesurés (FV P0 upwind)",
    )
    axis.grid(which="both", alpha=0.25)
    axis.legend()
    plt.show()