"""Generalized chain DP for cut tires (handles branched H_d faces).

A cut tire T_d^{(f)} is treated as:
  - Edge set E_f = edges in the boundary walk of face f of H_d.
  - Vertex set V_f = vertices incident to E_f.
  - For each v in V_f with deg_{H_d}(v) = 2: one labeled pendant
    (in or out depending on the depth of the non-H_d edge at v).
  - For each v in V_f with deg_{H_d}(v) = 3: no pendant.

Proper 3-edge-coloring of T = (E_f ∪ pendants) such that all edges
incident at every v in V_f have distinct colors. (The "underlying"
graph treats E_f and pendants as edges in a planar multigraph; a
branch vertex has 3 incident E_f edges all distinct.)

Out-spokes projection: tuple of colors on labeled-OUT pendants
(depth d-1 edges).

Chain DP bottom-up:
  - At each tire, enumerate proper 3-edge-colorings.
  - Restrict to those compatible with at least one combination
    of out-spoke colorings from each child.
  - Project onto own out spokes.

Verify:
  (a) At each tire: |A(T)| > 0, contains full S_3 orbit.
  (b) R_0 ∩ R_1 non-empty at the cut.
"""
import os
import sys
import itertools
from collections import defaultdict, Counter

from sage.all import Graph, graphs

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from cut_depth_label import (
    parse_planar_code, HM_FILE,
    apply_procedure, compute_nice_layout,
)
from cut_tire_tree import build_tree
from tree_structure_sweep import all_six_edge_cuts


def cut_tire_structure(H, edge_depth, d, face):
    """Extract structure of cut tire T_d^{(face)}.

    Returns dict:
      'edges_indexed': list of (e_id, edge_tuple) for cycle edges.
                       Order: by appearance in face walk (so duplicates
                       can happen for branch points, but we use each
                       edge once).
      'pendants_in':   dict v -> (pendant_id, depth=d+1 edge)
      'pendants_out':  dict v -> (pendant_id, depth=d-1 edge)
      'edges_at_v':    dict v -> list of edge_ids incident to v
                       (each appearance of v in walk = 1 cycle edge)
      'cycle_edges':   list of (cycle_edge_id, edge_tuple)
    """
    # First normalize: count each unique cycle edge once
    cycle_e_set = set()
    cycle_edges = []
    for u, v in face:
        e = tuple(sorted((u, v)))
        if e not in cycle_e_set:
            cycle_e_set.add(e)
            cycle_edges.append(e)
    # Vertices on this face boundary
    verts_set = set()
    for u, v in cycle_edges:
        verts_set.add(u)
        verts_set.add(v)
    # Edges incident to each vertex (cycle edges only)
    edges_at_v = defaultdict(list)
    for ce_id, (a, b) in enumerate(cycle_edges):
        edges_at_v[a].append(ce_id)
        edges_at_v[b].append(ce_id)
    # Pendants: at each degree-2 vertex of H, identify the third edge
    pendants_in = {}
    pendants_out = {}
    pendant_counter = [len(cycle_edges)]  # next available edge id
    for v in verts_set:
        H_deg = sum(1 for nb in H.neighbors(v)
                    if tuple(sorted((v, nb))) in cycle_e_set)
        # Wait — that's just degree restricted to face cycle edges.
        # But pendant attaches to the H-edge of different depth.
        # Look at ALL edges of H incident to v:
        for nb in H.neighbors(v):
            e_full = tuple(sorted((v, nb)))
            d_e = edge_depth.get(e_full)
            if d_e is None:
                continue
            if d_e == d:
                continue  # skip cycle edges
            elif d_e == d + 1:
                pid = pendant_counter[0]
                pendant_counter[0] += 1
                pendants_in[v] = (pid, e_full)
                edges_at_v[v].append(pid)
            elif d_e == d - 1:
                pid = pendant_counter[0]
                pendant_counter[0] += 1
                pendants_out[v] = (pid, e_full)
                edges_at_v[v].append(pid)
    return {
        'cycle_edges': cycle_edges,
        'pendants_in': pendants_in,
        'pendants_out': pendants_out,
        'edges_at_v': dict(edges_at_v),
        'n_edges': pendant_counter[0],
    }


