"""Plane depth sequencing on maximal planar graphs."""
import os, sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")))  # repo root for `lib`
import random
from typing import Any, TypedDict
from sage.all import Graph, graphs  # type: ignore[attr-defined]  # pylint: disable=no-name-in-module
from lib.colored_graphs import canonize_and_save_graph


class DeeplyEmbeddedGraph(TypedDict):
    graph: Graph
    outer_cycle: list[Any]
    plane_depth_labelling: dict[Any, int]
    deep_embedding: Graph


class QuadrilateralSequence(TypedDict):
    deep_embedding: Graph
    triangular_faces: list[frozenset[Any]]
    depth_labelling: dict[Any, int]
    outer_cap_vertex: Any
    quadrilaterals: list[frozenset[frozenset[Any]]]
    sequence: list[frozenset[frozenset[Any]]]
    move_codes: list[int]


def get_plane_depth_labelling(g: Graph, outer_cycle: list[Any]) -> dict[Any, int]:
    """Return the plane depth of each vertex relative to the given outer cycle."""
    # equivalent to the commented out naive implementation:
    # return {v: min(g.distance(v, u) for u in outer_cycle) for v in g.vertices()}
    source = max(g.vertices()) + 1
    g_prime = g.copy()
    g_prime.add_vertex(source)
    g_prime.add_edges([(source, v) for v in outer_cycle])
    distances = g_prime.breadth_first_search(source, report_distance=True)
    return {v: d - 1 for v, d in distances if v != source}


def deep_embedding(g: Graph, outer_cycle: list[Any], plane_depth_labelling: dict[Any, int] | None = None) -> Graph:
    """
    Return the deep embedding of g relative to outer_cycle.

    For every interior triangular face whose three vertices all have the same
    plane depth, a new vertex is added adjacent to each vertex of that face.
    """
    if plane_depth_labelling is None:
        plane_depth_labelling = get_plane_depth_labelling(g, outer_cycle)
    outer_vertices = set(outer_cycle)

    embedding = g.get_embedding()
    if embedding is None:
        g.is_planar(set_embedding=True)
        embedding = g.get_embedding()

    g_prime = g.copy()
    new_vertex = max(g.vertices()) + 1

    for face in g.faces(embedding):
        face_vertices = [u for u, v in face]
        if set(face_vertices) == outer_vertices:
            continue
        if plane_depth_labelling[face_vertices[0]] == plane_depth_labelling[face_vertices[1]] == plane_depth_labelling[face_vertices[2]]:
            g_prime.add_vertex(new_vertex)
            g_prime.add_edges([(new_vertex, v) for v in face_vertices])
            new_vertex += 1

    return g_prime


def _triangle_type(face: frozenset[Any], depth_labelling: dict[Any, int]) -> str:
    """Return 'up' or 'down' for a triangular face (up: {d, d+1, d+1}; down: {d, d, d+1})."""
    depths = sorted(depth_labelling[v] for v in face)
    if depths[0] == depths[1]:
        return 'down'
    return 'up'


def _level_edge_of_face(face: frozenset[Any], depth_labelling: dict[Any, int]) -> frozenset[Any]:
    """Return the unique level edge of an up/down triangular face."""
    vs = list(face)
    for i in range(3):
        for j in range(i + 1, 3):
            if depth_labelling[vs[i]] == depth_labelling[vs[j]]:
                return frozenset([vs[i], vs[j]])
    raise ValueError(f"Face {set(face)} has no level edge (not up or down)")


def _quad_vertices(quad: frozenset[frozenset[Any]]) -> frozenset[Any]:
    """Return the 4 vertices of a quadrilateral."""
    f1, f2 = list(quad)
    return f1 | f2


def _quad_perimeter_edges(
    quad: frozenset[frozenset[Any]],
    depth_labelling: dict[Any, int],
) -> list[frozenset[Any]]:
    """Return the 4 perimeter edges (non-level) of a quadrilateral."""
    f1 = next(iter(quad))
    level_edge = _level_edge_of_face(f1, depth_labelling)
    edges: list[frozenset[Any]] = []
    for f in quad:
        vs = list(f)
        for i in range(3):
            edge = frozenset([vs[i], vs[(i + 1) % 3]])
            if edge != level_edge:
                edges.append(edge)
    return edges


