"""Lift the recoloured + bridged H_1 back to G' (dual of the n=14
triangulation), producing a proper 3-edge-colouring of the modified G'.

The modified G' = G' with: subdivisions Y_1, Y_2 of the two green witness
edges, plus a new red edge Y_1-Y_2.

Run with: sage experiments/draw_lift_to_Gprime.py
"""
from sage.all import Graph
from sage.graphs.graph_generators import graphs
import matplotlib.pyplot as plt
import math
import os


def tutte_layout(G_sage, avoid_verts=None, iterations=300):
    """Same Tutte barycentric layout as in draw_iterated_reduction_n14.py."""
    avoid = set(avoid_verts or ())
    candidates = []
    for face in G_sage.faces():
        verts = [u for (u, v) in face]
        if not (set(verts) & avoid):
            candidates.append(verts)
    if not candidates:
        outer = [u for (u, v) in max(G_sage.faces(), key=len)]
    else:
        outer = max(candidates, key=len)
    n_outer = len(outer)
    pos = {}
    for k, v in enumerate(outer):
        ang = 2 * math.pi * k / n_outer + math.pi / 2
        pos[v] = (math.cos(ang), math.sin(ang))
    interior = [v for v in G_sage.vertex_iterator() if v not in pos]
    for v in interior: pos[v] = (0.0, 0.0)
    for _ in range(iterations):
        new_pos = dict(pos)
        for v in interior:
            nbrs = list(G_sage.neighbor_iterator(v))
            sx = sum(pos[w][0] for w in nbrs) / len(nbrs)
            sy = sum(pos[w][1] for w in nbrs) / len(nbrs)
            new_pos[v] = (sx, sy)
        pos = new_pos
    return pos

OUT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

C = ['#dc2626', '#16a34a', '#2563eb']
DARK = '#374151'
HIGHLIGHT = '#fef3c7'


def dual_of(G):
    G.is_planar(set_embedding=True)
    faces = G.faces()
    edge_to_faces = {}
    for fi, face in enumerate(faces):
        for u, v in face:
            edge_to_faces.setdefault(frozenset((u, v)), []).append(fi)
    return Graph(
        [(fs[0], fs[1]) for fs in edge_to_faces.values() if len(fs) == 2],
        multiedges=False, loops=False)


def proper_3_edge_colorings(G):
    edges = list(G.edges(labels=False))
    n = len(edges)
    adj = [[] for _ in range(n)]
    for i in range(n):
        u, v = edges[i][0], edges[i][1]
        for j in range(i):
            x, y = edges[j][0], edges[j][1]
            if u in (x, y) or v in (x, y):
                adj[i].append(j); adj[j].append(i)
    coloring = [-1] * n
    results = []

    def back(k):
        if k == n:
            results.append(tuple(coloring)); return
        for c in range(3):
            if all(coloring[j] != c for j in adj[k]):
                coloring[k] = c
                back(k + 1)
        coloring[k] = -1
    back(0)
    return edges, results


def kempe_cycle_set(edges, col_list, start_idx, color_pair):
    a, b = color_pair
    if col_list[start_idx] not in (a, b): return set()
    in_sub = set(i for i in range(len(edges)) if col_list[i] in (a, b))
    visited = {start_idx}; stack = [start_idx]
    while stack:
        cur = stack.pop()
        u, v = edges[cur][0], edges[cur][1]
        for j in in_sub:
            if j in visited: continue
            x, y = edges[j][0], edges[j][1]
            if u in (x, y) or v in (x, y):
                visited.add(j); stack.append(j)
    return visited


def edge_idx(edges, e_frozen):
    for i, e in enumerate(edges):
        if frozenset((e[0], e[1])) == e_frozen:
            return i
    return None


