"""Generate the labelled quadrilateral sequencing example figure for the paper.

Renders the deep embedding G' of a small maximal planar graph, fills each
quadrilateral with a color encoding its type, and overlays the index it
occupies in the canonical sequence Q_1, ..., Q_N. The output is a PDF placed
next to paper.tex.
"""
import argparse
from pathlib import Path
from typing import Any

import matplotlib
matplotlib.use("Agg")
import matplotlib.patches as patches  # noqa: E402  pylint: disable=wrong-import-position
import matplotlib.pyplot as plt  # noqa: E402  pylint: disable=wrong-import-position

from sage.all import Graph, graphs  # type: ignore[attr-defined]  # pylint: disable=no-name-in-module,wrong-import-position  # noqa: E402
from sage.misc.randstate import set_random_seed  # type: ignore[import-not-found]  # noqa: E402

from lib.tutte_embedding import tutte_embedding  # noqa: E402
from plane_depth_sequencing import (  # noqa: E402
    _level_edge_of_face,
    _quad_type,
    _quad_vertices,
    quadrilateral_sequencing,
)


def _planar_pos(
    g_prime: Graph,
    outer_face: list[Any],
    layout: str,
) -> dict[Any, tuple[float, float]]:
    """Compute positions for g_prime under either 'tutte' or 'sage_planar' layout."""
    if layout == "tutte":
        return tutte_embedding(g_prime, outer_face)
    if layout == "sage_planar":
        g_prime.is_planar(set_embedding=True, set_pos=True)
        pos = g_prime.get_pos()
        assert pos is not None, "Sage failed to set planar positions"
        return {v: (float(p[0]), float(p[1])) for v, p in pos.items()}
    raise ValueError(f"Unknown layout: {layout}")


MOVE_NAMES = {0: "AD", 1: "LA", 2: "J", 3: "RC"}
TYPE_COLORS = {
    "shallow_diamond": "#FFE0B2",
    "deep_diamond": "#B2DFDB",
    "s_quad": "#F8BBD0",
}


def _ordered_quad_vertices(quad: frozenset, depth_labelling: dict[Any, int]) -> list[Any]:
    """Return the 4 quad vertices in cyclic order: e1, a, e2, b."""
    f1, f2 = list(quad)
    level_edge = _level_edge_of_face(f1, depth_labelling)
    e1, e2 = list(level_edge)
    a = next(v for v in f1 if v not in level_edge)
    b = next(v for v in f2 if v not in level_edge)
    return [e1, a, e2, b]


def _pick_outer_cap_face(
    g_prime: Graph,
    outer_cap_vertex: Any,
    outer_cycle: list[Any],
) -> list[Any]:
    """Return the vertex list of an outer-cap face (the face used as outer in the plane drawing)."""
    g_prime.is_planar(set_embedding=True)
    embedding = g_prime.get_embedding()
    outer_set = set(outer_cycle)
    for face in g_prime.faces(embedding):
        verts = [u for u, _ in face]
        if outer_cap_vertex in verts and len(set(verts) & outer_set) == 2:
            return verts
    raise RuntimeError("No outer-cap face found in G' embedding")