def _quad_type(quad: frozenset[frozenset[Any]], depth_labelling: dict[Any, int]) -> str:
    """Return 'shallow_diamond', 'deep_diamond', or 's_quad' for a quadrilateral."""
    types = tuple(sorted(_triangle_type(f, depth_labelling) for f in quad))
    if types == ('up', 'up'):
        return 'shallow_diamond'
    if types == ('down', 'down'):
        return 'deep_diamond'
    return 's_quad'


def extended_deep_embedding(
    g: Graph,
    outer_cycle: list[Any],
    plane_depth_labelling: dict[Any, int] | None = None,
) -> tuple[Graph, list[frozenset[Any]], dict[Any, int], Any, dict[Any, list[Any]]]:
    """
    Return the extended deep embedding of g (including the outer face).

    Subdivides every neutral triangular face --- including the outer face,
    which is always neutral for a maximal planar graph --- by adding a new
    interior vertex. The vertex added inside the outer face is the
    outer-cap vertex x*.

    Returns (g_prime, faces, depth_labelling, outer_cap_vertex, embedding) where:
      - g_prime: the deep embedding graph G'
      - faces: list of triangular faces of G' (each a frozenset of 3 vertices),
        derived from Sage's planar embedding of g_prime
      - depth_labelling: depth of every vertex in G' (extends the original)
      - outer_cap_vertex: the new vertex placed inside the outer face
      - embedding: the CCW rotation system of g_prime (vertex -> ordered
        neighbors) as computed by Sage's planarity routine
    """
    if plane_depth_labelling is None:
        plane_depth_labelling = get_plane_depth_labelling(g, outer_cycle)
    outer_vertices = frozenset(outer_cycle)

    embedding_g = g.get_embedding()
    if embedding_g is None:
        g.is_planar(set_embedding=True)
        embedding_g = g.get_embedding()

    g_prime = g.copy()
    next_vertex = max(g.vertices()) + 1
    depth_labelling = dict(plane_depth_labelling)
    outer_cap_vertex: Any = None

    for face in g.faces(embedding_g):
        face_vertices = [u for u, v in face]
        a, b, c = face_vertices
        if depth_labelling[a] == depth_labelling[b] == depth_labelling[c]:
            x = next_vertex
            next_vertex += 1
            g_prime.add_vertex(x)
            g_prime.add_edges([(x, a), (x, b), (x, c)])
            if frozenset(face_vertices) == outer_vertices:
                outer_cap_vertex = x
                depth_labelling[x] = -1  # Outer-cap vertex has depth -1
            else:
                depth_labelling[x] = depth_labelling[a] + 1

    g_prime.is_planar(set_embedding=True)
    embedding = g_prime.get_embedding()
    assert embedding is not None, "g_prime must be planar after construction"
    faces = [frozenset([u for u, v in face]) for face in g_prime.faces(embedding)]

    return g_prime, faces, depth_labelling, outer_cap_vertex, embedding


def quadrilateral_decomposition(
    faces: list[frozenset[Any]],
    depth_labelling: dict[Any, int],
) -> tuple[list[frozenset[frozenset[Any]]], dict[frozenset[Any], frozenset[frozenset[Any]]]]:
    """
    Pair each face with the face on the other side of its level edge.

    Returns (quads, quad_of_face) where:
      - quads: list of {F1, F2} pairs (each a quadrilateral of the decomposition)
      - quad_of_face: dict mapping each face to its containing quad
    """
    faces_by_edge: dict[frozenset[Any], list[frozenset[Any]]] = {}
    for face in faces:
        vs = list(face)
        for i in range(3):
            edge = frozenset([vs[i], vs[(i + 1) % 3]])
            faces_by_edge.setdefault(edge, []).append(face)

    quads: list[frozenset[frozenset[Any]]] = []
    quad_of_face: dict[frozenset[Any], frozenset[frozenset[Any]]] = {}
    seen: set[frozenset[Any]] = set()
    for face in faces:
        if face in seen:
            continue
        level_edge = _level_edge_of_face(face, depth_labelling)
        adjacent = faces_by_edge[level_edge]
        assert len(adjacent) == 2, f"Level edge {set(level_edge)} has {len(adjacent)} adjacent faces"
        other = adjacent[0] if adjacent[1] == face else adjacent[1]
        quad = frozenset([face, other])
        quads.append(quad)
        quad_of_face[face] = quad
        quad_of_face[other] = quad
        seen.add(face)
        seen.add(other)

    return quads, quad_of_face


