"""Plane diamond coloring 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`
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


def search_counterexample_comprehensive(max_order: int, min_order: int = 4) -> list[Graph]:
    """
    Iterate through every maximal planar graph of order in [min_order, max_order]
    and return all those without a plane diamond coloring.
    """
    counterexamples: list[Graph] = []
    for n in range(min_order, max_order + 1):
        checked = 0
        for g in graphs.planar_graphs(n, minimum_connectivity=3, maximum_face_size=3):
            checked += 1
            if not has_plane_diamond_coloring(g):
                print(f"Counterexample at order {n} (graph #{checked}): {g.graph6_string()}")
                counterexamples.append(g)
            if checked % 100 == 0:
                print(f"  order {n}: checked {checked} graphs, {len(counterexamples)} counterexamples so far")
        print(f"order {n} done: {checked} triangulations checked")
    return counterexamples


def search_min_degree_counterexample_comprehensive(max_order: int, minimum_degree: int, min_order: int = 4) -> list[Graph]:
    """
    Iterate through every maximal planar graph of order in [min_order, max_order]
    with the given minimum degree, and return all those without a plane diamond
    coloring.
    """
    counterexamples: list[Graph] = []
    for n in range(min_order, max_order + 1):
        checked = 0
        for g in graphs.planar_graphs(n, minimum_connectivity=3, maximum_face_size=3, minimum_degree=minimum_degree):
            checked += 1
            if not has_plane_diamond_coloring(g):
                print(f"Counterexample at order {n}, min_degree {minimum_degree} (graph #{checked}): {g.graph6_string()}")
                counterexamples.append(g)
            if checked % 100 == 0:
                print(f"  order {n}: checked {checked} graphs, {len(counterexamples)} counterexamples so far")
        print(f"order {n} done: {checked} triangulations of min degree {minimum_degree} checked")
    return counterexamples


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "comprehensive":
        max_order = int(sys.argv[2]) if len(sys.argv) > 2 else 10
        min_order = int(sys.argv[3]) if len(sys.argv) > 3 else 4
        counterexamples = search_counterexample_comprehensive(max_order, min_order)
        print(f"Found {len(counterexamples)} counterexamples in orders {min_order}..{max_order}")
        for g in counterexamples:
            canonize_and_save_graph(g)
    elif len(sys.argv) > 1 and sys.argv[1] == "min-degree":
        max_order = int(sys.argv[2]) if len(sys.argv) > 2 else 13
        minimum_degree = int(sys.argv[3]) if len(sys.argv) > 3 else 5
        min_order = int(sys.argv[4]) if len(sys.argv) > 4 else 4
        counterexamples = search_min_degree_counterexample_comprehensive(max_order, minimum_degree, min_order)
        print(f"Found {len(counterexamples)} counterexamples in orders {min_order}..{max_order} with min degree {minimum_degree}")
        for g in counterexamples:
            canonize_and_save_graph(g)
    else:
        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}")