def _draw_figure(
    g: Graph,
    outer_cycle: list[Any],
    out_path: Path,
    layout: str = "sage_planar",
) -> dict[str, Any]:
    """Render and save the labelled sequencing figure. Returns summary stats."""
    result = quadrilateral_sequencing(g, outer_cycle)
    g_prime: Graph = result["deep_embedding"]
    depth_labelling: dict[Any, int] = result["depth_labelling"]
    sequence = result["sequence"]
    move_codes = result["move_codes"]
    outer_cap_vertex = result["outer_cap_vertex"]

    outer_face = _pick_outer_cap_face(g_prime, outer_cap_vertex, outer_cycle)
    pos = _planar_pos(g_prime, outer_face, layout)

    fig, ax = plt.subplots(figsize=(9, 9))

    for i, quad in enumerate(sequence):
        ordered = _ordered_quad_vertices(quad, depth_labelling)
        poly_pts = [pos[v] for v in ordered]
        qt = _quad_type(quad, depth_labelling)
        polygon = patches.Polygon(
            poly_pts,
            closed=True,
            facecolor=TYPE_COLORS[qt],
            edgecolor="none",
            alpha=0.65,
            zorder=1,
        )
        ax.add_patch(polygon)

    for u, v in g_prime.edges(labels=False):
        x1, y1 = pos[u]
        x2, y2 = pos[v]
        if depth_labelling[u] == depth_labelling[v]:
            ax.plot([x1, x2], [y1, y2], linestyle=(0, (3, 2)), color="#666666", linewidth=1.0, zorder=2)
        else:
            ax.plot([x1, x2], [y1, y2], "-", color="black", linewidth=1.0, zorder=2)

    outer_set = set(outer_cycle)
    for v, (x, y) in pos.items():
        if v == outer_cap_vertex:
            color = "#D32F2F"
        elif v in outer_set:
            color = "#1976D2"
        else:
            color = "black"
        ax.scatter([x], [y], s=28, color=color, zorder=5, edgecolors="white", linewidths=0.8)
        name = "$x^{*}$" if v == outer_cap_vertex else str(v)
        ax.annotate(
            name,
            (x, y),
            xytext=(7, 6),
            textcoords="offset points",
            fontsize=9,
            zorder=8,
            color="#222222",
            bbox={"boxstyle": "round,pad=0.05", "facecolor": "white", "edgecolor": "none", "alpha": 0.85},
        )

    for i, quad in enumerate(sequence):
        ordered = _ordered_quad_vertices(quad, depth_labelling)
        poly_pts = [pos[v] for v in ordered]
        cx = sum(p[0] for p in poly_pts) / 4
        cy = sum(p[1] for p in poly_pts) / 4
        if i == 0:
            label = "$Q_{1}$"
        else:
            label = f"$Q_{{{i + 1}}}^{{\\mathrm{{{MOVE_NAMES[move_codes[i - 1]]}}}}}$"
        ax.text(
            cx,
            cy,
            label,
            ha="center",
            va="center",
            fontsize=12,
            zorder=7,
            bbox={"boxstyle": "round,pad=0.22", "facecolor": "white", "edgecolor": "#444444", "alpha": 0.95},
        )

    legend_handles = [
        patches.Patch(color=TYPE_COLORS["shallow_diamond"], label="shallow diamond"),
        patches.Patch(color=TYPE_COLORS["deep_diamond"], label="deep diamond"),
        patches.Patch(color=TYPE_COLORS["s_quad"], label="S quad"),
    ]
    ax.legend(handles=legend_handles, loc="lower right", fontsize=10, framealpha=0.95)

    ax.set_aspect("equal")
    ax.axis("off")
    plt.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(out_path, format="pdf", bbox_inches="tight")
    plt.close(fig)

    return {
        "n_vertices_g": g.order(),
        "n_vertices_gprime": g_prime.order(),
        "n_quads": len(sequence),
        "move_codes": move_codes,
        "outer_cycle": outer_cycle,
        "outer_cap_vertex": outer_cap_vertex,
        "depth_labelling": depth_labelling,
    }


def _build_example(seed: int, n: int) -> tuple[Graph, list[Any]]:
    """Build a random triangulation and pick a deterministic outer cycle."""
    set_random_seed(seed)
    g = graphs.RandomTriangulation(n)
    g.is_planar(set_embedding=True)
    embedding = g.get_embedding()
    faces = g.faces(embedding)
    outer_cycle = [u for u, _ in faces[0]]
    return g, outer_cycle


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seed", type=int, default=141)
    parser.add_argument("--n", type=int, default=9)
    parser.add_argument(
        "--out",
        type=Path,
        default=Path("papers/plane_depth_sequencing/example_figure.pdf"),
    )
    parser.add_argument("--layout", choices=("tutte", "sage_planar"), default="sage_planar")
    args = parser.parse_args()

    g, outer_cycle = _build_example(args.seed, args.n)
    summary = _draw_figure(g, outer_cycle, args.out, layout=args.layout)

    print(f"Wrote {args.out}")
    print(f"  |V(G)|={summary['n_vertices_g']}, |V(G')|={summary['n_vertices_gprime']}, "
          f"N={summary['n_quads']}")
    print(f"  outer cycle: {summary['outer_cycle']}, x* = {summary['outer_cap_vertex']}")
    code_str = "".join(str(c) for c in summary["move_codes"])
    print(f"  move-code string: {code_str}")


if __name__ == "__main__":
    main()