"""Chain DP with JOINT projection tracking + rigorous edge-based
parent verification + ground-truth comparison.

Key fixes vs. chain_dp_general.py:

(1) DP state per tire = SET of full proper edge 3-colorings of the
    tire (cycle + pendants), indexed by edge id within the tire.
    No projection to spokes; the full coloring is the state.

(2) Composition between parent T_p and child T_c is via SHARED
    EDGES (= full G'_i edge tuples appearing in both tires).
    Parent coloring c_p is compatible with child coloring c_c iff
    they agree on every shared edge.  No need to figure out which
    in-spoke maps to which cycle edge — equality of full edge
    tuples does the work.

(3) Parent-finding sanity: after build_tree's heuristic, verify
    each (child, parent) pair shares at least one edge.  If not,
    flag as bug.

(4) Ground truth: compute ALL proper 3-edge-colorings of G'_i via
    backtracking with constraint propagation.  Project to cut
    edges.  Compare R_i^DP with R_i^direct.  They should match.

The empirical test: for 3-edge-colorable G (dodecahedron, HM #0
are all 3-edge-colorable by Vizing/Tait), R_0 ∩ R_1 should be
non-empty (= projection of a full G coloring to cut).
"""
import os
import sys
import itertools
from collections import defaultdict

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 tire_edges(H, edge_depth, d, face):
    """Return ordered list of edges in cut tire T_d^{(face)}.
    Cycle edges (depth d) + pendants (depth d-1 OUT and d+1 IN).
    Each edge is a sorted tuple of vertices.
    Returns (edges_in_order, vertex_to_edge_ids, in_pendant_ids,
    out_pendant_ids, cycle_edge_ids)."""
    # Cycle edges (dedup)
    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 cycle
    verts = set()
    for u, v in cycle_edges:
        verts.add(u); verts.add(v)
    # Pendants
    pendants_in, pendants_out = [], []
    for v in verts:
        for nb in H.neighbors(v):
            e = tuple(sorted((v, nb)))
            if e in cycle_e_set:
                continue
            de = edge_depth.get(e)
            if de == d + 1:
                pendants_in.append((v, e))
            elif de == d - 1:
                pendants_out.append((v, e))
    # Edge list: cycle then in pendants then out pendants
    all_edges = cycle_edges + [e for (_, e) in pendants_in] + \
                [e for (_, e) in pendants_out]
    edge_id = {e: i for i, e in enumerate(all_edges)}
    # Per vertex, edges at v
    edges_at_v = defaultdict(list)
    for (a, b) in cycle_edges:
        edges_at_v[a].append(edge_id[(a, b)])
        edges_at_v[b].append(edge_id[(a, b)])
    for (v, e) in pendants_in:
        edges_at_v[v].append(edge_id[e])
    for (v, e) in pendants_out:
        edges_at_v[v].append(edge_id[e])
    cycle_eids = [edge_id[e] for e in cycle_edges]
    in_eids = [edge_id[e] for (_, e) in pendants_in]
    out_eids = [edge_id[e] for (_, e) in pendants_out]
    return {
        'all_edges': all_edges,
        'edge_id': edge_id,
        'edges_at_v': dict(edges_at_v),
        'cycle_eids': cycle_eids,
        'in_eids': in_eids,
        'out_eids': out_eids,
    }


def enumerate_proper(struct):
    """Enumerate proper 3-edge-colorings of the tire as tuples of
    colors (length = n_edges). Return None if too big."""
    n_e = len(struct['all_edges'])
    if n_e > 14:
        return None
    out = []
    for assign in itertools.product([0, 1, 2], repeat=n_e):
        ok = True
        for v, eids in struct['edges_at_v'].items():
            cs = [assign[i] for i in eids]
            if len(set(cs)) != len(cs):
                ok = False; break
        if ok:
            out.append(assign)
    return out


