"""Empirical test of the rainbow conjecture and König-lift conjecture
on cut tires from the new (pendant-redefined) definition.

For each cut tire in our Holton-McKay #0 example:
  - Build the cut tire as a graph (face boundary + labelled pendants).
  - Enumerate proper edge 3-colorings via Sage chromatic_polynomial
    on the line graph (or direct brute force for small tires).
  - Compute the projection support onto (out spoke colors, in spoke
    colors) jointly.

Tests:
  Rainbow conjecture analog:
    For tires whose face boundary is a simple cycle (spoke-only),
    by Prop 1.13 of paper.tex, # colorings = 2^n + 2(-1)^n.
    Spoke colors at each face-boundary vertex are uniquely
    determined by the cycle coloring (the "third color"). So
    projection support has specific structure.

  Stronger rainbow: face boundary closed walk that's not simple
    (= theta-graph-like). Apply Prop's from rainbow_proof.tex.

  König-lift conjecture analog:
    For pairs of adjacent cut tires (T_d, T_{d+1}) sharing in-spokes
    of T_d with face-boundary edges of T_{d+1}, do the projection
    supports overlap on the shared edges?
"""
import os
import sys
from itertools import permutations, product
from collections import defaultdict

from sage.all import Graph

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from cut_depth_label import (
    parse_planar_code, HM_FILE, find_six_edge_cut,
    apply_procedure, compute_nice_layout,
)
from cut_tire import cut_tire_at


def build_cut_tire_graph(tire, d):
    """Build a Sage Graph for the cut tire (face boundary + pendants),
    returning (G, edge_labels) where edge_labels[e] ∈
    {'face', 'out', 'in'} classifies the edge."""
    G = Graph(multiedges=False, loops=False)
    edge_labels = {}
    for (u, v) in tire['face_edges']:
        G.add_edge(u, v)
        edge_labels[(min(u, v), max(u, v))] = 'face'
    # Add pendants
    next_id = max(set(tire['face_vertices_unique']) | {0}) + 1000
    pendant_map = {}
    for (v, pid) in tire['out_spokes']:
        nv = next_id; next_id += 1
        G.add_edge(v, nv)
        edge_labels[(min(v, nv), max(v, nv))] = 'out'
        pendant_map[pid] = (v, nv)
    for (v, pid) in tire['in_spokes']:
        nv = next_id; next_id += 1
        G.add_edge(v, nv)
        edge_labels[(min(v, nv), max(v, nv))] = 'in'
        pendant_map[pid] = (v, nv)
    return G, edge_labels, pendant_map


def enumerate_proper_edge_3colorings(G):
    """Enumerate proper edge 3-colorings by computing the line graph
    and its proper 3-vertex-colorings."""
    L = G.line_graph(labels=False)
    edges = list(G.edges(labels=False))
    edges_norm = [(min(u, v), max(u, v)) for (u, v) in edges]
    n_edges = len(edges_norm)
    # Map L's vertices to edge indices
    L_verts = list(L.vertices())
    vertex_to_edge_idx = {}
    for i, e in enumerate(edges):
        e_can = frozenset(e)
        for v in L_verts:
            if frozenset(v) == e_can:
                vertex_to_edge_idx[v] = i
                break
    # Brute force over 3^n_edges, filtering for proper-coloring constraint
    colorings = []
    for assignment in product((1, 2, 3), repeat=n_edges):
        # Check: at each vertex of G, all incident edges have distinct colors
        good = True
        for v in G.vertices():
            incident = [edges_norm.index((min(v, w), max(v, w)))
                        for w in G.neighbors(v)]
            colors_at_v = [assignment[i] for i in incident]
            if len(set(colors_at_v)) != len(colors_at_v):
                good = False
                break
        if good:
            colorings.append(assignment)
    return edges_norm, colorings


def cut_tire_analysis(tire, d):
    """Analyze a single cut tire: enumerate proper 3-edge-colorings,
    compute spoke projections, check for S_3-orbit closure."""
    G, edge_labels, _ = build_cut_tire_graph(tire, d)
    n = G.size()
    if n > 18:  # too slow
        return None
    edges_norm, colorings = enumerate_proper_edge_3colorings(G)
    n_colorings = len(colorings)
    # Projection: spoke (out + in) tuple
    out_edge_indices = [i for i, e in enumerate(edges_norm)
                        if edge_labels[e] == 'out']
    in_edge_indices = [i for i, e in enumerate(edges_norm)
                       if edge_labels[e] == 'in']
    proj = set()
    for c in colorings:
        sigma = (tuple(c[i] for i in out_edge_indices),
                 tuple(c[i] for i in in_edge_indices))
        proj.add(sigma)
    # S_3-closure check
    s3_closed = True
    for sigma in proj:
        for pi in permutations((1, 2, 3)):
            mapped = (tuple(pi[c - 1] for c in sigma[0]),
                      tuple(pi[c - 1] for c in sigma[1]))
            if mapped not in proj:
                s3_closed = False
                break
        if not s3_closed:
            break
    # Count S_3-orbits
    seen = set()
    orbits = []
    for sigma in sorted(proj):
        if sigma in seen:
            continue
        orbit = set()
        for pi in permutations((1, 2, 3)):
            mapped = (tuple(pi[c - 1] for c in sigma[0]),
                      tuple(pi[c - 1] for c in sigma[1]))
            orbit.add(mapped)
        orbit_in_proj = orbit & proj
        orbits.append(sorted(orbit_in_proj))
        seen |= orbit_in_proj
    return {
        'n_edges': n,
        'n_colorings': n_colorings,
        'projection_size': len(proj),
        'S3_closed': s3_closed,
        'n_orbits': len(orbits),
        'orbit_sizes': sorted(len(o) for o in orbits),
        'projection': proj,
        'out_count': len(out_edge_indices),
        'in_count': len(in_edge_indices),
    }


