"""Test whether the Heawood sum on each Kempe cycle (K_b, K_c through
merged) is always 0, on chord-apex+Kempe colourings of reduced duals.

If sum_{v in V(K)} h_phi(v) = 0 for every Kempe cycle K, then h_phi cannot
be constant on V(K) (since |V(K)| > 0 and h takes values +-1), and hence
not constant on V(K_b) U V(K_c) either. By Lemma 5.3's contrapositive,
this proves Conjecture 5.1.

Run with: sage experiments/check_heawood_per_kempe_cycle.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,
)
from check_heawood_on_kempe import (
    dual_of,
    heawood_numbers,
    vertices_of_kempe,
)


def test_one(D):
    """Check Heawood sums on K_b, K_c for every chord-apex+Kempe colouring."""
    D.is_planar(set_embedding=True)
    n_col = 0
    sums_Kb = {}
    sums_Kc = {}
    nonzero_examples = []

    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
                try:
                    h = heawood_numbers(H, edges, col)
                except RuntimeError:
                    continue
                merged_idx = edge_idx(edges, named['merged'])
                a = col[merged_idx]
                bs = [c for c in range(3) if c != a]
                # K_b is the {a, b[0]}-Kempe cycle; K_c is the {a, b[1]}-cycle.
                kc_b = kempe_cycle_set(edges, col, merged_idx, (a, bs[0]))
                kc_c = kempe_cycle_set(edges, col, merged_idx, (a, bs[1]))
                Vb = vertices_of_kempe(edges, kc_b)
                Vc = vertices_of_kempe(edges, kc_c)
                sb = sum(h[v] for v in Vb)
                sc = sum(h[v] for v in Vc)
                sums_Kb[sb] = sums_Kb.get(sb, 0) + 1
                sums_Kc[sc] = sums_Kc.get(sc, 0) + 1
                if (sb != 0 or sc != 0) and len(nonzero_examples) < 5:
                    nonzero_examples.append({
                        'sum_Kb': sb, 'sum_Kc': sc,
                        '|V(K_b)|': len(Vb), '|V(K_c)|': len(Vc),
                    })
    return n_col, sums_Kb, sums_Kc, nonzero_examples


def main(max_n=18, time_budget_per_n=1800):
    print(f"Heawood sum per Kempe cycle (through merged), "
          f"n in [12, {max_n}]\n")
    grand_col = 0
    grand_Kb = {}
    grand_Kc = {}
    nonzero_total = 0
    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
        Kb_n = {}; Kc_n = {}
        exs = []
        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, Kb_i, Kc_i, nz_i = test_one(D)
            n_col_n += n_col_i
            for k, v in Kb_i.items(): Kb_n[k] = Kb_n.get(k, 0) + v
            for k, v in Kc_i.items(): Kc_n[k] = Kc_n.get(k, 0) + v
            if nz_i and len(exs) < 3:
                exs.extend(nz_i[:3 - len(exs)])
        elapsed = time.time() - start
        nz_Kb_n = sum(v for k, v in Kb_n.items() if k != 0)
        nz_Kc_n = sum(v for k, v in Kc_n.items() if k != 0)
        print(f"n={n}: {n_col_n} col.  "
              f"sum_Kb != 0: {nz_Kb_n}, dist={sorted(Kb_n.items())}  "
              f"sum_Kc != 0: {nz_Kc_n}, dist={sorted(Kc_n.items())}  "
              f"[{elapsed:.0f}s]")
        if exs:
            print(f"  examples of nonzero sums: {exs}")
            nonzero_total += len(exs)
        sys.stdout.flush()
        grand_col += n_col_n
        for k, v in Kb_n.items(): grand_Kb[k] = grand_Kb.get(k, 0) + v
        for k, v in Kc_n.items(): grand_Kc[k] = grand_Kc.get(k, 0) + v

    print()
    print("=" * 78)
    print(f"Grand totals (n in [12, {max_n}]):")
    print(f"  colourings tested: {grand_col}")
    print(f"  sum on K_b distribution: {sorted(grand_Kb.items())}")
    print(f"  sum on K_c distribution: {sorted(grand_Kc.items())}")
    nz_Kb = sum(v for k, v in grand_Kb.items() if k != 0)
    nz_Kc = sum(v for k, v in grand_Kc.items() if k != 0)
    print(f"  sum_Kb != 0: {nz_Kb}/{grand_col} ({100*nz_Kb/max(1,grand_col):.2f}%)")
    print(f"  sum_Kc != 0: {nz_Kc}/{grand_col} ({100*nz_Kc/max(1,grand_col):.2f}%)")


if __name__ == '__main__':
    main()