def chain_dp_joint(faces_by_depth, parent_of, H, edge_depth):
    """Chain DP with full per-tire colorings tracked.
    Returns:
      A: node -> list of full colorings (tuples)
      tire_struct: node -> structure dict
      issues: list of warnings
    """
    tire_struct = {}
    for d, faces in faces_by_depth.items():
        for fi, face in enumerate(faces):
            tire_struct[(d, fi)] = tire_edges(H, edge_depth, d, face)

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

    issues = []
    # Verify parent-child have shared edges
    for ch_node, p_node in parent_of.items():
        if p_node is None or p_node == ('cut', None):
            continue
        ch_edges = set(tire_struct[ch_node]['all_edges'])
        p_edges = set(tire_struct[p_node]['all_edges'])
        shared = ch_edges & p_edges
        if not shared:
            issues.append({
                'type': 'no_shared',
                'child': ch_node, 'parent': p_node,
                'n_ch_edges': len(ch_edges),
                'n_p_edges': len(p_edges),
            })

    max_d = max(d for (d, _) in tire_struct)
    A = {}
    for d in range(max_d, 0, -1):
        for node in [n for n in tire_struct if n[0] == d]:
            cs = enumerate_proper(tire_struct[node])
            if cs is None:
                A[node] = None
                continue
            kids = children.get(node, [])
            if not kids:
                A[node] = cs
                continue
            # For each kid, build {shared_color_tuple: count} so
            # parent can check existence quickly
            kid_shared = []
            for ch in kids:
                if A.get(ch) is None:
                    kid_shared.append(None)
                    continue
                ch_struct = tire_struct[ch]
                # Shared edges between parent and child
                p_edges = tire_struct[node]['all_edges']
                p_eid = tire_struct[node]['edge_id']
                ch_eid = ch_struct['edge_id']
                shared_pairs = []  # (p_eid, ch_eid)
                for e in p_edges:
                    if e in ch_eid:
                        shared_pairs.append((p_eid[e], ch_eid[e]))
                if not shared_pairs:
                    kid_shared.append(None)
                    continue
                # Build set of child shared-tuples (just keys we care)
                ch_shared_set = set()
                for ch_assign in A[ch]:
                    key = tuple(ch_assign[c_eid] for (_, c_eid)
                                in shared_pairs)
                    ch_shared_set.add(key)
                kid_shared.append((shared_pairs, ch_shared_set))
            # Filter parent colorings by all kid constraints
            achievable = []
            for assign in cs:
                ok = True
                for ks in kid_shared:
                    if ks is None:
                        continue
                    shared_pairs, ch_shared_set = ks
                    key = tuple(assign[p_eid] for (p_eid, _) in shared_pairs)
                    if key not in ch_shared_set:
                        ok = False; break
                if ok:
                    achievable.append(assign)
            A[node] = achievable
    return A, tire_struct, issues


def proper_3_edge_colorings_full(G, max_n=100):
    """Brute-force enumerate all proper 3-edge-colorings of G via
    backtracking + constraint propagation.  Returns list of dicts
    {edge_tuple: color}.  Only feasible for small G."""
    edges = [tuple(sorted(e[:2])) for e in G.edges()]
    n = len(edges)
    if n > max_n:
        return None
    edge_idx = {e: i for i, e in enumerate(edges)}
    # Edges incident to each vertex
    eav = defaultdict(list)
    for i, (u, v) in enumerate(edges):
        eav[u].append(i)
        eav[v].append(i)
    # Backtracking
    colors = [-1] * n
    results = []

    def backtrack(i):
        if i == n:
            results.append(tuple(colors))
            return
        u, v = edges[i]
        used = set()
        for j in eav[u]:
            if colors[j] != -1:
                used.add(colors[j])
        for j in eav[v]:
            if colors[j] != -1:
                used.add(colors[j])
        for c in range(3):
            if c in used:
                continue
            colors[i] = c
            backtrack(i + 1)
            colors[i] = -1

    backtrack(0)
    return [{edges[i]: c for i, c in enumerate(r)} for r in results]


def project_to_cut(colorings, cut_edges):
    """Project a list of full colorings to the given cut edges.
    cut_edges is a list of original-graph edges (tuples).
    Returns set of color tuples."""
    cut_norm = [tuple(sorted(e[:2])) for e in cut_edges]
    out = set()
    for col in colorings:
        try:
            out.add(tuple(col[ce] for ce in cut_norm))
        except KeyError:
            continue
    return out


def s3_orbit(t):
    """Return S_3 orbit of tuple under color permutation."""
    return {tuple(p[c] for c in t) for p in itertools.permutations([0, 1, 2])}


