"""Example: colored pentagon reduction on a random 20-vertex triangulation."""
import os, sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")))  # repo root for `lib`
from collections import defaultdict
from pathlib import Path
from typing import Any, cast, TypedDict, Literal
from sage.all import graphs, Graph  # type: ignore[attr-defined]  # pylint: disable=no-name-in-module
from lib.planar_embedding import get_embedding_from_pos
from lib.colored_graphs import ColoredGraphId, VertexColoring, save_colored_graph
from lib.operations import Operation, save_operation_sequence

DIR = Path(__file__).parent

def _neighbors_form_cycle(g: Graph, v: Any) -> bool:
    """Return True if the neighbors of v induce a cycle in g."""
    return bool(cast(Graph, g.subgraph(g.neighbors(v))).is_cycle())

class PluckMeta(TypedDict):
    """Meta information about the pluck operation"""
    v0: Any

class PluckOperation(Operation):
    """Info about an operation in which a vertex v0 and its incident edges is removed from G"""
    name: Literal['pluck']
    meta: PluckMeta

def pluck(g: Graph, coloring: VertexColoring, v0: Any) -> tuple[Graph, VertexColoring]:
    """Delete v0 and all its incident edges from g"""
    g_prime = g.copy()
    g_prime.delete_vertex(v0)
    if (pos := g.get_pos()) is not None:
        g_prime.set_pos({v: p for v, p in pos.items() if v != v0})
        g_prime.set_embedding(get_embedding_from_pos(g_prime))
    coloring_prime = coloring.copy()
    del coloring_prime[v0]
    return g_prime, coloring_prime


class SquishMeta(TypedDict):
    """Meta information about the squish operation"""
    v0: Any
    v_merged: list[Any]

class SquishOperation(Operation):
    """
    Info about an operation in which two same colored neighbors of a vertex v0 are merged
    into v0
    """
    name: Literal['squish']
    meta: SquishMeta

def squish(g: Graph, coloring: VertexColoring, v0: Any) -> tuple[Graph, VertexColoring, Any, Any]:
    """
    Contract two same-colored neighbors of v0 into v0 and return a new valid
    coloring along with the new graph.
    NOTE: assumes g is a maximal planar graph
    """
    neighbor_by_color: defaultdict[Any, list[Any]] = defaultdict(list)
    for v in g.neighbors(v0):
        neighbor_by_color[coloring[v]].append(v)

    v1, v2 = next(
        (vs[0], vs[1]) for vs in neighbor_by_color.values() if len(vs) >= 2
    )

    g_prime = g.copy()
    g_prime.merge_vertices([v0, v1, v2])
    if (pos := g.get_pos()) is not None:
        g_prime.set_pos({v: p for v, p in pos.items() if v not in (v1, v2)})
        g_prime.set_embedding(get_embedding_from_pos(g_prime))
    coloring_prime = {v: c for v, c in coloring.items() if v not in (v1, v2)}
    coloring_prime[v0] = coloring[v1]
    return g_prime, coloring_prime, v1, v2


ReduceOperation = SquishOperation | PluckOperation

def reduce(
    g: Graph,
    coloring: VertexColoring,
    source_canon_id: ColoredGraphId,
    op_sequence: list[ReduceOperation] | None = None,
) -> list[ReduceOperation]:
    """Repeatedly apply pluck/squish reductions until no candidates remain."""
    if op_sequence is None:
        op_sequence = []

    print(f"Coloring: {coloring}")

    degree_4_candidates: list[Any] = []
    degree_5_candidates: list[Any] = []

    for v in g.vertices():
        if g.degree(v) == 3 and _neighbors_form_cycle(g, v):
            result_graph, result_coloring = pluck(g, coloring, v)
            print(f"\nG' (after pluck v0={v}): {result_graph.order()} vertices, {result_graph.size()} edges")
            _, _, result_canon_id = save_colored_graph(result_graph, result_coloring)
            op_sequence.append(PluckOperation(name='pluck', meta=PluckMeta(v0=v), source_graph=g, source_canon_id=source_canon_id, result_graph=result_graph, result_coloring=result_coloring, result_canon_id=result_canon_id))
            return reduce(result_graph, result_coloring, result_canon_id, op_sequence)
        if g.degree(v) == 4 and _neighbors_form_cycle(g, v):
            degree_4_candidates.append(v)
        elif g.degree(v) == 5 and _neighbors_form_cycle(g, v):
            degree_5_candidates.append(v)

    if degree_4_candidates:
        v0 = degree_4_candidates[0]
        result_graph, result_coloring, v1, v2 = squish(g, coloring, v0)
        print(f"Shared-color neighbors: v1={v1}, v2={v2}  (color {coloring[v1]})")
        print(f"\nG' (after squish v0={v0}): {result_graph.order()} vertices, {result_graph.size()} edges")
        _, _, result_canon_id = save_colored_graph(result_graph, result_coloring)
        op_sequence.append(SquishOperation(name='squish', meta=SquishMeta(v0=v0, v_merged=[v1, v2]), source_graph=g, source_canon_id=source_canon_id, result_graph=result_graph, result_coloring=result_coloring, result_canon_id=result_canon_id))
        return reduce(result_graph, result_coloring, result_canon_id, op_sequence)

    if degree_5_candidates:
        v0 = degree_5_candidates[0]
        result_graph, result_coloring, v1, v2 = squish(g, coloring, v0)
        print(f"Shared-color neighbors: v1={v1}, v2={v2}  (color {coloring[v1]})")
        print(f"\nG' (after squish v0={v0}): {result_graph.order()} vertices, {result_graph.size()} edges")
        _, _, result_canon_id = save_colored_graph(result_graph, result_coloring)
        op_sequence.append(SquishOperation(name='squish', meta=SquishMeta(v0=v0, v_merged=[v1, v2]), source_graph=g, source_canon_id=source_canon_id, result_graph=result_graph, result_coloring=result_coloring, result_canon_id=result_canon_id))
        return reduce(result_graph, result_coloring, result_canon_id, op_sequence)

    print("DONE")
    return op_sequence


G = next(graphs.planar_graphs(20, minimum_degree=5))
print(f"G: {G.order()} vertices, {G.size()} edges")
print(f"Degree sequence: {sorted(G.degree_sequence(), reverse=True)}")
starting_coloring_classes = G.coloring()
starting_coloring = {v: i for i, cls in enumerate(starting_coloring_classes) for v in cls}
_, _, initial_canon_id = save_colored_graph(G, starting_coloring)
G.is_planar(set_embedding=True, set_pos=True)

op_sequence = reduce(G, starting_coloring, initial_canon_id)
save_operation_sequence(op_sequence, G, starting_coloring, DIR)
