"""Quad mesh of the ITER free-boundary Grad-Shafranov benchmark.

The domain of Serino et al., *An adaptive Newton-based free-boundary
Grad-Shafranov solver* (arXiv:2407.03499): the half disk ``r^2 + z^2 <= 16^2,
r >= 0`` of the poloidal plane, containing the first wall (the limiter region,
where the plasma lives) and the eleven coils (five central-solenoid modules,
six PF coils). The geometry is :mod:`scimba_jax.domains.tokamak.iter_geometry`,
copied from the authors' mesh generator.

:func:`~scimba_jax.mapping.macro_mesh.macro_mesh_from_regions` fragments the
half disk with the wall and the coils, so every interface is a chain of mesh
edges, and returns the region of each cell: the coils' current densities
``I_i / |Omega_Ci|`` and the plasma source can then be integrated cell by cell
with no cell straddling an interface. The check printed below is that the mesh
reproduces each coil's area exactly (they are rectangles) and the half disk's
up to the order-2 arc.

The boundaries are named after the fact, as elsewhere in scimba: ``"axis"``
(``r = 0``, where ``psi = 0``) and ``"farfield"`` (the arc, which carries the
free-boundary terms).

Needs the optional ``mesh`` extra (gmsh). Opens the figure (zoom with the
toolbar) and saves it as ``iter_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 DOMAIN_RADIUS, ITER_COILS, iter_first_wall
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

# Mesh size (m): fine over the machine (wall and coils), coarsening linearly
# with the distance to it, capped in the far field.
H_MACHINE = 0.25
GROWTH = 0.25
H_FAR = 2.0
MACHINE_BOX = (0.5, 13.5, -9.0, 9.0)  # 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)))


# %% Mesh the domain, regions in the paper's order: limiter, then I_1 .. I_11.
names = ["vacuum", "limiter", *(coil.name for coil in ITER_COILS)]
macro, region = macro_mesh_from_regions(
    ExactHalfDisk(DOMAIN_RADIUS),
    [iter_first_wall(), *(coil.corners() for coil in ITER_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})"
)
wall = iter_first_wall()
wall_area = 0.5 * abs(
    np.dot(wall[:, 0], np.roll(wall[:, 1], -1))
    - np.dot(wall[:, 1], np.roll(wall[:, 0], -1))
)  # shoelace
exact = [wall_area, *(coil.area for coil in ITER_COILS)]
print(f"{'region':8s} {'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:8s} {len(cells):6d} {area:10.6f} {target:10.6f} "
        f"{abs(area / target - 1):9.1e}"
    )

# %% Plot: whole domain, and the machine.
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)]
    + [palette(k) for k in range(len(ITER_COILS))]
)

fig, axes = plt.subplots(1, 2, figsize=(13, 8))
for ax, (title, limits) in zip(
    axes,
    [
        ("computational domain", (-0.5, 16.5, -16.5, 16.5)),
        ("the machine", MACHINE_BOX),
    ],
):
    ax.add_collection(
        PolyCollection(quads, facecolors=colors[region], edgecolors="k", linewidths=0.2)
    )
    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 ITER_COILS:
    axes[1].text(
        (coil.r_min + coil.r_max) / 2,
        (coil.z_min + coil.z_max) / 2,
        coil.name,
        ha="center",
        va="center",
        fontsize=8,
    )
fig.suptitle(
    f"ITER free-boundary mesh (arXiv:2407.03499): {mesh.n_cells_total} quads, "
    "limiter in blue, coils coloured"
)
fig.tight_layout()
out = Path(__file__).with_name("iter_free_boundary_mesh.png")
fig.savefig(out, dpi=140)
print(f"Saved figure to {out}")

plt.show()