r"""3D visualisation of the extruded toroidal volume.

For each machine: one figure with 3 panels:

  A (left)    — poloidal cross-section (R, Z)
  B (centre)  — 3D torus: wall surface + interior points (sector)
  C (right)   — top view (x, y) coloured by 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  # 3D interior points
N_WALL_PHI = 60  # number of angles for the wall surface
PHI_SECTOR = np.pi  # half-angle of the sector shown in 3D (π = half torus)
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):
    """Draw the toroidal wall surface (extrusion of the poloidal contour)."""
    phi = np.linspace(0, 2 * np.pi, n_phi)
    # poloidal lines (wall contour at each φ)
    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)
    # toroidal circles along the wall
    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: poloidal cross-section (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}\n(R, Z) cross-section coloured by φ", fontsize=9)

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

    # Toroidal wall surface (wireframe)
    draw_toroidal_wall(ax3d, R_wall, Z_wall, n_phi=N_WALL_PHI)

    # Interior points — half torus only (φ ∈ [0, π]) to see the inside
    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("3D torus (half torus, coloured by Z)", fontsize=9)
    ax3d.view_init(elev=25, azim=-50)

    # ── C: top view (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)
    # R_min and R_max circles
    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("Top view\n(coloured by 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"3D sampling — {name}\n"
        r"toroidal extrusion  (R, Z, φ)  with Jacobian $\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.")