"""Plane diamond coloring on maximal planar graphs."""
from typing import Any
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


def get_plane_diamond_scaffold(g: Graph, v: Any) -> Graph:
    """
    Return the diamond scaffold of g relative to v.

    The diamond scaffold is the spanning subgraph of g obtained by removing
    every edge whose endpoints lie in the same level of the distance
    partition of g from v.
    """
    distances = dict(g.breadth_first_search(v, report_distance=True))
    scaffold = g.copy()
    scaffold.delete_edges([(x, y) for x, y in g.edges(labels=False) if distances[x] == distances[y]])
    return scaffold


def has_plane_diamond_coloring_at_root(g: Graph, u: Any) -> bool:
    """
    Return True iff g admits a proper 4-coloring with two color classes
    parity-separated by the BFS distance partition from u.

    Equivalent to 4-colorability of the auxiliary graph H_u obtained by
    adjoining vertices v_a, v_b to g, with v_a adjacent to every odd-parity
    layer vertex, v_b adjacent to every even-parity layer vertex, and a
    v_a v_b edge.
    """
    distances = dict(g.breadth_first_search(u, report_distance=True))
    odd_vertices = [v for v in g.vertices() if distances[v] % 2 == 1]
    even_vertices = [v for v in g.vertices() if distances[v] % 2 == 0]
    h = g.copy()
    v_a = max(g.vertices()) + 1
    v_b = v_a + 1
    h.add_vertex(v_a)
    h.add_vertex(v_b)
    h.add_edges([(v_a, w) for w in odd_vertices])
    h.add_edges([(v_b, w) for w in even_vertices])
    h.add_edge(v_a, v_b)
    return h.chromatic_number() <= 4


def has_plane_diamond_coloring(g: Graph) -> bool:
    """Return True iff some root vertex of g witnesses a plane diamond coloring."""
    return any(has_plane_diamond_coloring_at_root(g, u) for u in g.vertices())


def search_counterexample(n: int, num_trials: int) -> Graph | None:
    """
    Sample random maximal planar graphs of order n and return the first one
    with no plane diamond coloring, or None if none is found.
    """
    for trial in range(num_trials):
        g = graphs.RandomTriangulation(n)
        if not has_plane_diamond_coloring(g):
            print(f"Counterexample found on trial {trial + 1}")
            return g
        if (trial + 1) % 10 == 0:
            print(f"Checked {trial + 1}/{num_trials} graphs of order {n}, no counterexample yet")
    return None


if __name__ == "__main__":
    import sys
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 12
    num_trials = int(sys.argv[2]) if len(sys.argv) > 2 else 100
    counterexample = search_counterexample(n, num_trials)
    if counterexample is None:
        print(f"No counterexample found in {num_trials} random triangulations of order {n}")
    else:
        canonical, graph_dir = canonize_and_save_graph(counterexample)
        print(f"Counterexample saved to {graph_dir}")