r"""Visualisation 2D des coupes poloidales tokamak.

Pour chaque machine on trace sur une même figure :
  - les points intérieurs (bleu, sampling par rejet dans le mur)
  - les points de bord (rouge) avec les normales sortantes (orange)
  - le contour du mur (noir)

Usage::

    python examples/examples_jax/geometry/tokamak_plot_2d.py
"""

from __future__ import annotations

import sys
from pathlib import Path

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

sys.path.insert(0, str(Path(__file__).parent.parent.parent.parent / "src"))
jax.config.update("jax_enable_x64", True)

from tokamak_domains import TOKAMAKS  # noqa: E402

# ── Config ─────────────────────────────────────────────────────────────────────
N_INT = 6_000
N_BC = 600
KEY = jax.random.PRNGKey(0)

# ── Figure ─────────────────────────────────────────────────────────────────────
n = len(TOKAMAKS)
ncols = min(3, n)
nrows = (n + ncols - 1) // ncols

fig, axes = plt.subplots(
    nrows, ncols, figsize=(6 * ncols, 6 * nrows), constrained_layout=True
)
fig.suptitle("Sampling 2D — coupes poloidales tokamak", fontsize=14, y=1.01)

axes_flat = np.array(axes).ravel()
key = KEY

for ax, (name, s2d, _, R_wall, Z_wall) in zip(axes_flat, TOKAMAKS):
    key, sample = s2d.sample(key, n=N_INT, n_bc=N_BC)
    x_int = np.array(sample["interior"][0])
    x_bc, nrm = sample["boundary"]
    x_bc = np.array(x_bc)
    nrm = np.array(nrm)

    # interior
    ax.scatter(
        x_int[:, 0],
        x_int[:, 1],
        s=0.8,
        c="steelblue",
        alpha=0.35,
        rasterized=True,
        label=f"intérieur ({N_INT})",
    )

    # boundary + normals
    ax.scatter(
        x_bc[:, 0], x_bc[:, 1], s=6, c="crimson", zorder=4, label=f"bord ({N_BC})"
    )
    sk = max(1, N_BC // 40)
    scale = 0.05 * (R_wall.max() - R_wall.min())
    ax.quiver(
        x_bc[::sk, 0],
        x_bc[::sk, 1],
        nrm[::sk, 0] * scale,
        nrm[::sk, 1] * scale,
        angles="xy",
        scale_units="xy",
        scale=1.0,
        color="darkorange",
        width=0.004,
        alpha=0.85,
        label="normales",
    )

    # wall contour
    R_c = np.append(R_wall, R_wall[0])
    Z_c = np.append(Z_wall, Z_wall[0])
    ax.plot(R_c, Z_c, "k-", lw=1.5, zorder=5)

    ax.set_title(name, fontsize=10)
    ax.set_xlabel("R [m]")
    ax.set_ylabel("Z [m]")
    ax.set_aspect("equal")
    ax.legend(markerscale=4, fontsize=7, loc="upper right")
    ax.text(
        0.02,
        0.02,
        f"mur : {len(R_wall)} pts",
        transform=ax.transAxes,
        fontsize=7,
        color="gray",
    )

for ax in axes_flat[n:]:
    ax.set_visible(False)

plt.savefig("tokamak_sampling_2d.png", dpi=150, bbox_inches="tight")
plt.show()
print("Saved: tokamak_sampling_2d.png")