def _boundary_edges(
    slice_faces: set[frozenset[Any]],
    faces_by_edge: dict[frozenset[Any], list[frozenset[Any]]],
) -> set[frozenset[Any]]:
    """Return the set of edges that lie on the boundary of the slice."""
    boundary: set[frozenset[Any]] = set()
    for edge, fs in faces_by_edge.items():
        inside = sum(1 for f in fs if f in slice_faces)
        if inside == 1:
            boundary.add(edge)
    return boundary


def _face_of_wedge(
    embedding: dict[Any, list[Any]],
) -> dict[tuple[Any, Any, Any], frozenset[Any]]:
    """For each (v, n_a, n_b) where n_a, n_b are CCW-consecutive in v's rotation,
    return the triangular face at that wedge."""
    fow: dict[tuple[Any, Any, Any], frozenset[Any]] = {}
    for v, rotation in embedding.items():
        k = len(rotation)
        for i in range(k):
            n_a = rotation[i]
            n_b = rotation[(i + 1) % k]
            fow[(v, n_a, n_b)] = frozenset({v, n_a, n_b})
    return fow


def _find_starting_boundary_edge(
    slice_faces: set[frozenset[Any]],
    embedding: dict[Any, list[Any]],
    face_of_wedge: dict[tuple[Any, Any, Any], frozenset[Any]],
) -> tuple[Any, Any]:
    """Find a directed boundary edge (v, n_b) with the slice on its LEFT."""
    for v, rotation in embedding.items():
        k = len(rotation)
        for i in range(k):
            n_a = rotation[i]
            n_b = rotation[(i + 1) % k]
            n_c = rotation[(i + 2) % k]
            left_wedge = face_of_wedge[(v, n_a, n_b)]
            right_wedge = face_of_wedge[(v, n_b, n_c)]
            if left_wedge in slice_faces and right_wedge not in slice_faces:
                return (v, n_b)
    raise RuntimeError("No boundary edge found for slice")


def _boundary_walk(
    slice_faces: set[frozenset[Any]],
    embedding: dict[Any, list[Any]],
    face_of_wedge: dict[tuple[Any, Any, Any], frozenset[Any]],
) -> list[tuple[Any, Any]]:
    """Trace the slice's boundary CCW with the slice on the LEFT.

    Returns an ordered list of directed edges forming the closed boundary walk.
    """
    start = _find_starting_boundary_edge(slice_faces, embedding, face_of_wedge)
    walk: list[tuple[Any, Any]] = [start]
    curr_u, curr_v = start
    while True:
        rotation = embedding[curr_v]
        k = len(rotation)
        i = rotation.index(curr_u)
        next_v = None
        for j in range(1, k + 1):
            n_jm1 = rotation[(i + j - 1) % k]
            n_j = rotation[(i + j) % k]
            if face_of_wedge[(curr_v, n_jm1, n_j)] not in slice_faces:
                next_v = n_jm1
                break
        if next_v is None:
            raise RuntimeError(f"All wedges at {curr_v} are in slice; no boundary edge")
        next_edge = (curr_v, next_v)
        if next_edge == start:
            break
        walk.append(next_edge)
        curr_u, curr_v = next_edge
    return walk


