"""Quad mesh of MAST-U for a free-boundary Grad-Shafranov problem.

The MAST-U-like machine of FreeGSNKE (the geometry of arXiv:2503.05674), meshed
as the ITER case of ``iter_free_boundary_mesh.py``: the half disk
``r^2 + z^2 <= 5^2, r >= 0`` of the poloidal plane, fragmented with

* the limiter region (where the plasma lives) and the two divertor chambers,
  i.e. the wall minus the limiter (the limiter is drawn on the wall, only two
  chords close the chambers);
* the 23 active coil sides, each the rectangle enveloping its filaments with a
  uniform current density, as in the ITER benchmark.

The 138 passive structures (vessel, coil cases, 2-3 mm thick) are read but not
meshed: they carry induced currents only, i.e. they matter for the evolution,
and would need millimetre cells.

⚠ The machine description is not shipped with scimba. Copy the four
``MAST-U_like_*.pickle`` files of ``machine_configs/MAST-U/`` from
https://github.com/FusionComputingLab/freegsnke into
``src/scimba_jax/domains/tokamak/data/freegsnke/MAST-U/`` (ignored by git).

Needs the optional ``mesh`` extra (gmsh). Opens the figure (zoom with the
toolbar) and saves it as ``mastu_free_boundary_mesh.png`` next to this file.
"""

# %%
from pathlib import Path

import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import PolyCollection

from scimba_jax.domains.tokamak import read_freegsnke_machine
from scimba_jax.linear_approximation.meshes.unstructured_mesh import UnstructuredMesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.mapping.macro_mesh import ExactHalfDisk, macro_mesh_from_regions

ORDER = 2  # geometry degree: the far-field arc is the only curved boundary
DOMAIN_RADIUS = 5.0  # m, about twice the machine's half height

# Mesh size (m): fine over the machine, coarsening linearly away from it. The
# 1.2 cm solenoid and the 2 mm wall edges are finer than that: each of their
# edges still gets two segments (see `macro_mesh_from_regions`).
H_MACHINE = 0.03
GROWTH = 0.3
H_FAR = 0.5
MACHINE_BOX = (0.0, 2.1, -2.3, 2.3)  # r_min, r_max, z_min, z_max


def mesh_size(r, z):
    r_min, r_max, z_min, z_max = MACHINE_BOX
    dr = max(r_min - r, 0.0, r - r_max)
    dz = max(z_min - z, 0.0, z - z_max)
    return min(H_FAR, H_MACHINE + GROWTH * float(np.hypot(dr, dz)))


def polygon_area(p):
    return 0.5 * abs(
        np.dot(p[:, 0], np.roll(p[:, 1], -1)) - np.dot(p[:, 1], np.roll(p[:, 0], -1))
    )


# %% Mesh the domain. Regions: limiter first, so that the wall contributes only
# what lies outside it (the divertor chambers); then the coils.
machine = read_freegsnke_machine("MAST-U")
names = ["vacuum", "limiter", "divertor", *(coil.name for coil in machine.coils)]
macro, region = macro_mesh_from_regions(
    ExactHalfDisk(DOMAIN_RADIUS),
    [machine.limiter, machine.wall, *(coil.corners() for coil in machine.coils)],
    order=ORDER,
    mesh_size=mesh_size,
)
mesh = UnstructuredMesh.from_macro_mesh(macro, UnitSquareTensorized(dim=2, order=4))
axis_faces = mesh.label_boundary("axis", lambda x: x[:, 0] < 1e-8)
farfield_faces = mesh.label_boundary("farfield", lambda x: x[:, 0] >= 1e-8)
# The regions become named groups of the mesh, what a weak form reads with
# `self.per_region({"PF1": J1, ...})` or `self.in_region("limiter")`. The
# vacuum is the cells of no group (the `default` of per_region).
for k, name in enumerate(names[1:], start=1):
    mesh.label_cells(name, np.flatnonzero(region == k))

weights, _ = mesh.evaluate_mesh_weights_points()
cell_area = np.asarray(jnp.sum(weights, axis=1))
print(f"{mesh.n_cells_total} cells (order {ORDER})")
print(f"boundary faces: {len(axis_faces)} on the axis, {len(farfield_faces)} far field")
half_disk_area = np.pi * DOMAIN_RADIUS**2 / 2
print(
    f"half disk: area {cell_area.sum():.6f} vs {half_disk_area:.6f} "
    f"(relative {abs(cell_area.sum() / half_disk_area - 1):.1e})"
)
exact = [
    polygon_area(machine.limiter),
    polygon_area(machine.wall) - polygon_area(machine.limiter),
    *(coil.area for coil in machine.coils),
]
print(f"{'region':9s} {'cells':>6s} {'area':>10s} {'exact':>10s} {'rel. err':>9s}")
for k, (name, target) in enumerate(zip(names[1:], exact), start=1):
    cells = mesh.cell_groups[name]
    area = cell_area[cells].sum()
    print(
        f"{name:9s} {len(cells):6d} {area:10.6f} {target:10.6f} "
        f"{abs(area / target - 1):9.1e}"
    )

# %% Plot: whole domain, the machine, and the lower divertor.
corner_ids = [0, ORDER, (ORDER + 1) ** 2 - 1, ORDER * (ORDER + 1)]
nodes = np.asarray(macro.nodes)
quads = nodes[np.asarray(macro.cells)[:, corner_ids]]
palette = plt.get_cmap("tab20")
colors = np.array(
    [(0.95, 0.95, 0.95, 1.0), (0.55, 0.75, 0.95, 1.0), (0.75, 0.85, 0.65, 1.0)]
    + [palette(k % 20) for k in range(len(machine.coils))]
)

fig, axes = plt.subplots(1, 3, figsize=(16, 8))
views = [
    ("computational domain", (-0.2, 5.2, -5.2, 5.2)),
    ("the machine", MACHINE_BOX),
    ("lower divertor", (0.15, 1.95, -2.15, -1.3)),
]
for ax, (title, limits) in zip(axes, views):
    ax.add_collection(
        PolyCollection(quads, facecolors=colors[region], edgecolors="k", linewidths=0.1)
    )
    ax.set_xlim(limits[0], limits[1])
    ax.set_ylim(limits[2], limits[3])
    ax.set_aspect("equal")
    ax.set_xlabel("r [m]")
    ax.set_ylabel("z [m]")
    ax.set_title(title)
for coil in machine.coils:
    if coil.z_max < 0 or coil.name == "Solenoid":
        continue  # label the upper sides only, the machine is up-down symmetric
    axes[1].text(
        coil.r_max + 0.02,
        (coil.z_min + coil.z_max) / 2,
        coil.name.removesuffix("_1"),
        va="center",
        fontsize=7,
    )
fig.suptitle(
    f"MAST-U free-boundary mesh (FreeGSNKE geometry): {mesh.n_cells_total} quads, "
    "limiter in blue, divertor in green, coils coloured"
)
fig.tight_layout()
out = Path(__file__).with_name("mastu_free_boundary_mesh.png")
fig.savefig(out, dpi=140)
print(f"Saved figure to {out}")

plt.show()