"""For every chord-apex+Kempe colouring, walk the Kempe cycle K_b
through the merged edge and at each consecutive pair (v_k, v_{k+1})
record:
  - Is h_phi(v_k) == h_phi(v_{k+1}) (same Heawood)?
  - Is the c-edge at v_k on the same side of K_b as the c-edge at v_{k+1}?

The "side" of a c-edge at v_k is identified by removing the K_b edges
from H and asking which connected component of the result contains the
non-v_k endpoint of v_k's c-edge.

Lemma A (proposed) says: same Heawood at consecutive K-vertices ==>
their c-edges are on the same side of K_b. We test the empirical
correlation across all chord-apex+Kempe colourings.

Run with: sage experiments/check_heawood_side_correlation.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,
    trace_kempe_cycle,
    edge_idx,
)
from check_heawood_on_kempe import dual_of, heawood_numbers


def kempe_sides(H, kempe_edge_frozensets):
    """Remove the kempe edges from H. Return v -> component index (in
    H \\ K)."""
    H2 = H.copy()
    for fs in kempe_edge_frozensets:
        u, v = tuple(fs)
        if H2.has_edge(u, v):
            H2.delete_edge(u, v)
    comps = H2.connected_components()
    v_to_comp = {}
    for i, comp in enumerate(comps):
        for v in comp:
            v_to_comp[v] = i
    return v_to_comp, len(comps)


def test_one(D):
    """Tally (same-h, same-side) correlation over consecutive K_b pairs."""
    D.is_planar(set_embedding=True)
    n_col = 0
    # 4-way tally of (same_h, same_side):
    counts = {(True, True): 0, (True, False): 0,
              (False, True): 0, (False, False): 0}
    n_2components = 0  # how often does H \ K_b have exactly 2 comps?
    skipped = 0
    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:
                    skipped += 1
                    continue
                merged_idx = edge_idx(edges, named['merged'])
                a = col[merged_idx]
                for b in range(3):
                    if b == a: continue
                    c_color = 3 - a - b
                    walk = trace_kempe_cycle(edges, col, merged_idx, (a, b))
                    walk_edge_indices = [w[0] for w in walk]
                    walk_edge_fs = [frozenset(edges[i])
                                    for i in walk_edge_indices]
                    # Vertex sequence along the cycle
                    walk_vertices = []
                    cur_v = list(edges[walk_edge_indices[0]])[0]
                    for k, (ei, leave) in enumerate(walk):
                        walk_vertices.append(leave)
                    # walk_vertices[k] is the vertex we end at after edge k
                    # equivalently, walk_vertices[k] = first vertex of edge
                    # k+1.  The starting vertex of edge 0 is the cycle's
                    # closing point.
                    L = len(walk_vertices)
                    v_to_comp, n_comp = kempe_sides(H, walk_edge_fs)
                    if n_comp == 2:
                        n_2components += 1
                    # For each vertex on K_b, find its c-edge endpoint
                    v_side = {}    # vertex -> component index of c-neighbor
                    for v in walk_vertices:
                        for u in H.neighbor_iterator(v):
                            ei = edge_idx(edges, frozenset((u, v)))
                            if col[ei] == c_color:
                                v_side[v] = v_to_comp.get(u, -1)
                                break
                    # Tally consecutive pairs
                    for k in range(L):
                        v0 = walk_vertices[k]
                        v1 = walk_vertices[(k + 1) % L]
                        if v0 not in v_side or v1 not in v_side:
                            continue
                        same_h = (h[v0] == h[v1])
                        same_side = (v_side[v0] == v_side[v1])
                        counts[(same_h, same_side)] += 1
    return n_col, counts, n_2components, skipped


def main(max_n=18, time_budget_per_n=1800):
    print(f"(same-h, same-side) correlation on consecutive K_b pairs, "
          f"n in [12, {max_n}]\n")
    grand_col = 0
    grand_counts = {(True, True): 0, (True, False): 0,
                    (False, True): 0, (False, False): 0}
    grand_2comp = 0
    grand_skipped = 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
        counts_n = {(True, True): 0, (True, False): 0,
                    (False, True): 0, (False, False): 0}
        n_2comp_n = 0; n_skipped_n = 0
        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, c_i, n_2c_i, sk_i = test_one(D)
            n_col_n += n_col_i
            for k, v in c_i.items(): counts_n[k] = counts_n.get(k, 0) + v
            n_2comp_n += n_2c_i; n_skipped_n += sk_i
        elapsed = time.time() - start
        total_pairs = sum(counts_n.values())
        print(f"n={n}: {n_col_n} col., {total_pairs} pairs  [{elapsed:.0f}s]")
        print(f"  (same_h, same_side): {counts_n}")
        if (counts_n[(True, False)] == 0 and counts_n[(False, True)] == 0
                and total_pairs > 0):
            print(f"  *** perfect biconditional: same_h <==> same_side ***")
        sys.stdout.flush()
        grand_col += n_col_n
        for k, v in counts_n.items(): grand_counts[k] = grand_counts.get(k, 0) + v
        grand_2comp += n_2comp_n; grand_skipped += n_skipped_n

    print()
    print("=" * 78)
    print(f"Grand totals (n in [12, {max_n}], {grand_col} col., "
          f"{sum(grand_counts.values())} consecutive K_b pairs):")
    print(f"  K_b separates H into 2 components in "
          f"{grand_2comp}/{2*grand_col} K_b's "
          f"({100*grand_2comp/max(1,2*grand_col):.1f}%)")
    print(f"  (same_h=Y, same_side=Y): {grand_counts[(True, True)]}")
    print(f"  (same_h=Y, same_side=N): {grand_counts[(True, False)]}")
    print(f"  (same_h=N, same_side=Y): {grand_counts[(False, True)]}")
    print(f"  (same_h=N, same_side=N): {grand_counts[(False, False)]}")
    total = sum(grand_counts.values())
    if total > 0:
        for k, v in grand_counts.items():
            print(f"    {k}: {100*v/total:.2f}%")
    # Test the biconditional empirically.
    if grand_counts[(True, False)] == 0 and grand_counts[(False, True)] == 0:
        print(f"  *** EMPIRICAL BICONDITIONAL HOLDS: same h <==> same side ***")
    elif grand_counts[(True, False)] == 0:
        print(f"  empirical IMPLICATION holds: same h ==> same side")
    elif grand_counts[(False, True)] == 0:
        print(f"  empirical IMPLICATION holds: same side ==> same h")
    else:
        print(f"  no biconditional or one-way implication")


if __name__ == '__main__':
    main()