def _canonicalize_walk(
    walk: list[tuple[Any, Any]],
    depth_labelling: dict[Any, int],
) -> list[tuple[Any, Any]]:
    """Rotate walk to start at the (smallest depth, smallest vertex id) position.

    This fixes the canonical 'top' of the slice for the top-to-bottom scan.
    """
    starts = [(depth_labelling[u], u) for u, _ in walk]
    canon_idx = min(range(len(walk)), key=lambda i: starts[i])
    return walk[canon_idx:] + walk[:canon_idx]


def _attachment_position(
    quad: frozenset[frozenset[Any]],
    walk: list[tuple[Any, Any]],
    depth_labelling: dict[Any, int],
) -> int:
    """Return the latest index in walk at which a perimeter edge of quad appears.

    Realizes 'bottommost attachment on the right boundary scanned top-to-bottom'.
    Returns -1 if the quad has no perimeter edge on the boundary walk.
    """
    perimeter = {frozenset(e) for e in _quad_perimeter_edges(quad, depth_labelling)}
    latest = -1
    for i, (u, v) in enumerate(walk):
        if frozenset([u, v]) in perimeter:
            latest = i
    return latest


def _boundary_deep_diamonds(
    quads: list[frozenset[frozenset[Any]]],
    outer_cycle: list[Any],
    outer_cap_vertex: Any,
) -> list[frozenset[frozenset[Any]]]:
    """Return the three boundary deep diamonds (each spans an edge of C)."""
    outer_set = set(outer_cycle)
    diamonds: list[frozenset[frozenset[Any]]] = []
    for q in quads:
        vs = _quad_vertices(q)
        if outer_cap_vertex in vs and len(vs & outer_set) == 2:
            diamonds.append(q)
    assert len(diamonds) == 3, f"Expected 3 boundary deep diamonds, got {len(diamonds)}"
    return diamonds


def _run_sequence(
    initial_quad: frozenset[frozenset[Any]],
    quads: list[frozenset[frozenset[Any]]],
    depth_labelling: dict[Any, int],
    faces_by_edge: dict[frozenset[Any], list[frozenset[Any]]],
    embedding: dict[Any, list[Any]],
    face_of_wedge: dict[tuple[Any, Any, Any], frozenset[Any]],
) -> tuple[list[frozenset[frozenset[Any]]], list[int]]:
    """Run the sequencing loop from initial_quad and return (sequence, move_codes)."""
    sequence: list[frozenset[frozenset[Any]]] = [initial_quad]
    move_codes: list[int] = []
    slice_quads: set[frozenset[frozenset[Any]]] = {initial_quad}
    slice_faces: set[frozenset[Any]] = set(initial_quad)

    while len(slice_quads) < len(quads):
        slice_v: set[Any] = set().union(*slice_faces)
        boundary = _boundary_edges(slice_faces, faces_by_edge)
        walk = _canonicalize_walk(
            _boundary_walk(slice_faces, embedding, face_of_wedge),
            depth_labelling,
        )

        anchor_drop: list[frozenset[frozenset[Any]]] = []
        level_add: list[frozenset[frozenset[Any]]] = []
        join: list[frozenset[frozenset[Any]]] = []
        ring_completion: list[frozenset[frozenset[Any]]] = []

        for q in quads:
            if q in slice_quads:
                continue
            vs = _quad_vertices(q)
            k = len(vs & slice_v)
            edges = _quad_perimeter_edges(q, depth_labelling)
            j = sum(1 for e in edges if e in boundary)
            qt = _quad_type(q, depth_labelling)

            if qt == 's_quad' and k == 2 and j == 1:
                anchor_drop.append(q)
            if k == 3 and j == 2:
                level_add.append(q)
            if qt == 'deep_diamond' and k == 2 and j == 1:
                join.append(q)
            if k == 4:
                ring_completion.append(q)

        def pick(cands: list[frozenset[frozenset[Any]]]) -> frozenset[frozenset[Any]]:
            return max(cands, key=lambda q: _attachment_position(q, walk, depth_labelling))

        if anchor_drop:
            next_quad = pick(anchor_drop)
            next_code = 0
        elif level_add:
            next_quad = pick(level_add)
            next_code = 1
        elif join:
            next_quad = pick(join)
            next_code = 2
        elif ring_completion:
            next_quad = pick(ring_completion)
            next_code = 3
        else:
            raise RuntimeError(
                f"No applicable move at step {len(sequence)}; slice has {len(slice_quads)}/{len(quads)} quads"
            )

        sequence.append(next_quad)
        move_codes.append(next_code)
        slice_quads.add(next_quad)
        slice_faces.update(next_quad)

    return sequence, move_codes


