r"""Visualisation 3D du volume toroïdal extrudé.

Pour chaque machine : une figure avec 3 panneaux :

  A (gauche)  — coupe poloidale (R, Z)
  B (centre)  — tore 3D : surface du mur + points intérieurs (secteur)
  C (droite)  — vue du dessus (x, y) colorée par Z

Usage::

    python examples/examples_jax/geometry/tokamak_plot_3d.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  # points intérieurs 3D
N_WALL_PHI = 60  # nb d'angles pour la surface du mur
PHI_SECTOR = np.pi  # demi-angle du secteur affiché en 3D (π = demi-tore)
KEY = jax.random.PRNGKey(42)


def rz_phi_to_xyz(R, Z, phi):
    return R * np.cos(phi), R * np.sin(phi), Z


def draw_toroidal_wall(ax, R_wall, Z_wall, alpha_surface=0.15, n_phi=60):
    """Dessine la surface toroïdale du mur (extrusión du contour poloidale)."""
    phi = np.linspace(0, 2 * np.pi, n_phi)
    # lignes poloïdales (contour du mur à chaque φ)
    for ph in phi[:: n_phi // 8]:
        Rw = np.append(R_wall, R_wall[0])
        Zw = np.append(Z_wall, Z_wall[0])
        ax.plot(Rw * np.cos(ph), Rw * np.sin(ph), Zw, "k-", lw=0.6, alpha=0.4)
    # cercles toroïdaux le long du mur
    for i in range(0, len(R_wall), max(1, len(R_wall) // 10)):
        ax.plot(
            R_wall[i] * np.cos(phi),
            R_wall[i] * np.sin(phi),
            np.full_like(phi, Z_wall[i]),
            "k-",
            lw=0.4,
            alpha=0.25,
        )


def plot_one_tokamak(fig, row, nrows, name, s3d, R_wall, Z_wall, key):
    key, sample = s3d.sample(key, n=N_INT)
    pts = np.array(sample["interior"][0])  # (N, 3) = (R, Z, phi)
    R, Z, phi = pts[:, 0], pts[:, 1], pts[:, 2]
    xc, yc, zc = rz_phi_to_xyz(R, Z, phi)

    Rw = np.append(R_wall, R_wall[0])
    Zw = np.append(Z_wall, Z_wall[0])
    theta_full = np.linspace(0, 2 * np.pi, 400)

    # ── A : coupe poloidale (R, Z) ─────────────────────────────────────────────
    ax_rz = fig.add_subplot(nrows, 3, row * 3 + 1)
    ax_rz.scatter(R, Z, s=1.0, c=phi, cmap="hsv", alpha=0.4, rasterized=True)
    ax_rz.plot(Rw, Zw, "k-", lw=1.5)
    ax_rz.set_xlabel("R [m]")
    ax_rz.set_ylabel("Z [m]")
    ax_rz.set_aspect("equal")
    ax_rz.set_title(f"{name}\ncoupe (R, Z) colorée par φ", fontsize=9)

    # ── B : tore 3D ────────────────────────────────────────────────────────────
    ax3d = fig.add_subplot(nrows, 3, row * 3 + 2, projection="3d")

    # Surface du mur toroïdal (fil de fer)
    draw_toroidal_wall(ax3d, R_wall, Z_wall, n_phi=N_WALL_PHI)

    # Points intérieurs — seulement demi-tore (φ ∈ [0, π]) pour voir l'intérieur
    mask = phi >= 0
    sk = max(1, mask.sum() // 800)
    ax3d.scatter(
        xc[mask][::sk],
        yc[mask][::sk],
        zc[mask][::sk],
        s=2.5,
        c=zc[mask][::sk],
        cmap="coolwarm",
        alpha=0.6,
        depthshade=True,
    )

    ax3d.set_xlabel("x [m]", fontsize=7)
    ax3d.set_ylabel("y [m]", fontsize=7)
    ax3d.set_zlabel("Z [m]", fontsize=7)
    ax3d.tick_params(labelsize=6)
    ax3d.set_title("Tore 3D (demi-tore, coloré par Z)", fontsize=9)
    ax3d.view_init(elev=25, azim=-50)

    # ── C : vue du dessus (x, y) ───────────────────────────────────────────────
    ax_top = fig.add_subplot(nrows, 3, row * 3 + 3)
    sk2 = max(1, len(xc) // 2000)
    sc = ax_top.scatter(
        xc[::sk2],
        yc[::sk2],
        s=1.5,
        c=zc[::sk2],
        cmap="viridis",
        alpha=0.5,
        rasterized=True,
    )
    fig.colorbar(sc, ax=ax_top, label="Z [m]", fraction=0.046, pad=0.04)
    # cercles R_min et R_max
    for R_ref, ls, lw in [(R_wall.min(), "--", 0.8), (R_wall.max(), "-", 1.2)]:
        ax_top.plot(
            R_ref * np.cos(theta_full),
            R_ref * np.sin(theta_full),
            "k" + ls,
            lw=lw,
            alpha=0.6,
        )
    ax_top.set_xlabel("x [m]")
    ax_top.set_ylabel("y [m]")
    ax_top.set_aspect("equal")
    ax_top.set_title("Vue du dessus\n(colorée par Z, R_min/R_max)", fontsize=9)

    return key


# ── Main ───────────────────────────────────────────────────────────────────────

key = KEY
for idx, (name, _s2d, s3d, R_wall, Z_wall) in enumerate(TOKAMAKS):
    print(f"  {name} …")
    fig = plt.figure(figsize=(17, 6), constrained_layout=True)
    fig.suptitle(
        f"Sampling 3D — {name}\n"
        r"extrusion toroïdale  (R, Z, φ)  avec jacobien $\propto R$",
        fontsize=11,
    )
    key = plot_one_tokamak(fig, 0, 1, name, s3d, R_wall, Z_wall, key)
    fname = f"tokamak_3d_{idx + 1}_{name.split()[0].replace('/', '_')}.png"
    plt.savefig(fname, dpi=130, bbox_inches="tight")
    print(f"    → {fname}")

plt.show()
print("Done.")