"""Macro-mesh of a smooth 3-petal flower, via the parametrized-curve API.

Same one-call pipeline as ``example_macro_mesh_from_curve`` but on a non-convex,
smooth boundary — a 3-petal rose ``r(theta) = 1 + a cos(3 theta)``. Shows the
block-structured :class:`~scimba_jax.mapping.macro_mesh.MacroMesh` on a curved
domain: bilinear (straight) cells blue, curved order-``p`` cells orange, physical
boundary red, plus the macro-cell adjacency graph.

Needs the optional ``mesh`` extra (pygmsh/meshio/gmsh). Saves
``macro_mesh_flower.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_curve

ORDER = 3
PETAL_AMPLITUDE = 0.35  # petal depth; boundary r = 1 + a cos(3 theta)


def flower(t):
    """Smooth 3-petal rose boundary t -> (x, y), t in [0, 1)."""
    theta = 2 * np.pi * t
    r = 1.0 + PETAL_AMPLITUDE * np.cos(3.0 * theta)
    return r * np.cos(theta), r * np.sin(theta)


# One call: parametrized boundary curve -> block-structured MacroMesh.
mm = macro_mesh_from_curve(flower, order=ORDER, mesh_size=0.7, n_boundary=400)
print(mm)

# %% Drawing helpers.
_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))


centroids = np.stack(
    [
        np.asarray(mm.cell_mapping(i).forward(jnp.array([0.5, 0.5])))
        for i in range(mm.n_cells)
    ]
)

# %% Plot: (left) mesh + mappings, (right) macro adjacency graph.
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 7))
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)
    ax1.fill(bnd[:, 0], bnd[:, 1], color=col, alpha=0.18)
    ax1.plot(bnd[:, 0], bnd[:, 1], color=col, lw=0.8)
    grid = np.asarray(jax.vmap(m.forward)(_REF_GRID))
    ax1.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)])
    )
    ax1.plot(seg[:, 0], seg[:, 1], "r-", lw=2.0)
ax1.set_title(
    f"3-petal flower, order {ORDER}: {mm.n_cells} macro-cells\n"
    "blue=bilinear, orange=curved, red=boundary"
)
ax1.set_aspect("equal")

for a, b in mm.cell_edges:
    ax2.plot(*zip(centroids[a], centroids[b]), "-", color="0.6", lw=0.8, zorder=1)
ax2.scatter(centroids[:, 0], centroids[:, 1], s=45, c="#54A24B", zorder=2)
ax2.set_title(
    f"Macro-cell adjacency graph\n{mm.n_cells} patches, {mm.n_interfaces} interfaces"
)
ax2.set_aspect("equal")

out = Path(__file__).with_name("macro_mesh_flower.png")
fig.tight_layout()
fig.savefig(out, dpi=140)
print(f"\nSaved figure to {out}")

plt.show()