"""Empirically test, on every chord-apex+Kempe colouring of every
reduced dual we can enumerate:

  (a) Does the Heawood-sum identity sum_v h_phi(v) = 0 hold?
  (b) Is h_phi ever constant on V(K_b) U V(K_c), the union of the two
      Kempe cycles through the merged edge?

(b) is the precise hypothesis of Lemma 5.3's conclusion. If (b) is
*never* witnessed on any chord-apex+Kempe colouring, then by the
contrapositive of Lemma 5.3 we have empirical verification of
Conjecture 5.1 on the chord-apex+Kempe surrogates.

Run with: sage experiments/check_heawood_on_kempe.py
"""
import os
import sys
import time

from sage.all import Graph
from sage.graphs.graph_generators import graphs

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

from check_conj_final_scaled import (
    apply_reduction,
    proper_3_edge_colorings,
    matches_chord_apex_kempe,
    kempe_cycle_set,
    edge_idx,
)


def dual_of(G):
    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)
    return Graph(
        [(fs[0], fs[1]) for fs in edge_to_faces.values() if len(fs) == 2],
        multiedges=False, loops=False)


def heawood_numbers(H, edges, col_list):
    """Return dict v -> h_phi(v) in {+1, -1}.

    The Heawood number is +1 iff the clockwise cyclic colour order at v
    is an even cyclic permutation of (0, 1, 2).
    """
    H.is_planar(set_embedding=True)
    emb = H.get_embedding()   # v -> list of neighbours in CW order
    edge_color = {frozenset(e): col_list[i] for i, e in enumerate(edges)}
    out = {}
    for v in H.vertex_iterator():
        nbrs = emb[v]
        if len(nbrs) != 3:
            raise RuntimeError(f"vertex {v} not cubic: {len(nbrs)}")
        cs = tuple(edge_color[frozenset((v, w))] for w in nbrs)
        # cs is a permutation of (0, 1, 2). Check parity of its cyclic
        # equivalence class compared to (0, 1, 2).
        # (0,1,2), (1,2,0), (2,0,1) are even cyclic -> +1
        # (0,2,1), (2,1,0), (1,0,2) are odd  cyclic -> -1
        if cs in [(0, 1, 2), (1, 2, 0), (2, 0, 1)]:
            out[v] = +1
        elif cs in [(0, 2, 1), (2, 1, 0), (1, 0, 2)]:
            out[v] = -1
        else:
            raise RuntimeError(f"bad colour tuple at {v}: {cs}")
    return out


def vertices_of_kempe(edges, kc_edge_indices):
    """Vertices touched by the edges with indices in kc_edge_indices."""
    vs = set()
    for i in kc_edge_indices:
        u, v = edges[i][0], edges[i][1]
        vs.add(u); vs.add(v)
    return vs


def test_one(D, name=""):
    """Run the empirical check on one cubic plane graph D = G'."""
    D.is_planar(set_embedding=True)
    n_col = 0
    n_sum_zero = 0
    n_constant_union = 0
    sum_distribution = {}
    examples = []   # (n_pent_face, i_red, coloring index, sum_h, constant?)

    for face in D.faces():
        if len(face) != 5: continue
        for i_red in range(5):
            res = apply_reduction(D, face, i_red, 9999)
            if res is None: continue
            H = res['H']; named = res['named']
            H.is_planar(set_embedding=True)
            edges, colorings = proper_3_edge_colorings(H)
            cand = [c for c in colorings
                    if matches_chord_apex_kempe(edges, c, named)]
            for col in cand:
                n_col += 1
                # Compute Heawood numbers
                try:
                    h = heawood_numbers(H, edges, col)
                except RuntimeError:
                    continue
                # Part (a): sum_v h(v)
                s = sum(h.values())
                sum_distribution[s] = sum_distribution.get(s, 0) + 1
                if s == 0:
                    n_sum_zero += 1
                # Part (b): is h constant on V(K_b) U V(K_c)?
                merged_idx = edge_idx(edges, named['merged'])
                a = col[merged_idx]
                K_unions = set()
                for b in range(3):
                    if b == a: continue
                    kc = kempe_cycle_set(edges, col, merged_idx, (a, b))
                    K_unions |= vertices_of_kempe(edges, kc)
                h_on_union = {h[v] for v in K_unions}
                constant = (len(h_on_union) == 1)
                if constant:
                    n_constant_union += 1
                    if len(examples) < 5:
                        examples.append((len(K_unions), s, list(h_on_union)[0]))
    return n_col, n_sum_zero, n_constant_union, sum_distribution, examples


def main(max_n=20, time_budget_per_n=1800):
    print(f"Checking Heawood identities on chord-apex+Kempe colourings, "
          f"n in [12, {max_n}]\n")
    total_col = 0
    total_sum_zero = 0
    total_constant = 0
    overall_dist = {}
    for n in range(12, max_n + 1):
        start = time.time()
        try:
            triangulations = list(graphs.triangulations(n, minimum_degree=5))
        except Exception as ex:
            print(f"n={n}: cannot enumerate ({ex})")
            continue
        n_col_n = 0
        n_sum_zero_n = 0
        n_constant_n = 0
        dist_n = {}
        examples_n = []
        for tri_idx, G in enumerate(triangulations):
            if time.time() - start > time_budget_per_n:
                print(f"  n={n}: timeout at tri {tri_idx}/{len(triangulations)}")
                break
            G.is_planar(set_embedding=True)
            D = dual_of(G)
            n_col_i, n_sz_i, n_cu_i, dist_i, exs = test_one(D, name=f"n{n}t{tri_idx}")
            n_col_n += n_col_i
            n_sum_zero_n += n_sz_i
            n_constant_n += n_cu_i
            for k, v in dist_i.items():
                dist_n[k] = dist_n.get(k, 0) + v
            if exs and len(examples_n) < 5:
                examples_n.extend(exs[:5 - len(examples_n)])
        elapsed = time.time() - start
        print(f"n={n}: {n_col_n} colourings  "
              f"sum=0: {n_sum_zero_n}/{n_col_n}  "
              f"constant on V(K_b)UV(K_c): {n_constant_n}/{n_col_n}  "
              f"sum-dist: {sorted(dist_n.items())}  [{elapsed:.0f}s]")
        if examples_n and n_constant_n:
            print(f"  examples of constant-union colourings: {examples_n}")
        sys.stdout.flush()
        total_col += n_col_n
        total_sum_zero += n_sum_zero_n
        total_constant += n_constant_n
        for k, v in dist_n.items():
            overall_dist[k] = overall_dist.get(k, 0) + v

    print()
    print("=" * 78)
    print(f"Grand totals (n in [12, {max_n}]):")
    print(f"  colourings tested: {total_col}")
    print(f"  sum_v h(v) = 0: {total_sum_zero} ({100*total_sum_zero/max(1,total_col):.2f}%)")
    print(f"  h constant on V(K_b)UV(K_c): {total_constant} ({100*total_constant/max(1,total_col):.2f}%)")
    print(f"  sum distribution: {sorted(overall_dist.items())}")


if __name__ == '__main__':
    main()