def proper_3_colorings_general(struct):
    """Enumerate all proper 3-edge-colorings of cut tire.
    Edge ids 0..n_cycle_edges-1 are cycle edges; n_cycle...n_edges-1 are
    pendants. Constraint: at each vertex v, the colors of edges in
    edges_at_v[v] are all distinct."""
    n_e = struct['n_edges']
    edges_at_v = struct['edges_at_v']
    # Edge-coloring constraint reduces to: per vertex, colors distinct.
    # Since vertex can have ≤ 3 incident edges in cubic (since each
    # vertex in original G' is cubic), and our cut tire is subgraph,
    # each v has at most 3 incident edges in the cut tire.
    # We need all colors at each v distinct.
    out = []
    # Naive enumeration: 3^n_e then filter.  n_e is small (~10-50)
    # for our test cases, so doable.
    if n_e > 12:
        return None  # too big
    for assign in itertools.product([0, 1, 2], repeat=n_e):
        ok = True
        for v, eids in edges_at_v.items():
            colors = [assign[eid] for eid in eids]
            if len(set(colors)) != len(colors):
                ok = False
                break
        if ok:
            out.append(assign)
    return out


def s3_orbit(tup):
    """S_3 orbit of color tuple under permutations of {0,1,2}."""
    orbit = set()
    for perm in itertools.permutations([0, 1, 2]):
        orbit.add(tuple(perm[c] for c in tup))
    return orbit


def has_full_s3_orbit(s):
    """Does set s of tuples contain a 6-element S_3 orbit?"""
    if not s:
        return False
    for tup in s:
        orb = s3_orbit(tup)
        if len(orb) == 6 and orb <= s:
            return True
    return False


def chain_dp_general(faces_by_depth, parent_of, H, edge_depth):
    """Run chain DP, returns dict node -> A(T) (set of out-spoke
    color tuples)."""
    # Build structures for each tire
    structs = {}
    for d, faces in faces_by_depth.items():
        for f_idx, face in enumerate(faces):
            structs[(d, f_idx)] = cut_tire_structure(H, edge_depth, d, face)

    # Build children list
    children = defaultdict(list)
    for node, p in parent_of.items():
        if p is not None and p != ('cut', None):
            children[p].append(node)

    # Order by depth (process leaves first)
    max_d = max(d for (d, _) in structs)
    A = {}
    for d in range(max_d, 0, -1):
        for (dd, f_idx), struct in list(structs.items()):
            if dd != d:
                continue
            node = (d, f_idx)
            cs = proper_3_colorings_general(struct)
            if cs is None:
                A[node] = None  # too big
                continue
            # Pendant-IN edge ids: those whose pendant_id is in pendants_in
            in_pids = {v: pid for v, (pid, _) in struct['pendants_in'].items()}
            out_pids = {v: pid for v, (pid, _) in struct['pendants_out'].items()}
            in_v_sorted = sorted(in_pids.keys())
            out_v_sorted = sorted(out_pids.keys())

            # Find correspondence: each pendant-IN edge of this tire
            # = full edge in G'_i. We need to find which child cycle
            # edge it corresponds to.
            # The pendant-IN at v has full edge e (in struct).
            # Among children, find child c whose cycle_edges contains e.
            child_links = []  # (child_node, parent_v -> child_edge_id)
            for ch in children.get(node, []):
                if ch not in A:
                    continue
                if A[ch] is None:
                    child_links.append((ch, None))  # too big, skip
                    continue
                ch_struct = structs[ch]
                ch_out_v = sorted(ch_struct['pendants_out'].keys())
                ch_out_full = {v: ch_struct['pendants_out'][v][1]
                               for v in ch_out_v}
                # Match: parent's pendant-IN edge = child's pendant-OUT edge?
                # No wait — parent's pendant-IN is a depth-(d+1) edge.
                # Child T_{d+1}^{(f')} has cycle edges of depth d+1.
                # So parent's pendant-IN edge = a cycle edge of child.
                p_v_to_child_cycle_eid = {}
                for v_p, (pid_p, e_full_p) in struct['pendants_in'].items():
                    # Find edge in child's cycle_edges with same e_full
                    for ce_id, e_c in enumerate(ch_struct['cycle_edges']):
                        if e_c == e_full_p:
                            p_v_to_child_cycle_eid[v_p] = ce_id
                            break
                child_links.append((ch, p_v_to_child_cycle_eid))

            # For each coloring, check compatibility with each child
            achievable = set()
            for assign in cs:
                # Determine parent's pendant-IN colors
                p_in_colors = {v: assign[in_pids[v]] for v in in_pids}
                # For each child, check if there's a child-coloring with
                # the same colors at the corresponding cycle edges.
                # We use child's A in terms of OUT-SPOKE colors, but we
                # need to constrain CYCLE colors. So we need to recompute
                # per-child compatibility from child's allowed colorings.
                # ALTERNATE: instead of using A(ch) (out-spoke restriction),
                # track full child colorings. This blows up but is correct.
                # Simpler: track child's set of valid (cycle, out) colorings.
                # For now we use the WEAKER check: out-spoke compatibility
                # via "any child coloring at all", since A(ch) doesn't
                # carry cycle-edge info.
                # Since the chain compatibility is on PARENT'S IN spokes =
                # CHILD'S CYCLE EDGES (not out spokes), we actually need
                # the child's projection on cycle edges, not out spokes.
                # For this test, let's use child's full coloring set.
                ok = True
                for ch, link in child_links:
                    if link is None:
                        # child too big — skip compatibility check
                        continue
                    ch_full_colorings = proper_3_colorings_general(structs[ch])
                    if ch_full_colorings is None:
                        continue
                    # Find a ch_coloring with matching cycle-edge colors
                    found = False
                    for ch_assign in ch_full_colorings:
                        match = True
                        for v_p, ce_id_c in link.items():
                            if ch_assign[ce_id_c] != p_in_colors[v_p]:
                                match = False
                                break
                        if match:
                            found = True
                            break
                    if not found:
                        ok = False
                        break
                if ok:
                    # Project to out spokes
                    if out_v_sorted:
                        tup = tuple(assign[out_pids[v]] for v in out_v_sorted)
                    else:
                        tup = ()
                    achievable.add(tup)
            A[node] = achievable
    return A, structs


