"""Visualise cell and face classification for a circle level-set on a Cartesian mesh."""

import time

import matplotlib.collections as mc
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np

from scimba_jax.linear_approximation.meshes.levelset_classifier import (
    LevelSetClassifier,
)
from scimba_jax.linear_approximation.meshes.mesh import Mesh
from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized
from scimba_jax.mapping.mapping import InvertibleFunction, Mapping

# ── Mesh & level-set setup ──────────────────────────────────────────────────

n = 20  # cells per direction
mapping_id = InvertibleFunction(lambda x: x, lambda y: y)
mesh = Mesh(
    dim=2,
    n_cells=[n, n],
    ref_quad=UnitSquareTensorized(dim=2, order=3),
    mapping=Mapping(mappings=[mapping_id]),
)
print(mesh.h)
cx, cy, r = 0.5, 0.5, 2 ** (1 / 2) / 4


def phi(x):
    return (x[0] - cx) ** 2 + (x[1] - cy) ** 2 - r**2


t0 = time.perf_counter()
classifier = LevelSetClassifier(mesh, phi)
print(
    "len(self.interior_cells) + len(self.cut_cells)",
    len(classifier.interior_cells) + len(classifier.cut_cells),
)
t1 = time.perf_counter()
print(f"LevelSetClassifier init : {t1 - t0:.4f} s")

print("=== Cell classification ===")
t0 = time.perf_counter()
interior_cells = classifier.interior_cells
cut_cells = classifier.cut_cells
exterior_cells = classifier.exterior_cells
t1 = time.perf_counter()
print(f"  Interior cells        : {len(interior_cells)}")
print(f"  Cut cells             : {len(cut_cells)}")
print(f"  Exterior cells        : {len(exterior_cells)}")
print(f"  Time                  : {t1 - t0:.4f} s")

print("\n=== Face classification ===")
t0 = time.perf_counter()
interior_adj_cut_faces = classifier.interior_adj_cut_faces
exterior_adj_cut_faces = classifier.exterior_adj_cut_faces
interior_faces = classifier.interior_faces
exterior_faces = classifier.exterior_faces
cut_faces = classifier.cut_faces
t1 = time.perf_counter()
print(f"  F_INTERIOR_ADJ_CUT (1): {len(interior_adj_cut_faces)}")
print(f"  F_EXTERIOR_ADJ_CUT (2): {len(exterior_adj_cut_faces)}")
print(f"  F_INTERIOR         (3): {len(interior_faces)}")
print(f"  F_EXTERIOR         (4): {len(exterior_faces)}")
print(f"  F_CUT              (5): {len(cut_faces)}")
print(f"  Time                  : {t1 - t0:.4f} s")

# ── Colour maps ────────────────────────────────────────────────────────────

CELL_COLORS = {
    LevelSetClassifier.INTERIOR: "#4CAF50",  # green
    LevelSetClassifier.CUT: "#FF9800",  # orange
    LevelSetClassifier.EXTERIOR: "#BBDEFB",  # light blue
}

FACE_COLORS = {
    LevelSetClassifier.F_INTERIOR_ADJ_CUT: "#1B5E20",  # dark green
    LevelSetClassifier.F_EXTERIOR_ADJ_CUT: "#0D47A1",  # dark blue
    LevelSetClassifier.F_INTERIOR: "#81C784",  # medium green
    LevelSetClassifier.F_EXTERIOR: "#90CAF9",  # medium blue
    LevelSetClassifier.F_CUT: "#E53935",  # red
}

FACE_LW = {
    LevelSetClassifier.F_INTERIOR_ADJ_CUT: 2.5,
    LevelSetClassifier.F_EXTERIOR_ADJ_CUT: 2.5,
    LevelSetClassifier.F_INTERIOR: 0.8,
    LevelSetClassifier.F_EXTERIOR: 0.8,
    LevelSetClassifier.F_CUT: 2.5,
}

# ── Plot ───────────────────────────────────────────────────────────────────

fig, (ax_cells, ax_faces) = plt.subplots(1, 2, figsize=(13, 6))

nx, ny = int(mesh.n_cells[0]), int(mesh.n_cells[1])
hx, hy = 1.0 / nx, 1.0 / ny

theta = np.linspace(0, 2 * np.pi, 300)
circle_x = cx + r * np.cos(theta)
circle_y = cy + r * np.sin(theta)

# ── Left plot: cells ────────────────────────────────────────────────────────
for cell_idx in range(mesh.n_cells_total):
    row = cell_idx // ny
    col = cell_idx % ny
    x0, y0 = col * hx, row * hy
    tag = int(classifier.cell_tags[cell_idx])
    rect = mpatches.FancyBboxPatch(
        (x0, y0),
        hx,
        hy,
        boxstyle="square,pad=0",
        linewidth=0.4,
        edgecolor="k",
        facecolor=CELL_COLORS[tag],
        alpha=0.8,
    )
    ax_cells.add_patch(rect)

ax_cells.plot(circle_x, circle_y, "k-", lw=1.5)
ax_cells.set_xlim(0, 1)
ax_cells.set_ylim(0, 1)
ax_cells.set_aspect("equal")
ax_cells.set_title(f"Cell tags — {n}×{n} mesh")

legend_cells = [
    mpatches.Patch(
        color=CELL_COLORS[LevelSetClassifier.INTERIOR], label="Interior (1)"
    ),
    mpatches.Patch(color=CELL_COLORS[LevelSetClassifier.CUT], label="Cut (0)"),
    mpatches.Patch(
        color=CELL_COLORS[LevelSetClassifier.EXTERIOR], label="Exterior (-1)"
    ),
]
ax_cells.legend(handles=legend_cells, loc="upper right", fontsize=8)

# ── Right plot: faces ────────────────────────────────────────────────────────

# Pre-compute endpoints for all faces
all_face_idx = np.arange(mesh.n_faces)
endpoints = classifier.face_endpoints(all_face_idx)  # (n_faces, 2, 2)

# Group faces by tag and draw as LineCollections (efficient)
for tag, color in FACE_COLORS.items():
    mask = classifier.face_tags == tag
    if not np.any(mask):
        continue
    segs = endpoints[mask]  # (k, 2, 2)
    segs_list = [(segs[i, 0], segs[i, 1]) for i in range(len(segs))]
    lc = mc.LineCollection(segs_list, colors=color, linewidths=FACE_LW[tag], alpha=0.9)
    ax_faces.add_collection(lc)

ax_faces.plot(circle_x, circle_y, "k-", lw=1.5)
ax_faces.set_xlim(0, 1)
ax_faces.set_ylim(0, 1)
ax_faces.set_aspect("equal")
ax_faces.set_title(f"Face tags — {n}×{n} mesh")

legend_faces = [
    mpatches.Patch(
        color=FACE_COLORS[LevelSetClassifier.F_INTERIOR_ADJ_CUT],
        label="Interior adj. cut (1)",
    ),
    mpatches.Patch(
        color=FACE_COLORS[LevelSetClassifier.F_EXTERIOR_ADJ_CUT],
        label="Exterior adj. cut (2)",
    ),
    mpatches.Patch(
        color=FACE_COLORS[LevelSetClassifier.F_INTERIOR], label="Interior (3)"
    ),
    mpatches.Patch(
        color=FACE_COLORS[LevelSetClassifier.F_EXTERIOR], label="Exterior (4)"
    ),
    mpatches.Patch(color=FACE_COLORS[LevelSetClassifier.F_CUT], label="Cut (5)"),
]
ax_faces.legend(handles=legend_faces, loc="upper right", fontsize=8)

plt.tight_layout()
plt.show()