r"""Système advection-diffusion multi-espace + r-adaptation par Monge-Ampère.

Reprend ``solve_2d_system_diffusion_advection_multi_space.py`` mais :

  1.  Grille GROSSIÈRE, et coefficients de diffusion DIFFÉRENTS par variable :
      -ε_u Δu + b₁·∇u + c₁·∇v = f     (ε_u GRAND  → u lisse)
      -ε_v Δv + b₂·∇v + c₂·∇u = f     (ε_v PETIT  → transport domine → v raide)

  2.  On construit la densité monitrice  g(y) = 1 + α|∇v_h(y)|²  à partir de la
      solution FE de v, en passant par ``local_variables`` (ParamScalarFunction)
      puis ``.gradient("x")`` (ParamVecFunction), le tout encapsulé dans un
      callable  g : y ↦ scalaire  (comme demandé).

  3.  On résout un Monge-Ampère (PINN, ICNN convexe, linéarisation de Picard)
      g(∇u) det(∇²u) = f=1, dont la carte de transport  T = ∇u  pousse la
      densité uniforme vers g : un maillage uniforme transporté par T se
      concentre là où |∇v|² est grand. Comme dans
      ``pinns/optimal_transport/density_transport_2d.py``, on apprend ENSUITE
      la carte inverse T⁻¹, puis on assemble ``InvertibleNetBasedICNN`` (T et
      T⁻¹ exacts par Newton).

  4.  On relance le problème d'origine, mais le maillage de v utilise le
      ``Mapping`` adapté (T) — donc le maillage de v s'adapte aux zones raides,
      tandis que le maillage de u reste cartésien (identité).
"""

import time

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