def quadrilateral_sequencing(
    g: Graph,
    outer_cycle: list[Any],
    plane_depth_labelling: dict[Any, int] | None = None,
) -> QuadrilateralSequence:
    """
    Build the quadrilateral sequence of g relative to outer_cycle.

    Constructs the extended deep embedding G' (including the outer-cap vertex
    x* in the outer face), decomposes G' into quadrilaterals via level-edge
    pairing, and produces the deterministic sequence Q_1, ..., Q_N by
    repeatedly applying the move-selection rule:

        anchor drop (0) > level add (1) > join (2) > ring completion (3).

    Within each move's candidate list, the bottommost attachment on the
    canonical boundary walk is selected (the largest index in the walk where
    a perimeter edge of the candidate appears). Among the three boundary deep
    diamonds, Q_1 is the start that produces the lexicographically smallest
    move-code string.

    Simplification: anchor drop / join detection uses (k, j) = (vertices in
    slice, perimeter edges on slice boundary) only; the orientation-specific
    'left edge = right edge' clause is not separately enforced.
    """
    g_prime, faces, depth_labelling, outer_cap_vertex, embedding = extended_deep_embedding(
        g, outer_cycle, plane_depth_labelling
    )
    quads, _ = quadrilateral_decomposition(faces, depth_labelling)

    faces_by_edge: dict[frozenset[Any], list[frozenset[Any]]] = {}
    for face in faces:
        vs = list(face)
        for i in range(3):
            edge = frozenset([vs[i], vs[(i + 1) % 3]])
            faces_by_edge.setdefault(edge, []).append(face)

    face_of_wedge = _face_of_wedge(embedding)
    candidates = _boundary_deep_diamonds(quads, outer_cycle, outer_cap_vertex)

    best_sequence: list[frozenset[frozenset[Any]]] | None = None
    best_codes: list[int] | None = None
    for q1 in candidates:
        seq, codes = _run_sequence(q1, quads, depth_labelling, faces_by_edge, embedding, face_of_wedge)
        if best_codes is None or codes < best_codes:
            best_sequence = seq
            best_codes = codes
    assert best_sequence is not None and best_codes is not None

    return QuadrilateralSequence(
        deep_embedding=g_prime,
        triangular_faces=faces,
        depth_labelling=depth_labelling,
        outer_cap_vertex=outer_cap_vertex,
        quadrilaterals=quads,
        sequence=best_sequence,
        move_codes=best_codes,
    )


def generate_example(n: int) -> DeeplyEmbeddedGraph:
    """Generate a random maximal planar graph of size n and return the triangulation, outer cycle, and deep embedding."""
    g = graphs.RandomTriangulation(n)
    g.is_planar(set_embedding=True)
    embedding = g.get_embedding()
    faces = g.faces(embedding)
    outer_cycle = [u for u, v in random.choice(faces)]
    plane_depth_labelling = get_plane_depth_labelling(g, outer_cycle)
    return DeeplyEmbeddedGraph(graph=g, outer_cycle=outer_cycle, plane_depth_labelling=plane_depth_labelling, deep_embedding=deep_embedding(g, outer_cycle, plane_depth_labelling))


if __name__ == "__main__":
    example = generate_example(10)
    result = quadrilateral_sequencing(example['graph'], example['outer_cycle'])
    canonical, graph_dir = canonize_and_save_graph(example['graph'])
    (graph_dir / "plane_depth_sequence").mkdir(parents=True, exist_ok=True)
    print(canonical)
    print(f"Number of quadrilaterals: {len(result['quadrilaterals'])}")
    print(f"Sequence length: {len(result['sequence'])}")
    print(f"Move codes: {result['move_codes']}")