"""Check empirically whether all proper 3-edge-colorings of a cubic plane
graph form a single Kempe equivalence class.

For each test graph:
  1. enumerate all proper 3-edge-colorings;
  2. build a multigraph K whose vertices are the colorings and whose edges
     connect pairs related by a single Kempe swap (swapping two colors on
     one cycle of the 2-color subgraph);
  3. report the number of connected components of K (= number of Kempe
     classes).

Test inputs:
  - Dodecahedron's reduced dual (16 v, 24 e), built as in
    check_dodecahedron_kempe.py;
  - The n=14 reduced dual found by search_kempe_property.py (20 v, 30 e).
"""
from sage.all import Graph
from sage.graphs.graph_generators import graphs
import sys


def all_proper_3_edge_colorings(G):
    edges = list(G.edges(labels=False))
    n = len(edges)
    adj = [[] for _ in range(n)]
    for i in range(n):
        u, v = edges[i][0], edges[i][1]
        for j in range(i):
            x, y = edges[j][0], edges[j][1]
            if u in (x, y) or v in (x, y):
                adj[i].append(j)
                adj[j].append(i)
    coloring = [-1] * n
    results = []

    def back(k):
        if k == n:
            results.append(tuple(coloring))
            return
        for c in range(3):
            if all(coloring[j] != c for j in adj[k]):
                coloring[k] = c
                back(k + 1)
        coloring[k] = -1
    back(0)
    return edges, results


def kempe_cycles_of(edges, coloring, color_pair):
    """List of connected components (as sorted tuples of edge indices) in
    the subgraph of edges with color in color_pair."""
    a, b = color_pair
    in_sub = [i for i in range(len(edges)) if coloring[i] in (a, b)]
    in_set = set(in_sub)
    visited = set()
    components = []
    for start in in_sub:
        if start in visited:
            continue
        comp = []
        stack = [start]
        visited.add(start)
        while stack:
            cur = stack.pop()
            comp.append(cur)
            u, v = edges[cur][0], edges[cur][1]
            for j in in_set:
                if j in visited:
                    continue
                x, y = edges[j][0], edges[j][1]
                if u in (x, y) or v in (x, y):
                    visited.add(j)
                    stack.append(j)
        components.append(tuple(sorted(comp)))
    return components


def swap_on_cycle(coloring, cycle_edges, color_pair):
    a, b = color_pair
    swap = list(coloring)
    for i in cycle_edges:
        if swap[i] == a:
            swap[i] = b
        elif swap[i] == b:
            swap[i] = a
    return tuple(swap)


def kempe_components(G):
    edges, colorings = all_proper_3_edge_colorings(G)
    print(f"  {len(colorings)} proper 3-edge-colorings")
    idx = {c: i for i, c in enumerate(colorings)}
    parent = list(range(len(colorings)))

    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(x, y):
        rx, ry = find(x), find(y)
        if rx != ry:
            parent[rx] = ry

    for col in colorings:
        for a, b in [(0, 1), (0, 2), (1, 2)]:
            for cyc in kempe_cycles_of(edges, col, (a, b)):
                new_col = swap_on_cycle(col, cyc, (a, b))
                if new_col == col:
                    continue
                if new_col in idx:
                    union(idx[col], idx[new_col])
    n_classes = len({find(i) for i in range(len(colorings))})
    print(f"  Kempe equivalence classes: {n_classes}")
    return n_classes


def dodecahedron_reduced_dual():
    edges = []
    for k in range(5):
        edges.append((('b', k), ('c', k)))
        edges.append((('b', k), ('c', (k - 1) % 5)))
        edges.append((('c', k), ('d', k)))
        edges.append((('d', k), ('d', (k + 1) % 5)))
    edges.append(('v_n', ('b', 0)))
    edges.append(('v_n', ('b', 1)))
    edges.append(('v_n', ('b', 2)))
    edges.append((('b', 3), ('b', 4)))
    return Graph(edges, multiedges=False, loops=False)


def n14_reduced_dual():
    # First min-degree-5 plantri triangulation on 14 vertices, reduced at the
    # first pentagonal face with i_red = 1 (matches search_kempe_property.py).
    G = next(graphs.triangulations(14, minimum_degree=5))
    G.is_planar(set_embedding=True)
    faces = G.faces()
    edge_to_faces = {}
    for fi, face in enumerate(faces):
        for u, v in face:
            edge_to_faces.setdefault(frozenset((u, v)), []).append(fi)
    dual_edges = [(fs[0], fs[1]) for fs in edge_to_faces.values() if len(fs) == 2]
    D = Graph(dual_edges, multiedges=False, loops=False)
    D.is_planar(set_embedding=True)
    # use first pentagonal face, i_red = 1
    face = next(f for f in D.faces() if len(f) == 5)
    boundary = [u for (u, v) in face]
    A = []
    for B_k in boundary:
        outer = [w for w in D.neighbor_iterator(B_k) if w not in boundary]
        A.append(outer[0])
    H = D.copy()
    for v in boundary:
        H.delete_vertex(v)
    vn = '__v_n__'
    H.add_vertex(vn)
    H.add_edge(vn, A[1])             # spike
    H.add_edge(vn, A[0])             # side_0
    H.add_edge(vn, A[2])             # side_1
    H.add_edge(A[3], A[4])           # merged
    H.is_planar(set_embedding=True)
    return H


def main():
    print("Dodecahedron's reduced dual:")
    G1 = dodecahedron_reduced_dual()
    kempe_components(G1)

    print()
    print("n=14 reduced dual (first plantri triangulation, i_red=1):")
    G2 = n14_reduced_dual()
    kempe_components(G2)

    print()
    print("Small sanity check: K_4")
    K4 = graphs.CompleteGraph(4)
    kempe_components(K4)


if __name__ == '__main__':
    main()