def konig_lift_test(t1, t2, d1, d2, analysis1, analysis2):
    """For adjacent tires (t1 at depth d1, t2 at depth d2 = d1+1),
    test König-lift compatibility.

    Shared structure: in spokes of t1 correspond to a subset of
    face-boundary edges of t2 (those that are also incident to t1's
    face boundary).
    """
    # For now: compute the intersection of projections restricted to
    # the in-spoke side of t1.  (Full identification of shared edges
    # requires more careful bookkeeping; this is a first approximation.)
    if d2 != d1 + 1:
        return None
    in_projections_t1 = set(sigma[1] for sigma in analysis1['projection'])
    out_projections_t2 = set(sigma[0] for sigma in analysis2['projection'])
    # These don't have a natural bijection without more setup; instead
    # compare their sizes and S_3-symmetry.
    return {
        't1_in_size': len(in_projections_t1),
        't2_out_size': len(out_projections_t2),
        # The "shared boundary" between t1 and t2 isn't trivially
        # identifiable from this aggregated data; flag as needing
        # explicit shared-edge identification.
    }


def main():
    gs = parse_planar_code(HM_FILE)
    G = gs[0]
    S, cut = find_six_edge_cut(G)
    base_pos = compute_nice_layout(G)
    S1 = frozenset(G.vertices()) - frozenset(S)
    H1, pos1, ed1, _, _, _ = apply_procedure(
        G, S1, cut, base_pos, '1',
        pendant_start_id=max(G.vertices()) + 1 + 6)

    print(f"G'_1 depths range: 0 to {max(ed1.values())}\n")

    all_tires = {}  # (d, face_idx) -> tire dict + analysis
    print("=" * 80)
    print("Cut tire analysis (per tire)")
    print("=" * 80)
    for d in range(1, max(ed1.values()) + 1):
        tires = cut_tire_at(H1, ed1, d)
        for f_idx, tire in enumerate(tires):
            n_face = tire['face_length']
            n_out = len(tire['out_spokes'])
            n_in = len(tire['in_spokes'])
            n_edges = n_face + n_out + n_in
            print(f'\ndepth {d}, face {f_idx}: |face|={n_face}, '
                  f'out={n_out}, in={n_in}, total edges={n_edges}')
            if n_edges > 18:
                print(f'  ... too many edges to enumerate ({n_edges} > 18); skipping')
                continue
            analysis = cut_tire_analysis(tire, d)
            if analysis is None:
                continue
            print(f"  Proper edge 3-colorings: {analysis['n_colorings']}")
            print(f"  Projection (out, in) size: {analysis['projection_size']}")
            print(f"  S_3-closed: {analysis['S3_closed']}")
            print(f"  # S_3-orbits: {analysis['n_orbits']}, "
                  f"orbit sizes: {analysis['orbit_sizes']}")
            all_tires[(d, f_idx)] = (tire, analysis)

    # Rainbow analysis: for each tire, check if its projection is
    # related to a specific structured set (perms-per-face etc.)
    print("\n" + "=" * 80)
    print("Rainbow conjecture-style check")
    print("=" * 80)
    for (d, f_idx), (tire, analysis) in all_tires.items():
        n_out = analysis['out_count']
        n_in = analysis['in_count']
        if n_out == 0 or n_in == 0:
            print(f'\ndepth {d}, face {f_idx}: trivial '
                  f'(one side has 0 spokes), skipping rainbow check')
            continue
        # Expected: if both sides have 3 spokes each, check perms-per-face
        if n_out == 3 and n_in == 3:
            expected_size = 36  # perms × perms
            actual = analysis['projection_size']
            print(f'\ndepth {d}, face {f_idx}: out={n_out}, in={n_in}; '
                  f'predicted perms-per-face size = {expected_size}, '
                  f'actual = {actual}')


if __name__ == '__main__':
    main()