def test_on_graph(G, name, max_cuts=3):
    cuts = all_six_edge_cuts(G, max_cuts=max_cuts)
    print(f'\n=== {name}: V={G.order()}, {len(cuts)} cuts tested ===',
          flush=True)
    base_pos = compute_nice_layout(G)
    n_tests, n_pertire_orbit_ok, n_root_ok, n_cross_ok = 0, 0, 0, 0
    per_tire_failures = []
    cross_failures = []
    for cut_idx, (S, cut_edges) in enumerate(cuts):
        S0, S1 = S, frozenset(G.vertices()) - S
        if len(S0) < 4 or len(S1) < 4:
            continue
        Rs = {}
        for side, S_side in [('0', S0), ('1', S1)]:
            try:
                H, pos, ed, _, _, _ = apply_procedure(
                    G, S_side, cut_edges, base_pos, side,
                    pendant_start_id=max(G.vertices()) + 1 +
                    (0 if side == '0' else 200))
            except Exception:
                continue
            if not ed:
                continue
            faces_by_depth, parent_of, _ = build_tree(H, ed)
            try:
                A, structs = chain_dp_general(
                    faces_by_depth, parent_of, H, ed)
            except Exception as e:
                print(f'  cut #{cut_idx}/{side}: DP failed: {e}', flush=True)
                continue
            # Per-tire check
            for node, a in A.items():
                if a is None:
                    continue
                n_tests += 1
                if a and has_full_s3_orbit(a):
                    n_pertire_orbit_ok += 1
                else:
                    per_tire_failures.append({
                        'cut': cut_idx, 'side': side, 'node': node,
                        'n_edges': structs[node]['n_edges'],
                        '|A|': len(a),
                        'has_orbit': has_full_s3_orbit(a) if a else False,
                    })
            # Roots
            R_i = set()
            for node in A:
                if node[0] == 1 and A[node]:
                    R_i |= A[node]
            if R_i and has_full_s3_orbit(R_i):
                n_root_ok += 1
            Rs[side] = R_i
        if '0' in Rs and '1' in Rs:
            if Rs['0'] & Rs['1']:
                n_cross_ok += 1
            else:
                cross_failures.append({
                    'cut': cut_idx,
                    '|R0|': len(Rs['0']),
                    '|R1|': len(Rs['1']),
                })
    print(f'  Per-tire A(T) full-orbit count: {n_pertire_orbit_ok}/{n_tests}',
          flush=True)
    if per_tire_failures:
        print(f'  First per-tire failures:')
        for f in per_tire_failures[:5]:
            print(f'    {f}')
    print(f'  R_i has full orbit (root sides): {n_root_ok}', flush=True)
    print(f'  Cross-cut R_0 ∩ R_1 != empty: {n_cross_ok}', flush=True)
    if cross_failures:
        print(f'  Cross failures: {cross_failures[:3]}')


def main():
    G = graphs.DodecahedralGraph()
    G.is_planar(set_embedding=True)
    test_on_graph(G, 'Dodecahedron', max_cuts=3)
    gs = parse_planar_code(HM_FILE)
    test_on_graph(gs[0], 'HM_0', max_cuts=3)


if __name__ == '__main__':
    main()