def run_one_cut(G, cut_idx, cuts, base_pos, verbose=False):
    S, cut_edges = cuts[cut_idx]
    S0, S1 = S, frozenset(G.vertices()) - S
    if len(S0) < 4 or len(S1) < 4:
        return None
    if verbose:
        print(f'  cut #{cut_idx}: |S0|={len(S0)}, |S1|={len(S1)}',
              flush=True)
    result = {}
    # Ground truth: enumerate G's 3-edge-colorings, project to cut
    g_cols = proper_3_edge_colorings_full(G)
    if g_cols is None:
        result['ground_truth'] = None
    else:
        R_ground = project_to_cut(g_cols, cut_edges)
        result['R_ground'] = R_ground
        if verbose:
            print(f'    G has {len(g_cols)} proper 3-edge-colorings, '
                  f'|R_ground|={len(R_ground)}', flush=True)
    # DP per side
    for side, S_side in [('0', S0), ('1', S1)]:
        pendant_id = max(G.vertices()) + 1 + (0 if side == '0' else 500)
        try:
            H, pos, ed, _, _, _ = apply_procedure(
                G, S_side, cut_edges, base_pos, side,
                pendant_start_id=pendant_id)
        except Exception as e:
            result[f'side_{side}_error'] = str(e)
            continue
        if not ed:
            continue
        faces_by_depth, parent_of, _ = build_tree(H, ed)
        try:
            A, struct, issues = chain_dp_joint(
                faces_by_depth, parent_of, H, ed)
        except Exception as e:
            result[f'side_{side}_dp_error'] = str(e)
            continue
        if issues:
            result[f'side_{side}_issues'] = issues
        # Roots = depth-1 tires; project to cut (= depth-0 pendant edges
        # in the H subgraph).  Cut pendant edges are stored as pseudo
        # edges in H.
        # Cut edges in H: pendant edges at the original cut endpoints
        cut_eids_in_H = []
        for u, v in cut_edges:
            # In side-i view, only the side-i endpoint exists; the cut
            # edge in H is between the original endpoint and a new
            # pendant vertex.  We need to identify these pendants.
            pass
        # Simpler: for each root tire T_1^{(f)} on this side, its out
        # spokes (depth 0) ARE the cut edges (or some of them).
        R_dp = set()
        for node, a in A.items():
            if node[0] != 1 or a is None:
                continue
            # We want to identify cut edges in this tire's edges.
            # Cut edges are depth-0 pendants in H.  Filter struct's
            # all_edges to depth-0 edges.
            depth_0_eids = [i for i, e in enumerate(struct[node]['all_edges'])
                            if ed.get(e) == 0]
            if not depth_0_eids:
                continue
            for assign in a:
                tup = tuple(assign[i] for i in depth_0_eids)
                # Also store WHICH cut edges (by tuple)
                edges_used = [struct[node]['all_edges'][i]
                              for i in depth_0_eids]
                # Project: key by edge tuple, value by color
                R_dp.add(frozenset(zip(edges_used, tup)))
        result[f'R_dp_{side}_raw'] = R_dp
        # Convert frozenset-of-(edge,color) to a per-cut-edge color
        # tuple (in fixed order of cut_edges)
        cut_norm = [tuple(sorted(e[:2])) for e in cut_edges]
        # The cut edges in H may have different endpoints (pendant
        # vertices added by apply_procedure), so we need a mapping.
        # For now, count cardinality and S3 stats.
        result[f'R_dp_{side}_size'] = len(R_dp)
    return result


def main():
    print('=== Dodecahedron ===', flush=True)
    G = graphs.DodecahedralGraph()
    G.is_planar(set_embedding=True)
    base_pos = compute_nice_layout(G)
    cuts = all_six_edge_cuts(G, max_cuts=3)
    for cut_idx in range(min(3, len(cuts))):
        res = run_one_cut(G, cut_idx, cuts, base_pos, verbose=True)
        if res is None:
            continue
        if 'R_ground' in res:
            R_ground = res['R_ground']
            print(f'    ground truth |R|={len(R_ground)}, S_3 orbits='
                  f'{len({frozenset(s3_orbit(t)) for t in R_ground})}',
                  flush=True)
        for side in ['0', '1']:
            sz = res.get(f'R_dp_{side}_size', 'N/A')
            issues = res.get(f'side_{side}_issues')
            print(f'    side {side}: |R_dp|={sz}, '
                  f'issues={len(issues) if issues else 0}', flush=True)
            if issues:
                for iss in issues[:2]:
                    print(f'      {iss}', flush=True)


if __name__ == '__main__':
    main()
