"""Macro-mesh of real tokamak cross-sections (JET / MAST), via the point-cloud API.

Uses :func:`~scimba_jax.mapping.macro_mesh.macro_mesh_from_points` on the real
wall polygon read from an EQDSK equilibrium file
(:func:`~scimba_jax.domains.tokamak.read_eqdsk_wall`): one call turns the wall
point cloud into a block-structured
:class:`~scimba_jax.mapping.macro_mesh.MacroMesh` of the poloidal cross-section.

The wall polygon is a **straight-segment** shape, but some of its segments are
sparse (MAST has a ~1 m gap between two consecutive points): fitting a closed
spline (``smooth=True``, needed for GMSH's 1-D mesher to stay well-posed on a
loop this jagged) straight through the raw points would bulge away from the
true polygon on those sparse stretches. :func:`_densify_polygon` inserts extra
points *on* the existing straight segments (no new geometry) so the spline
hugs the true wall shape everywhere, corners and notches included.

Needs the optional ``mesh`` extra (pygmsh/meshio/gmsh). Saves
``macro_mesh_tokamak.png`` next to this file.
"""

# %%
from pathlib import Path

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

from scimba_jax.domains.tokamak import read_eqdsk_wall
from scimba_jax.mapping.macro_mesh import macro_mesh_from_points

ORDER = 3  # geometry degree: 2=quad9, 3=quad16, 4=quad25, ...

_DATA = Path(__file__).parents[4] / "src/scimba_jax/domains/tokamak/data"
_TOKAMAKS = {
    "JET": (_DATA / "eqdsk_jet_compare.dat", 0.20),
    "MAST": (_DATA / "g_p49320_t0.60000", 0.20),
}

# %% Reference-space samples reused for drawing.
_REF_GRID = jnp.stack(
    jnp.meshgrid(jnp.linspace(0, 1, 5), jnp.linspace(0, 1, 5)), axis=-1
).reshape(-1, 2)
_S = jnp.linspace(0.0, 1.0, 30)
_ZERO, _ONE = jnp.zeros_like(_S), jnp.ones_like(_S)
_REF_EDGES = {
    0: jnp.stack([_S, _ZERO], -1),
    1: jnp.stack([_ONE, _S], -1),
    2: jnp.stack([_S, _ONE], -1),
    3: jnp.stack([_ZERO, _S], -1),
}
STRAIGHT, CURVED = "#4C78A8", "#F58518"


def _dedupe_polygon(points: np.ndarray, tol: float = 1e-3) -> np.ndarray:
    """Drop near-duplicate consecutive points on a closed polygon.

    The raw EQDSK wall polygon has a few points closer together than ``tol``
    (down to ~1e-4 m); such near-zero-length segments make GMSH's 1-D mesher
    choke on the straight-segment loop, so they are merged away before meshing
    (the wall *shape* is unchanged — this only removes redundant sampling).
    """
    keep = [points[0]]
    for p in points[1:]:
        if np.hypot(*(p - keep[-1])) > tol:
            keep.append(p)
    if np.hypot(*(keep[0] - keep[-1])) <= tol:
        keep.pop()
    return np.array(keep)


def _densify_polygon(points: np.ndarray, max_seg: float) -> np.ndarray:
    """Insert linearly-interpolated points so no edge exceeds ``max_seg``.

    The EQDSK wall is itself a straight-segment polygon (the true wall shape
    *is* the piecewise-linear interpolation of its points), but some segments
    are sparse (MAST has a ~1 m straight run between two points). The closed
    spline fit through too few points bulges away from that straight polygon;
    inserting extra points *on* the existing segments (no new geometry, just
    denser sampling of the same straight lines) keeps the spline hugging the
    true wall shape.
    """
    out = []
    n = len(points)
    for i in range(n):
        a, b = points[i], points[(i + 1) % n]
        length = np.hypot(*(b - a))
        n_sub = max(1, int(np.ceil(length / max_seg)))
        for k in range(n_sub):
            out.append(a + (b - a) * (k / n_sub))
    return np.array(out)


def patch_boundary(mapping):
    """Physical polyline of a patch's boundary loop (exact, via its mapping)."""
    loop = jnp.concatenate(
        [_REF_EDGES[0], _REF_EDGES[1], _REF_EDGES[2][::-1], _REF_EDGES[3][::-1]]
    )
    return np.asarray(jax.vmap(mapping.forward)(loop))


def draw(ax, mm):
    """Draw a MacroMesh: patches, sub-grids, and physical-boundary faces."""
    for i in range(mm.n_cells):
        m = mm.cell_mapping(i)
        col = STRAIGHT if mm.cell_types[i] == "bilinear" else CURVED
        bnd = patch_boundary(m)
        ax.fill(bnd[:, 0], bnd[:, 1], color=col, alpha=0.18)
        ax.plot(bnd[:, 0], bnd[:, 1], color=col, lw=0.8)
        grid = np.asarray(jax.vmap(m.forward)(_REF_GRID))
        ax.plot(grid[:, 0], grid[:, 1], ".", color=col, ms=1.2)
    for patch, edge in mm.boundary_faces:
        seg = np.asarray(
            jax.vmap(mm.cell_mapping(int(patch)).forward)(_REF_EDGES[int(edge)])
        )
        ax.plot(seg[:, 0], seg[:, 1], "r-", lw=2.0)
    ax.set_aspect("equal")
    ax.set_xlabel("R [m]")
    ax.set_ylabel("Z [m]")


# %% Build + plot the JET and MAST macro-meshes from their real EQDSK wall.
fig, axes = plt.subplots(1, len(_TOKAMAKS), figsize=(7 * len(_TOKAMAKS), 7))
for ax, (name, (path, mesh_size)) in zip(axes, _TOKAMAKS.items()):
    R_wall, Z_wall = read_eqdsk_wall(str(path))
    boundary_raw = np.stack(
        [R_wall[:-1], Z_wall[:-1]], axis=1
    )  # drop repeated closing pt
    boundary_raw = _dedupe_polygon(boundary_raw)
    boundary = _densify_polygon(boundary_raw, max_seg=0.03)
    mm = macro_mesh_from_points(boundary, order=ORDER, mesh_size=mesh_size, smooth=True)
    print(
        f"{name}: {mm} ({len(boundary_raw)} raw wall pts -> {len(boundary)} densified)"
    )
    draw(ax, mm)
    ax.plot(
        boundary[:, 0],
        boundary[:, 1],
        "k.",
        ms=2.5,
        label=f"EQDSK wall points ({len(boundary)}, densified)",
    )
    ax.legend(loc="upper right", fontsize=8)
    ax.set_title(f"{name} wall — {mm.n_cells} macro-cells (order {ORDER})")

fig.suptitle(
    "macro_mesh_from_points(EQDSK wall) — real tokamak cross-sections "
    "(blue=bilinear, orange=curved, red=boundary)",
    fontsize=13,
)
out = Path(__file__).with_name("macro_mesh_tokamak.png")
fig.tight_layout()
fig.savefig(out, dpi=140)
print(f"\nSaved figure to {out}")

plt.show()