def trace_kempe_cycle(edges, col_list, start_idx, color_pair):
    cycle_set = kempe_cycle_set(edges, col_list, start_idx, color_pair)
    incident_at = {}
    for ei in cycle_set:
        u, v = edges[ei][0], edges[ei][1]
        incident_at.setdefault(u, []).append(ei)
        incident_at.setdefault(v, []).append(ei)
    start_u, start_v = edges[start_idx][0], edges[start_idx][1]
    walk = [(start_idx, start_v)]
    cur_e = start_idx
    cur_leave = start_v
    while True:
        nbrs = incident_at[cur_leave]
        if len(nbrs) != 2: break
        nxt = nbrs[0] if nbrs[1] == cur_e else nbrs[1]
        u2, v2 = edges[nxt][0], edges[nxt][1]
        leave_next = v2 if u2 == cur_leave else u2
        if nxt == start_idx: break
        walk.append((nxt, leave_next))
        cur_e = nxt
        cur_leave = leave_next
    return walk


def matches_chord_apex_kempe(edges, col, named):
    idx = {role: edge_idx(edges, ns) for role, ns in named.items()}
    if any(v is None for v in idx.values()): return False
    c_spike  = col[idx['spike']]
    c_merged = col[idx['merged']]
    if c_spike != c_merged: return False
    c_s0 = col[idx['side_0']]; c_s1 = col[idx['side_1']]
    kc0 = kempe_cycle_set(edges, col, idx['spike'], (c_spike, c_s0))
    if idx['side_0'] not in kc0 or idx['merged'] not in kc0: return False
    kc1 = kempe_cycle_set(edges, col, idx['spike'], (c_spike, c_s1))
    if idx['side_1'] not in kc1 or idx['merged'] not in kc1: return False
    return True


def apply_reduction(G, face, i, v_n_label):
    boundary = [u for (u, v) in face]
    if len(set(boundary)) != 5: return None
    A = []
    for B_k in boundary:
        outer = [w for w in G.neighbor_iterator(B_k) if w not in boundary]
        if len(outer) != 1: return None
        A.append(outer[0])
    if len(set(A)) != 5 or A[(i+3) % 5] == A[(i+4) % 5]: return None
    H = G.copy()
    for v in boundary: H.delete_vertex(v)
    H.add_vertex(v_n_label)
    side_0 = (v_n_label, A[i])
    spike = (v_n_label, A[(i+1) % 5])
    side_1 = (v_n_label, A[(i+2) % 5])
    merged = (A[(i+3) % 5], A[(i+4) % 5])
    H.add_edges([side_0, spike, side_1, merged])
    if H.has_multiple_edges() or H.has_loops(): return None
    if not H.is_planar(set_embedding=True): return None
    if not all(H.degree(v) == 3 for v in H.vertex_iterator()): return None
    return {
        'H': H, 'A': A, 'boundary': boundary, 'face': face, 'i_red': i,
        'named': {
            'spike': frozenset(spike),
            'side_0': frozenset(side_0),
            'side_1': frozenset(side_1),
            'merged': frozenset(merged),
        },
    }


def find_first_match():
    for G in graphs.triangulations(14, minimum_degree=5):
        if not G.is_planar(set_embedding=True): continue
        D = dual_of(G); D.is_planar(set_embedding=True)
        for face in D.faces():
            if len(face) != 5: continue
            for i_red in range(5):
                res = apply_reduction(D, face, i_red, '__v_n_1__')
                if res is None: continue
                H, named = res['H'], res['named']
                edges, gen = proper_3_edge_colorings(H)
                for col in gen:
                    if matches_chord_apex_kempe(edges, col, named):
                        coloring_dict = {frozenset((e[0], e[1])): c
                                         for e, c in zip(edges, col)}
                        return G, D, res, H, coloring_dict
    return None


def find_conj_witness(H, edges, col_list, named):
    GREEN, BLUE = 1, 2
    merged_idx = edge_idx(edges, named['merged'])
    kc_gb = kempe_cycle_set(edges, col_list, merged_idx, (GREEN, BLUE))
    if merged_idx not in kc_gb: return None
    for face in H.faces():
        face_edge_ids = []
        for u, v in face:
            ei = edge_idx(edges, frozenset((u, v)))
            if ei is not None:
                face_edge_ids.append(ei)
        green_on_face_in_kc = [ei for ei in face_edge_ids
                               if col_list[ei] == GREEN
                               and ei in kc_gb
                               and ei != merged_idx]
        if len(green_on_face_in_kc) >= 2:
            return face, green_on_face_in_kc[0], green_on_face_in_kc[1], kc_gb
    return None