from scimba_jax.domains.meshless_domains.domains_2d import Square2D
from scimba_jax.linear_approximation.basis.analytic_bases import local_lagrange_basis
from scimba_jax.linear_approximation.basis.general_bases import AnalyticBasis
from scimba_jax.linear_approximation.galerkin.fem.elliptic_fe_scheme import (
    EllipticFEscheme,
)
from scimba_jax.linear_approximation.galerkin.fem.elliptic_fe_scheme_multi_space import (
    EllipticFEschemeMultipleSpaces,
)
from scimba_jax.linear_approximation.meshes.mesh import Mesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.linear_approximation.variables.variables_fe import VariablesFE
from scimba_jax.mapping.mapping import InvertibleFunction, Mapping
from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import (
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
    DomainSampler,
    TensorizedSampler,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo_parameters import (
    UniformParametricSampler,
)
from scimba_jax.nonlinear_approximation.model_class.funcparam_vectorial import (
    ParamVecFunction,
)
from scimba_jax.nonlinear_approximation.networks.icnn import ICNN
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.networks.structure_preserving_nets.invertiblenet_based_icnn import (
    InvertibleNetBasedICNN,
)
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.abstract_physical_weak_model import (
    AbstractPhysicalWeakModel,
)
from scimba_jax.physical_models.abstract_weak_form import AbstractWeakForm
from scimba_jax.physical_models.elliptic_pde.monge_ampere import (
    InverseFlowMongeAmpere2D,
    MongeAmperePicard2D,
)
from scimba_jax.physical_models.function_approximator.function_approximator import (
    FunctionApproximator,
)
from scimba_jax.physical_models.weak_boundary_conditions import Dirichlet

# ── Paramètres ────────────────────────────────────────────────────────────────

physical_dim = 2
quad_order = 3
poly_order = 1

# Grille grossière (r-adaptation : on veut peu de mailles → effet du mapping visible).
n_cells_u = 30
n_cells_v = 30

# Diffusion par variable : u lisse (ε grand), v dominé par le transport (ε petit).
eps_u = 0.3
eps_v = 0.02

sigma_gauss = 0.07
x0 = jnp.array([0.5, 0.5])


def gaussian_source(x):
    return jnp.exp(-jnp.sum((x - x0) ** 2) / (2.0 * sigma_gauss**2)) / (
        2.0 * jnp.pi * sigma_gauss**2
    )


# ── Forme faible : diffusion anisotrope par variable (A1 pour u, A2 pour v) ────


class SystemDiffAdvWeakForm(AbstractWeakForm):
    def __init__(self, dim, eps1, eps2, b1, b2, c1, c2, f):
        super().__init__(dim=dim)
        self.A1 = lambda x: eps1 * jnp.eye(dim)
        self.A2 = lambda x: eps2 * jnp.eye(dim)
        self.b1 = b1
        self.b2 = b2
        self.c1 = c1
        self.c2 = c2
        self.f = f

    def bilinear_form(
        self, u: ParamVecFunction, v: ParamVecFunction
    ) -> ParamVecFunction:
        u_1, u_2 = u  # un argument par espace
        v_1, v_2 = v

        grad_u1 = u_1.gradient("x")
        grad_u2 = u_2.gradient("x")
        grad_v1 = v_1.gradient("x")
        grad_v2 = v_2.gradient("x")

        fields = self.get_fields()
        A1 = fields["A1"]
        A2 = fields["A2"]
        b1 = fields["b1"]
        b2 = fields["b2"]
        c1 = fields["c1"]
        c2 = fields["c2"]

        diff_u = grad_v1.dot(A1 @ grad_u1)
        adv_u = grad_u1.dot(b1) * v_1
        cross_adv_u = grad_u2.dot(c1) * v_1

        diff_v = grad_v2.dot(A2 @ grad_u2)
        adv_v = grad_u2.dot(b2) * v_2
        cross_adv_v = grad_u1.dot(c2) * v_2

        return ParamVecFunction.cat(
            [diff_u + adv_u + cross_adv_u, diff_v + adv_v + cross_adv_v]
        )

    def linear_form(self, v: ParamVecFunction) -> ParamVecFunction:
        v_1, v_2 = v
        f = self.get_fields()["f"]
        return ParamVecFunction.cat([f * v_1, f * v_2])


# ── Advection / couplage (identiques à l'original) ────────────────────────────

b1_vec = jnp.array([1.0, 0.0])
b2_vec = jnp.array([-1.0, 0.0])
c1_vec = jnp.array([0.0, 0.3])
c2_vec = jnp.array([0.0, -0.3])

mapping_id = InvertibleFunction(lambda x: x, lambda y: y)

pde = SystemDiffAdvWeakForm(
    dim=physical_dim,
    eps1=eps_u,
    eps2=eps_v,
    b1=lambda x: b1_vec,
    b2=lambda x: b2_vec,
    c1=lambda x: c1_vec,
    c2=lambda x: c2_vec,
    f=gaussian_source,
)


def make_variables(n_cells, mapping, is_identity=True):
    mesh = Mesh(
        dim=physical_dim,
        n_cells=(n_cells, n_cells),
        ref_quad=UnitSquareTensorized(dim=physical_dim, order=quad_order),
        mapping=mapping,
        is_identity_mapping=is_identity,
    )
    basis = AnalyticBasis(
        nb_basis=(poly_order + 1) ** physical_dim,
        out_dim=1,
        mesh=mesh,
        local_basis=lambda coords, i, m: local_lagrange_basis(
            coords, i, m, order=poly_order, out_dim=1
        ),
        basis_type="scalar",
    )
    return VariablesFE(basis=basis, nb_variables=1)


def dirichlet_bc_zero(x):
    return jnp.zeros(1)


def solve_system(variables_u, variables_v, tag=""):
    model = AbstractPhysicalWeakModel.from_weak_form(pde)
    model.add_boundary_condition("0/boundary", Dirichlet(dirichlet_bc_zero))
    model.add_boundary_condition("1/boundary", Dirichlet(dirichlet_bc_zero))

    assembler = EllipticFEschemeMultipleSpaces(
        pde=model, variables_list=[variables_u, variables_v], equation_spaces=[0, 1]
    )
    print(f"[{tag}] résolution (Newton matrix-free, BiCGStab)…")
    assembler = EllipticFEschemeMultipleSpaces.solve(
        assembler,
        max_iter=1,
        matrix_free=True,
        cg_solver="bicgstab",
        tol=1e-7,
        verbose=False,
    )
    return assembler


# ══════════════════════════════════════════════════════════════════════════════
# ÉTAPE 1 — Solve sur maillage cartésien grossier
# ══════════════════════════════════════════════════════════════════════════════

print("══ Étape 1 : solve couplé (maillage cartésien grossier) ══════════════")

variables_u0 = make_variables(n_cells_u, Mapping(mappings=[mapping_id]))
variables_v0 = make_variables(n_cells_v, Mapping(mappings=[mapping_id]))
assembler0 = solve_system(variables_u0, variables_v0, tag="cartésien")
variables_v0 = assembler0.variables_list[1]


# ══════════════════════════════════════════════════════════════════════════════
# ÉTAPE 2 — Densité monitrice g(y) = 1 + α|∇v_h(y)|²  via ParamFunction
# ══════════════════════════════════════════════════════════════════════════════

print("\n══ Étape 2 : densité monitrice |∇v|² (ParamVecFunction → callable) ═══")

# Monitor arc-length adouci : g_bar = sqrt(1 + α|∇v|²). Le sqrt écrase la
# dynamique (|∇v|² pique fort près de la source) → carte MA régulière.
ALPHA_MONITOR = 10.0  # contraste du monitor : plus grand → compression plus marquée
SMOOTH_SIGMA_PX = 1.5  # écart-type (en pixels) du lissage gaussien du monitor
GRID_MON = 96  # résolution de la grille support du monitor lissé

# Schéma mono-espace autour de variables_v0 : sert uniquement à récupérer
# ``local_variables(idx)`` (ParamScalarFunction FE) puis ``.gradient("x")``.
scheme_v0 = EllipticFEscheme(
    pde=AbstractPhysicalWeakModel.from_weak_form(pde, dirichlet=dirichlet_bc_zero),
    variables=variables_v0,
)


def grad_v_squared_raw(y):
    """|∇v_h(y)|² : find_cell_index → local_variables (ParamFunction) → gradient.

    En Q1 le gradient est constant par maille → carte en escalier : d'où le
    lissage ci-dessous.
    """
    _, cell_idx = variables_v0.mesh.find_cell_index(y[jnp.newaxis, :])
    v_local = scheme_v0.local_variables(cell_idx)  # ParamScalarFunction, f_type "x"
    grad_v = v_local.gradient("x")  # ParamVecFunction (∇v)
    g = grad_v(scheme_v0, y)  # (dim,)
    return jnp.sum(g**2)


# --- Lissage : |∇v|² sur une grille, flou gaussien séparable, puis interpolation
#     bilinéaire différentiable (le monitor reste un callable pour le PINN MA). ---
_xg = jnp.linspace(0.0, 1.0, GRID_MON)
_XXg, _YYg = jnp.meshgrid(_xg, _xg, indexing="ij")
_pts_g = jnp.stack([_XXg.ravel(), _YYg.ravel()], axis=-1)
_gs_grid = jax.vmap(grad_v_squared_raw)(_pts_g).reshape(GRID_MON, GRID_MON)


def _gaussian_blur(img, sigma_px):
    ksize = int(3 * sigma_px)
    ax = jnp.arange(-ksize, ksize + 1)
    ker = jnp.exp(-0.5 * (ax / sigma_px) ** 2)
    ker = ker / ker.sum()
    img = jax.vmap(lambda row: jnp.convolve(row, ker, mode="same"))(img)  # sur x
    img = jax.vmap(lambda col: jnp.convolve(col, ker, mode="same"))(img.T).T  # sur y
    return img


_gbar_grid = jnp.sqrt(1.0 + ALPHA_MONITOR * _gaussian_blur(_gs_grid, SMOOTH_SIGMA_PX))
Z_monitor = float(jnp.mean(_gbar_grid))  # moyenne = 1 (solvabilité MA)


def g_monitor(y):
    """Densité cible MA : monitor LISSÉ, interpolé bilinéairement (callable)."""
    coord = jnp.stack([y[0] * (GRID_MON - 1), y[1] * (GRID_MON - 1)])[:, None]
    val = jax.scipy.ndimage.map_coordinates(_gbar_grid, coord, order=1, mode="nearest")[
        0
    ]
    return jnp.array([val / Z_monitor])


def f_uniform(x):
    return jnp.array([1.0])


# ══════════════════════════════════════════════════════════════════════════════
# ÉTAPE 3 — Monge-Ampère PINN (Picard) + apprentissage de l'inverse + mapping
# ══════════════════════════════════════════════════════════════════════════════

print("\n══ Étape 3 : Monge-Ampère (PINN Picard) + carte inverse ═════════════")

N_COLLOC = 3000
N_BC_COLLOC = 1500
N_EPOCHS_MA = 400
N_INV_EPOCHS = 300
N_NEWTON = 4
EPS_POST = 0.03
HIDDEN = [12] * 5

key = jax.random.PRNGKey(0)
dx = Square2D([(0.0, 1.0), (0.0, 1.0)], is_main_domain=True)
dx2 = Square2D([(-2.0, 2.0), (-2.0, 2.0)], is_main_domain=True)
domain_mu: list = []

sampler_ma = TensorizedSampler(
    [DomainSampler(dx), UniformParametricSampler(domain_mu)], bc=True
)
sampler_int = TensorizedSampler(
    [DomainSampler(dx), UniformParametricSampler(domain_mu)], bc=False
)

level_sets_square = [
    lambda x: x[0],
    lambda x: 1.0 - x[0],
    lambda x: x[1],
    lambda x: 1.0 - x[1],
]

# --- Pré-entraînement de l'ICNN à u = ½|x|² (départ convexe = identité) ---
print("  pré-entraînement ICNN → ½|x|² …")
key, subkey = jax.random.split(key)
nn_u = ICNN(
    in_size=2, out_size=1, hidden_sizes=HIDDEN, activation="softplus", key=subkey
)
space_u = ApproximationSpace({"x": 2}, [(nn_u, "scalar", 1)], model_type="x_mu")
key, proj_pre = Projector(
    FunctionApproximator(
        main_domain=dx2,
        size=1,
        model_type="x_mu",
        f_rhs=lambda *a: jnp.array([0.5 * jnp.dot(a[0], a[0])]),
    ),
    space_u,
    TensorizedSampler(
        [DomainSampler(dx2), UniformParametricSampler(domain_mu)], bc=False
    ),
).project(key, space_u, 10, N_COLLOC)
space_u = proj_pre.space

# --- Monge-Ampère Picard : g(∇u) det(∇²u) = f=1 ---
print(f"  Monge-Ampère Picard ({N_EPOCHS_MA} epochs) …")
model_ma = MongeAmperePicard2D(
    dx, f=f_uniform, g=g_monitor, eps_reg=1e-8, model_type="x_mu", bc="weak"
)
pinn_ma = Projector(
    model_ma,
    space_u,
    sampler_ma,
    weights={"interior": [1.0], "boundary": [6.0]},
    matrix_regularization=1.0e-5,
)
t0 = time.perf_counter()
key, pinn_ma = pinn_ma.project(key, space_u, N_EPOCHS_MA, N_COLLOC, N_BC_COLLOC)
print(
    f"    best loss: {pinn_ma.best_loss['total']:.3e}  ({time.perf_counter() - t0:.1f}s)"
)
space_u = pinn_ma.space

# --- Apprentissage de la carte inverse  v(T(x)) = x ---
print(f"  carte inverse ({N_INV_EPOCHS} epochs) …")
key, subkey = jax.random.split(key)
nn_v = MLP(in_size=2, out_size=2, hidden_sizes=[16, 16, 16], key=subkey)
space_v_inv = ApproximationSpace({"x": 2}, [(nn_v, "field", 2)], model_type="x_mu")
model_inv = InverseFlowMongeAmpere2D(
    dx, space_u=space_u, level_sets=level_sets_square, eps_post=EPS_POST
)
t0 = time.perf_counter()
key, pinn_inv = Projector(model_inv, space_v_inv, sampler_int).project(
    key, space_v_inv, N_INV_EPOCHS, N_COLLOC
)
print(
    f"    best loss: {pinn_inv.best_loss['total']:.3e}  ({time.perf_counter() - t0:.1f}s)"
)

# --- Assemblage de la carte inversible T (forward) / T⁻¹ (backward) ---
icnn_map = InvertibleNetBasedICNN(
    u=space_u.models[0],
    size=2,
    v=pinn_inv.space.models[0],
    n_newton=N_NEWTON,
    level_sets=level_sets_square,
    eps_post=EPS_POST,
)


# ══════════════════════════════════════════════════════════════════════════════
# ÉTAPE 4 — Re-solve avec maillage adapté pour v (mapping T), u inchangé
# ══════════════════════════════════════════════════════════════════════════════

print("\n══ Étape 4 : re-solve, maillage de v adapté par T (u cartésien) ═════")

# Mapping du maillage de v : cellules de référence [0,1]² → physique via T = ∇u.
mapping_adapt = Mapping(mappings=[InvertibleFunction(icnn_map, icnn_map.backward)])

variables_u1 = make_variables(n_cells_u, Mapping(mappings=[mapping_id]))
variables_v1 = make_variables(n_cells_v, mapping_adapt, is_identity=False)

t0 = time.perf_counter()
assembler1 = solve_system(variables_u1, variables_v1, tag="adapté")
print(f"    ({time.perf_counter() - t0:.1f}s)")


# ══════════════════════════════════════════════════════════════════════════════
# Évaluation & visualisation
# ══════════════════════════════════════════════════════════════════════════════

print("\n══ Visualisation ════════════════════════════════════════════════════")

n_plot = 60
xp = np.linspace(0.01, 0.99, n_plot)
XC, YC = np.meshgrid(xp, xp)
pts = jnp.array(np.stack([XC.ravel(), YC.ravel()], axis=-1))

U0 = np.array(assembler0.variables_list[0].evaluate(pts)).reshape(n_plot, n_plot)
V0 = np.array(assembler0.variables_list[1].evaluate(pts)).reshape(n_plot, n_plot)
U1 = np.array(assembler1.variables_list[0].evaluate(pts)).reshape(n_plot, n_plot)
V1 = np.array(assembler1.variables_list[1].evaluate(pts)).reshape(n_plot, n_plot)
GMON = np.array(jax.vmap(lambda y: g_monitor(y)[0])(pts)).reshape(n_plot, n_plot)


def draw_cartesian_mesh(ax, n_cells, **kwargs):
    style = {"color": "w", "linewidth": 0.4, "alpha": 0.5} | kwargs
    for k in range(n_cells + 1):
        ax.axhline(k / n_cells, **style)
        ax.axvline(k / n_cells, **style)


def draw_adapted_mesh(ax, T_fwd, n_cells, **kwargs):
    """Trace le maillage adapté : lignes de la grille de réf. transportées par T."""
    style = {"color": "w", "linewidth": 0.5, "alpha": 0.7} | kwargs
    s = np.linspace(0.0, 1.0, 120)
    cs = np.linspace(0.0, 1.0, n_cells + 1)
    for c in cs:
        line_h = jnp.array(np.stack([s, np.full_like(s, c)], axis=-1))
        line_v = jnp.array(np.stack([np.full_like(s, c), s], axis=-1))
        for line in (line_h, line_v):
            m = np.array(T_fwd(line))
            ax.plot(m[:, 0], m[:, 1], **style)


T_fwd = jax.jit(jax.vmap(icnn_map))
UMA = np.vectorize(lambda a, b: float(space_u.models[0](jnp.array([a, b]))[0]))(XC, YC)
dU = U1 - U0  # effet du couplage : le maillage de u ne change pas, mais v mieux
#                résolu → terme croisé c₁·∇v modifié → u change un peu.

# ══ 3 figures séparées (chacune tient à l'écran) ═══════════════════════════════

# ── FIGURE 1 : maillages de BASE (u, v, monitor lissé) ────────────────────────
figA, axA = plt.subplots(1, 3, figsize=(16, 5.2))
figA.suptitle(
    f"① Solutions sur maillages de BASE — u:{n_cells_u}²(ε={eps_u})  "
    f"v:{n_cells_v}²(ε={eps_v})",
    fontsize=12,
)
im = axA[0].pcolormesh(XC, YC, U0, shading="auto", cmap="turbo")
draw_cartesian_mesh(axA[0], n_cells_u)
axA[0].set_title(r"$u$ initial — mesh cartésien")
plt.colorbar(im, ax=axA[0])

im = axA[1].pcolormesh(XC, YC, V0, shading="auto", cmap="turbo")
draw_cartesian_mesh(axA[1], n_cells_v)
axA[1].set_title(r"$v$ initial — mesh cartésien")
plt.colorbar(im, ax=axA[1])

im = axA[2].pcolormesh(XC, YC, GMON, shading="auto", cmap="Reds")
axA[2].set_title(r"Monitor lissé $g=\sqrt{1+\alpha|\nabla v|^2}$")
plt.colorbar(im, ax=axA[2])
for ax in axA:
    ax.set_aspect("equal")
plt.tight_layout()

# ── FIGURE 2 : mapping (potentiel MA, carte T) + maillage adapté ──────────────
figB, axB = plt.subplots(1, 3, figsize=(16, 5.2))
figB.suptitle(
    "② Monge-Ampère : potentiel, carte de transport, maillage adapté", fontsize=12
)

im = axB[0].pcolormesh(XC, YC, UMA, shading="auto", cmap="viridis")
axB[0].contour(XC, YC, UMA, levels=12, colors="k", linewidths=0.4, alpha=0.5)
axB[0].set_title(r"Potentiel MA $u_{MA}$")
plt.colorbar(im, ax=axB[0])

draw_adapted_mesh(axB[1], T_fwd, n_cells_v, color="k", linewidth=0.5, alpha=0.7)
axB[1].set_title(r"Carte $T=\nabla u_{MA}$ (grille transportée)")
axB[1].set_xlim(-0.02, 1.02)
axB[1].set_ylim(-0.02, 1.02)

axB[2].pcolormesh(XC, YC, GMON, shading="auto", cmap="Reds", alpha=0.55)
draw_adapted_mesh(axB[2], T_fwd, n_cells_v, color="k", alpha=0.7)
axB[2].set_title("Maillage adapté sur le monitor")
axB[2].set_xlim(-0.02, 1.02)
axB[2].set_ylim(-0.02, 1.02)
for ax in axB:
    ax.set_aspect("equal")
plt.tight_layout()

# ── FIGURE 3 : RUN FINAL — u (mesh base), v (mesh adapté), Δu couplage ────────
figC, axC = plt.subplots(1, 3, figsize=(16, 5.2))
figC.suptitle(
    "③ Run final — $u$ sur maillage de BASE, $v$ sur maillage ADAPTÉ", fontsize=12
)
im = axC[0].pcolormesh(XC, YC, U1, shading="auto", cmap="turbo")
draw_cartesian_mesh(axC[0], n_cells_u)
axC[0].set_title(r"$u$ final — mesh cartésien (inchangé)")
plt.colorbar(im, ax=axC[0])

im = axC[1].pcolormesh(XC, YC, V1, shading="auto", cmap="turbo")
draw_adapted_mesh(axC[1], T_fwd, n_cells_v, color="w", linewidth=0.6, alpha=0.8)
axC[1].axhline(0.5, color="w", ls="--", lw=1.0)  # ligne de coupe
axC[1].set_title(r"$v$ final — mesh adapté ($T=\nabla u_{MA}$)")
axC[1].set_xlim(-0.02, 1.02)
axC[1].set_ylim(-0.02, 1.02)
plt.colorbar(im, ax=axC[1])

vlim = float(np.max(np.abs(dU)))
im = axC[2].pcolormesh(
    XC, YC, dU, shading="auto", cmap="seismic", vmin=-vlim, vmax=vlim
)
axC[2].set_title(r"$u_{final}-u_{initial}$ (effet couplage $c_1\!\cdot\!\nabla v$)")
plt.colorbar(im, ax=axC[2])
for ax in axC:
    ax.set_aspect("equal")
plt.tight_layout()

# ── FIGURE 4 : coupe 1D de v en y=0.5 — base (cartésien) vs adapté ─────────────
x_cut = jnp.linspace(0.0, 1.0, 400)
pts_cut = jnp.stack([x_cut, jnp.full_like(x_cut, 0.5)], axis=-1)
v0_cut = np.array(assembler0.variables_list[1].evaluate(pts_cut)).ravel()
v1_cut = np.array(assembler1.variables_list[1].evaluate(pts_cut)).ravel()
x_cut_np = np.array(x_cut)
# positions des nœuds du maillage v en y=0.5 : grille de réf. transportée par T.
nodes_ref = jnp.stack(
    [jnp.linspace(0.0, 1.0, n_cells_v + 1), jnp.full(n_cells_v + 1, 0.5)], axis=-1
)
nodes_adapt = np.array(T_fwd(nodes_ref))
nodes_base = np.linspace(0.0, 1.0, n_cells_v + 1)

figD, axD = plt.subplots(figsize=(11, 5.5))
figD.suptitle("④ Coupe $v(x,\\,y=0.5)$ — maillage de base vs adapté", fontsize=12)
axD.plot(x_cut_np, v0_cut, "-", color="tab:blue", lw=1.8, label="$v$ base (cartésien)")
axD.plot(x_cut_np, v1_cut, "-", color="tab:red", lw=1.8, label="$v$ adapté")
axD.plot(
    nodes_base,
    np.interp(nodes_base, x_cut_np, v0_cut),
    "|",
    color="tab:blue",
    ms=16,
    mew=1.6,
    label="nœuds base (uniformes)",
)
axD.plot(
    nodes_adapt[:, 0],
    np.interp(nodes_adapt[:, 0], x_cut_np, v1_cut),
    "|",
    color="tab:red",
    ms=16,
    mew=1.6,
    label="nœuds adaptés (concentrés au front)",
)
# rug des positions de nœuds en bas pour bien voir la densité
axD.plot(
    nodes_base, np.full_like(nodes_base, -0.15), "|", color="tab:blue", ms=12, mew=1.4
)
axD.plot(
    nodes_adapt[:, 0],
    np.full_like(nodes_adapt[:, 0], -0.30),
    "|",
    color="tab:red",
    ms=12,
    mew=1.4,
)
axD.set_xlabel("x")
axD.set_ylabel("v")
axD.legend(fontsize=9, loc="upper right")
axD.grid(alpha=0.3)
plt.tight_layout()

plt.show()