"""Macro-mesh of a **rectangle with a hole**: GMSH-auto vs hand-built O-grid.

Left panel: :func:`~scimba_jax.mapping.macro_mesh.macro_mesh_from_points` with
its ``holes`` argument — outer boundary a straight-edged rectangle
(``smooth=False``), hole sampled from a closed parametrized curve
``t -> (x, y)`` (``holes_smooth=True``, spline-fit) — tested here with a
circle, but ``hole_curve`` swaps for any closed curve (ellipse, flower, ...)
exactly like :func:`~scimba_jax.mapping.macro_mesh.macro_mesh_from_curve`. GMSH
meshes the annular region between the two loops with however many quads its
heuristic wants.

Right panel: a hand-built **O-grid** (same spirit as
:func:`~scimba_jax.mapping.macro_mesh.macro_mesh_ogrid_disk`, adapted to an
outer rectangle instead of an outer circle and to a *hole* instead of a filled
disk) — ``N_PER_SIDE=1`` gives exactly **4** radial blocks (one per rectangle
side, outer edge = the full straight side, inner edge = an exact quarter-circle
arc through 3 points via :class:`~scimba_jax.mapping.quad_mappings.TransfiniteQuadMapping`);
``N_PER_SIDE=2`` splits each side at its midpoint for **8** blocks (smaller
arcs, less quadratic-vs-true-circle deviation — see the NURBS-vs-polynomial
discussion this example grew out of). No GMSH involved.

Needs the optional ``mesh`` extra (pygmsh/meshio/gmsh) for the left panel only.
Saves ``macro_mesh_hole.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.mapping.macro_mesh import MacroMesh, macro_mesh_from_points
from scimba_jax.mapping.quad_mappings import TransfiniteQuadMapping

ORDER = 4  # geometry degree for the GMSH panel (all straight/circle-spline here)
MESH_SIZE = 0.9
N_HOLE_BOUNDARY = 60
N_PER_SIDE = 2  # O-grid blocks per rectangle side: 1 -> 4 blocks, 2 -> 8 blocks

W, H = 4.0, 2.0  # outer rectangle half-extents: [-W,W] x [-H,H]
RADIUS = 0.6
_RECTANGLE = np.array([[-W, -H], [W, -H], [W, H], [-W, H]])


def hole_curve(t):
    """Closed hole boundary t -> (x, y), t in [0, 1). Circle of radius RADIUS."""
    theta = 2 * np.pi * t
    return RADIUS * np.cos(theta), RADIUS * np.sin(theta)


_ts = np.linspace(0.0, 1.0, N_HOLE_BOUNDARY, endpoint=False)
_hole = np.array([hole_curve(float(t)) for t in _ts])


def rectangle_with_hole_ogrid(w, h, radius, n_per_side=1) -> MacroMesh:
    """Hand-built O-grid: rectangle [-w,w]x[-h,h] minus a disk of radius ``radius``.

    ``n_per_side`` outer points per rectangle side (always including both
    corners), each paired with an inner circle point at the *same angle* from
    the origin: consecutive outer points are collinear (a side, or part of one)
    so the block's outer edge is exactly straight, and its inner edge is the
    circular arc between the two angles — built as a quadratic curve through
    the arc's exact endpoints *and* true midpoint
    (:class:`TransfiniteQuadMapping`), not a full NURBS circle, so it is only
    an **approximation** of the arc (see module docstring).
    """
    corners = [(-w, -h), (w, -h), (w, h), (-w, h)]
    outer = []
    for i in range(4):
        p0, p1 = np.array(corners[i]), np.array(corners[(i + 1) % 4])
        outer += [p0 + (p1 - p0) * (k / n_per_side) for k in range(n_per_side)]
    outer = np.array(outer)
    n = len(outer)
    angles = np.arctan2(outer[:, 1], outer[:, 0])
    inner = np.stack([radius * np.cos(angles), radius * np.sin(angles)], axis=1)
    nodes = np.vstack([outer, inner])

    cells, mappings, cell_types = [], [], []
    for k in range(n):
        kp = (k + 1) % n
        p0, p1, p2, p3 = outer[k], outer[kp], inner[kp], inner[k]
        cells.append([k, n + k, kp, n + kp])  # tensor order [P0,P3,P1,P2]
        theta_mid = angles[k] + ((angles[kp] - angles[k]) % (2 * np.pi)) / 2
        q_mid = radius * np.array([np.cos(theta_mid), np.sin(theta_mid)])
        a = 2 * p3 - 4 * q_mid + 2 * p2
        b = -3 * p3 + 4 * q_mid - p2
        mappings.append(
            TransfiniteQuadMapping(np.stack([p0, p1, p2, p3]), np.stack([a, b, p3]))
        )
        cell_types.append("transfinite")
    return MacroMesh(
        nodes, np.array(cells), order=1, mappings=mappings, cell_types=cell_types
    )


# %% 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 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")


# %% Build + plot both versions side by side.
mm_gmsh = macro_mesh_from_points(
    _RECTANGLE,
    order=ORDER,
    mesh_size=MESH_SIZE,
    smooth=False,
    holes=[_hole],
    holes_smooth=True,
)
mm_ogrid = rectangle_with_hole_ogrid(W, H, RADIUS, n_per_side=N_PER_SIDE)
print("GMSH auto: ", mm_gmsh)
print("Hand O-grid:", mm_ogrid)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6.5))
draw(ax1, mm_gmsh)
ax1.plot(_hole[:, 0], _hole[:, 1], "k.", ms=2, label=f"hole curve ({len(_hole)} pts)")
ax1.legend(loc="upper right", fontsize=8)
ax1.set_title(f"macro_mesh_from_points (GMSH auto): {mm_gmsh.n_cells} cells")

draw(ax2, mm_ogrid)
ax2.set_title(f"hand-built O-grid: {mm_ogrid.n_cells} cells (N_PER_SIDE={N_PER_SIDE})")

fig.suptitle(
    "Rectangle with a circular hole — blue=bilinear, orange=curved, "
    "red=boundary (outer + hole)",
    fontsize=13,
)
out = Path(__file__).with_name("macro_mesh_hole.png")
fig.tight_layout()
fig.savefig(out, dpi=140)
print(f"\nSaved figure to {out}")

plt.show()