"""Macro-mesh from a parametrized boundary curve, at increasing refinement.

Uses the high-level API :func:`~scimba_jax.mapping.macro_mesh.macro_mesh_from_curve`:
one call turns a closed boundary curve (here the unit circle) into a
block-structured :class:`~scimba_jax.mapping.macro_mesh.MacroMesh`. The figure
shows the same disk at three mesh sizes (coarse → fine), each panel drawing, per
macro-cell, its mapping (interior sub-grid) and the exact curved boundary
(bilinear cells blue, curved order-``p`` cells orange, physical boundary red).

Swap the ``curve`` for any closed ``t -> (x, y)`` (or use
``macro_mesh_from_points`` on a boundary point cloud). Needs the optional
``mesh`` extra (pygmsh/meshio/gmsh). Saves ``macro_mesh_from_curve.png``.
"""

# %%
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 macro_mesh_from_curve

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


def circle(t):
    """Unit-circle boundary curve t -> (x, y), t in [0, 1)."""
    return np.cos(2 * np.pi * t), np.sin(2 * np.pi * t)


# Three mesh sizes giving increasingly fine disks (~5, ~10, ~20 macro-cells).
MESH_SIZES = [0.85, 0.72, 0.44]

# %% 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 the disk at three refinements.
fig, axes = plt.subplots(1, len(MESH_SIZES), figsize=(6 * len(MESH_SIZES), 6))
for ax, h in zip(axes, MESH_SIZES):
    mm = macro_mesh_from_curve(circle, order=ORDER, mesh_size=h)
    print(mm)
    draw(ax, mm)
    ax.set_title(f"{mm.n_cells} macro-cells  (mesh_size={h})")

fig.suptitle(
    f"macro_mesh_from_curve(unit circle, order={ORDER}) — refinement "
    "(blue=bilinear, orange=curved, red=boundary)",
    fontsize=13,
)
out = Path(__file__).with_name("macro_mesh_from_curve.png")
fig.tight_layout()
fig.savefig(out, dpi=140)
print(f"\nSaved figure to {out}")

plt.show()