def main():
    print("Setting up ...")
    G14, D, red_info, H, coloring = find_first_match()
    i_red = red_info['i_red']
    boundary_in_D = red_info['boundary']
    A_in_D = red_info['A']
    named_in_D = red_info['named']
    print(f"  Found: i_red = {i_red}")
    print(f"  G' has |V|={D.order()}, |E|={D.size()}")

    # Relabel H so all vertices are integers
    H_relabel_map = {v: i for i, v in enumerate(H.vertex_iterator())}
    inv_relabel = {i: v for v, i in H_relabel_map.items()}
    H.relabel(perm=H_relabel_map, inplace=True)
    coloring = {frozenset(H_relabel_map[u] for u in e): c
                for e, c in coloring.items()}
    named_H = {role: frozenset(H_relabel_map[u] for u in e)
               for role, e in named_in_D.items()}

    H.is_planar(set_embedding=True)
    edges_H = list(H.edges(labels=False))
    col_list = [coloring[frozenset((u, v))] for (u, v) in edges_H]
    witness = find_conj_witness(H, edges_H, col_list, named_H)
    face_w, e1, e2, kc_gb = witness
    e1_uv = tuple(edges_H[e1]); e2_uv = tuple(edges_H[e2])
    print(f"  Witness in H_1 (relabeled): e1 = {e1_uv}, e2 = {e2_uv}")

    e1_D = (inv_relabel[e1_uv[0]], inv_relabel[e1_uv[1]])
    e2_D = (inv_relabel[e2_uv[0]], inv_relabel[e2_uv[1]])
    print(f"  Witness in D labels: e1 = {e1_D}, e2 = {e2_D}")

    # Build modified G'
    G_mod = D.copy()
    max_label = max(v for v in D.vertices(sort=False) if isinstance(v, int))
    Y1 = max_label + 1
    Y2 = Y1 + 1
    G_mod.add_vertex(Y1); G_mod.add_vertex(Y2)
    G_mod.delete_edge(e1_D); G_mod.delete_edge(e2_D)
    G_mod.add_edges([(e1_D[0], Y1), (Y1, e1_D[1]),
                     (e2_D[0], Y2), (Y2, e2_D[1]),
                     (Y1, Y2)])
    assert G_mod.is_planar(set_embedding=True)
    print(f"  Modified G': |V|={G_mod.order()}, |E|={G_mod.size()}")

    # Trace Kempe cycle from merged
    merged_idx = edge_idx(edges_H, named_H['merged'])
    walk = trace_kempe_cycle(edges_H, col_list, merged_idx, (1, 2))
    walk_edges = [w[0] for w in walk]
    leave_at = [w[1] for w in walk]

    # Map relabeled-H vertex -> D label if not v_n_1, else None
    def H_to_D(v):
        d = inv_relabel[v]
        return d if d != '__v_n_1__' else None

    # Build new coloring on G_mod's edges.
    # Step A: copy non-named, non-e1, non-e2, non-v_n_1 edges' original colors.
    new_coloring = {}
    named_H_set = set(named_H.values())
    for e_fs, c in coloring.items():
        if e_fs in named_H_set: continue
        if e_fs == frozenset(e1_uv) or e_fs == frozenset(e2_uv): continue
        u, v = tuple(e_fs)
        du, dv = H_to_D(u), H_to_D(v)
        if du is None or dv is None: continue
        new_coloring[frozenset((du, dv))] = c

    # Step B: walk K' (subdivided cycle) starting from merged.
    # For each old cycle edge that is NOT spike/side_0/side_1/merged
    # (i.e., not involving v_n_1, not the chord), we OVERWRITE its color
    # via the alternation. For e1, e2 we instead emit the two subdivision
    # halves. For named edges (spike, side_1, merged), they don't exist
    # in G_mod, so we record the would-be color to use later for f-vector.
    pos = 0
    cycle_label = {}   # e_fs in D labels -> new color (for cycle edges in G_mod)
    half_label = {}    # halves' colors
    would_be = {}      # named role -> would-be color after recoloring
    for k, ei in enumerate(walk_edges):
        leaving_vertex = leave_at[k]
        u, v = edges_H[ei][0], edges_H[ei][1]
        entry_vertex = v if leaving_vertex == u else u
        e_fs_H = frozenset((u, v))
        if ei == e1:
            c0 = 2 if pos % 2 == 0 else 1; pos += 1
            c1 = 2 if pos % 2 == 0 else 1; pos += 1
            half_label[frozenset((inv_relabel[entry_vertex], Y1))] = c0
            half_label[frozenset((Y1, inv_relabel[leaving_vertex]))] = c1
        elif ei == e2:
            c0 = 2 if pos % 2 == 0 else 1; pos += 1
            c1 = 2 if pos % 2 == 0 else 1; pos += 1
            half_label[frozenset((inv_relabel[entry_vertex], Y2))] = c0
            half_label[frozenset((Y2, inv_relabel[leaving_vertex]))] = c1
        else:
            c = 2 if pos % 2 == 0 else 1; pos += 1
            # is this a named edge (spike/side_1/merged)? Note: it can't be
            # side_0 since side_0 is red and the cycle is green/blue.
            for role, ef in named_H.items():
                if ef == e_fs_H:
                    would_be[role] = c
                    break
            else:
                # Regular cycle edge: convert to D labels
                du, dv = H_to_D(u), H_to_D(v)
                if du is None or dv is None:
                    # Shouldn't happen for non-named edges
                    pass
                else:
                    cycle_label[frozenset((du, dv))] = c

    # Apply cycle relabeling
    for e_fs, c in cycle_label.items():
        new_coloring[e_fs] = c
    # Apply half colors
    for e_fs, c in half_label.items():
        new_coloring[e_fs] = c
    # New red edge
    new_coloring[frozenset((Y1, Y2))] = 0

    print(f"  Would-be colors of named H_1 edges after recolor: {would_be}")
    # would_be may not include side_0 (not on cycle); side_0's color is
    # whatever it was in original (= red = 0).
    if 'side_0' not in would_be:
        would_be['side_0'] = coloring[named_H['side_0']]   # 0 (red)

    # Step C: color the 5 externals f_k = B_k - A_k.
    # At each A_k, the third color = would_be color of the missing H_1 edge.
    def third(c1, c2): return ({0, 1, 2} - {c1, c2}).pop()
    externals_colors = [None] * 5
    for k, B_k in enumerate(boundary_in_D):
        A_k = A_in_D[k]
        # Which named role was at A_k in H_1?
        # A_in_D[i_red] = side_0's endpoint, A_in_D[i_red+1] = spike's,
        # A_in_D[i_red+2] = side_1's, A_in_D[i_red+3,+4] = merged.
        if k == i_red:        role = 'side_0'
        elif k == (i_red+1) % 5:  role = 'spike'
        elif k == (i_red+2) % 5:  role = 'side_1'
        else:                  role = 'merged'
        c_ext = would_be[role]
        new_coloring[frozenset((B_k, A_k))] = c_ext
        externals_colors[k] = c_ext
    print(f"  f-vector at F_v: {externals_colors}")

    # Step D: 5 boundary edges B_k - B_{k+1} via Lemma 2.4.
    boundary_colors = [None] * 5
    for k in range(5):
        f_k = externals_colors[k]
        f_kp1 = externals_colors[(k + 1) % 5]
        if f_k != f_kp1:
            boundary_colors[k] = third(f_k, f_kp1)
    for _ in range(20):
        changed = False
        for k in range(5):
            if boundary_colors[k] is not None: continue
            f_k = externals_colors[k]
            prev_c = boundary_colors[(k - 1) % 5]
            next_c = boundary_colors[(k + 1) % 5]
            forbidden = {f_k}
            if prev_c is not None: forbidden.add(prev_c)
            if next_c is not None: forbidden.add(next_c)
            allowed = list({0, 1, 2} - forbidden)
            if allowed:
                boundary_colors[k] = allowed[0]
                changed = True
        if not changed: break

    for k in range(5):
        if boundary_colors[k] is None:
            print(f"  WARNING: undetermined boundary color at k={k}")
            return
        B_k = boundary_in_D[k]; B_kp1 = boundary_in_D[(k + 1) % 5]
        new_coloring[frozenset((B_k, B_kp1))] = boundary_colors[k]

    # Sanity check
    bad = []
    for v in G_mod.vertices(sort=False):
        seen = []
        for w in G_mod.neighbor_iterator(v):
            seen.append(new_coloring.get(frozenset((v, w))))
        if None in seen or len(set(seen)) != len(seen):
            bad.append((v, seen))
    if bad:
        print(f"  PROPRIETY FAILS at vertices:")
        for v, s in bad[:5]:
            print(f"    {v}: {s}")
    else:
        print("  Propriety holds at every vertex of modified G'.")

    # Use H_1's Tutte layout for vertices shared between H_1 and G' (the 19
    # surviving vertices). Then add the 5 boundary vertices B_k of partial F_v
    # by running a local barycenter iteration with the shared vertices fixed.
    H.is_planar(set_embedding=True)
    H_pos = tutte_layout(H,
                         avoid_verts={H_relabel_map['__v_n_1__']})
    # Map H_1 positions back to D labels (skip v_n_1)
    pos_layout = {}
    for v_H, p in H_pos.items():
        v_D = inv_relabel[v_H]
        if v_D != '__v_n_1__':
            pos_layout[v_D] = p
    # Place B_k's by iterating barycenter of their 3 neighbours in G_mod
    # (B_{k-1}, B_{k+1}, A_k); start each B_k near its A_k.
    B_pos = {B_k: pos_layout[A_in_D[k]]
             for k, B_k in enumerate(boundary_in_D)}
    for _ in range(300):
        new_B = {}
        for k, B_k in enumerate(boundary_in_D):
            B_prev = boundary_in_D[(k - 1) % 5]
            B_next = boundary_in_D[(k + 1) % 5]
            A_k = A_in_D[k]
            sx = (B_pos[B_prev][0] + B_pos[B_next][0] + pos_layout[A_k][0]) / 3
            sy = (B_pos[B_prev][1] + B_pos[B_next][1] + pos_layout[A_k][1]) / 3
            new_B[B_k] = (sx, sy)
        B_pos = new_B
    for B_k, p in B_pos.items():
        pos_layout[B_k] = p
    # Y_1, Y_2 at midpoints of their subdivided edges e1, e2
    pos_layout[Y1] = ((pos_layout[e1_D[0]][0] + pos_layout[e1_D[1]][0]) / 2,
                      (pos_layout[e1_D[0]][1] + pos_layout[e1_D[1]][1]) / 2)
    pos_layout[Y2] = ((pos_layout[e2_D[0]][0] + pos_layout[e2_D[1]][0]) / 2,
                      (pos_layout[e2_D[0]][1] + pos_layout[e2_D[1]][1]) / 2)

    fig, ax = plt.subplots(figsize=(8, 8))
    for u, v, _ in G_mod.edges():
        e = frozenset((u, v))
        c = C[new_coloring[e]]
        lw = 3.4 if e == frozenset((Y1, Y2)) else 1.5
        (x0, y0), (x1, y1) = pos_layout[u], pos_layout[v]
        ax.plot([x0, x1], [y0, y1], color=c, lw=lw, zorder=2)
    for v in G_mod.vertices(sort=False):
        x, y = pos_layout[v]
        if v in (Y1, Y2):
            ax.scatter(x, y, s=140, color=DARK, edgecolors='white',
                       linewidths=1.6, zorder=6)
        else:
            ax.scatter(x, y, s=55, color=DARK, zorder=3)
    ax.set_aspect('equal')
    ax.axis('off')
    out_path = os.path.join(OUT_DIR, 'fig_lift_to_Gprime.png')
    fig.savefig(out_path, dpi=170, bbox_inches='tight')
    plt.close(fig)
    print(f"Wrote {out_path}")


if __name__ == '__main__':
    main()
