"""Macro-mesh of the classic **L-shaped domain**: GMSH-auto vs hand-built.

Left panel: :func:`~scimba_jax.mapping.macro_mesh.macro_mesh_from_points` on the
6 corners of an L-shape (the textbook re-entrant-corner test domain for
elliptic PDEs: the interior angle at the notch is ``3*pi/2``, source of the
classic ``r^(2/3)`` corner singularity). GMSH's automatic unstructured
recombination does **not** stop at the "obvious" 3-block split, even at a
coarse ``mesh_size`` — it tiles the domain with however many quads its
triangulate-then-recombine heuristic produces.

Right panel: the L-shape is a **rectilinear** polygon (axis-aligned edges
only), so it decomposes into exactly 3 blocks fully automatically via
:func:`~scimba_jax.mapping.macro_mesh.macro_mesh_rectilinear` — no GMSH, no
manual corner-picking: it grids the plane by the corners' distinct x/y values
and keeps the cells whose center is inside the polygon. Generalizes to any
rectilinear shape (staircases, plus/cross, ...), not just this one L-shape.

Needs the optional ``mesh`` extra (pygmsh/meshio/gmsh) for the left panel only.
Saves ``macro_mesh_lshape.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 macro_mesh_from_points, macro_mesh_rectilinear

ORDER = 3  # geometry degree for the GMSH panel (all straight edges here anyway)
MESH_SIZE = 0.8

# L-shape corners, CCW, re-entrant corner at the origin.
_CORNERS = np.array(
    [[0.0, 0.0], [2.0, 0.0], [2.0, 1.0], [1.0, 1.0], [1.0, 2.0], [0.0, 2.0]]
)


# %% 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(
    _CORNERS, order=ORDER, mesh_size=MESH_SIZE, smooth=False
)
mm_auto = macro_mesh_rectilinear(_CORNERS)
print("GMSH auto:  ", mm_gmsh)
print("Rectilinear:", mm_auto)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 7))
for ax, mm, title in [
    (ax1, mm_gmsh, f"macro_mesh_from_points (GMSH auto): {mm_gmsh.n_cells} cells"),
    (ax2, mm_auto, f"macro_mesh_rectilinear (auto): {mm_auto.n_cells} cells"),
]:
    draw(ax, mm)
    ax.plot(*np.vstack([_CORNERS, _CORNERS[:1]]).T, "kx", ms=6, mew=1.5)
    ax.set_title(title)

fig.suptitle(
    "L-shape (re-entrant corner) — blue=bilinear, orange=curved, red=boundary",
    fontsize=13,
)
out = Path(__file__).with_name("macro_mesh_lshape.png")
fig.tight_layout()
fig.savefig(out, dpi=140)
print(f"\nSaved figure to {out}")

plt.show()