"""Draw the iterated reduction trace on the smallest triangulation where
the chord-apex + Kempe-cycle property is satisfied: the first min-degree-5
plantri triangulation on n = 14 vertices, found by search_kempe_property.py.

Overwrites fig_alg_step{0,1,2}.png in the paper directory with this
triangulation's trace (replacing the dodecahedron version).

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


def tutte_layout(G_sage, avoid_verts=None, iterations=300):
    """Tutte's barycentric embedding: pick the largest face whose vertex set
    avoids `avoid_verts` as the outer face, place its vertices on a regular
    polygon, then iterate each interior vertex to the barycenter of its
    neighbors. For 3-connected planar graphs this converges to the unique
    straight-line planar embedding with the chosen outer face --- balanced
    by construction and free of edge crossings."""
    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']
GRAY = '#9ca3af'
DARK = '#374151'
HIGHLIGHT = '#fef3c7'
SHADE = '#fef3c7'


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


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:
        return None
    if 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 % 5])
    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
    named = {
        'spike':  frozenset(spike),
        'side_0': frozenset(side_0),
        'side_1': frozenset(side_1),
        'merged': frozenset(merged),
    }
    return H, named, boundary, A


def proper_3_edge_colorings(G):
    edges = list(G.edges(labels=False))
    n_edges = len(edges)
    adj = [[] for _ in range(n_edges)]
    for i in range(n_edges):
        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_edges

    def back(k):
        if k == n_edges:
            yield tuple(coloring)
            return
        for c in range(3):
            if all(coloring[j] != c for j in adj[k]):
                coloring[k] = c
                yield from back(k + 1)
        coloring[k] = -1

    return edges, back(0)


def kempe_cycle(edges, coloring, start_idx, color_pair):
    a, b = color_pair
    in_sub = [i for i in range(len(edges)) if coloring[i] in (a, b)]
    if start_idx not in in_sub:
        return None
    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 matches_property(edges, col, named):
    idx = {}
    for ii, e in enumerate(edges):
        es = frozenset((e[0], e[1]))
        for role, ns in named.items():
            if es == ns:
                idx[role] = ii
    if len(idx) != 4:
        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(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(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 find_first_match():
    """Iterate over (G, face, i_red, coloring) and return the first hit."""
    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, boundary, A = res
                edges, gen = proper_3_edge_colorings(H)
                for col in gen:
                    if matches_property(edges, col, named):
                        coloring_dict = {frozenset((e[0], e[1])): c
                                         for e, c in zip(edges, col)}
                        return G, D, face, i_red, H, named, boundary, A, coloring_dict
    return None


def draw_graph(ax, G, pos, *, coloring=None, protected=None,
               shade_vertices=None, vn_labels=None):
    if shade_vertices:
        poly = [pos[v] for v in shade_vertices]
        ax.add_patch(Polygon(poly, closed=True, facecolor=SHADE,
                             edgecolor='none', zorder=0))
    protected = protected or set()
    vn_labels = vn_labels or {}
    for u, v, _ in G.edges():
        e = frozenset([u, v])
        c = C[coloring[e]] if (coloring is not None and e in coloring) else GRAY
        lw = 3.8 if e in protected else 1.4
        (x0, y0), (x1, y1) = pos[u], pos[v]
        ax.plot([x0, x1], [y0, y1], color=c, lw=lw, zorder=2)
    for v in G.vertices(sort=False):
        x, y = pos[v]
        if v in vn_labels:
            ax.scatter(x, y, s=320, color=HIGHLIGHT, marker='s',
                       edgecolors='black', linewidths=1.2, zorder=4)
            ax.annotate(vn_labels[v], (x, y),
                        textcoords='offset points', xytext=(16, 16),
                        ha='left', fontsize=14, fontweight='bold',
                        color=DARK, zorder=6,
                        bbox=dict(boxstyle='round,pad=0.2', fc='white',
                                  ec=DARK, lw=0.6))
        else:
            ax.scatter(x, y, s=60, color=DARK, zorder=3)
    ax.set_aspect('equal')
    ax.axis('off')


def main():
    print("Searching for the first match at n = 14 ...")
    result = find_first_match()
    if result is None:
        print("No match found at n = 14.")
        return
    G14, D, face, i_red, H1, named1, boundary1, A1, coloring1 = result
    print(f"Found at i_red = {i_red}")
    print(f"  G (n=14):  |V|={G14.order()}, |E|={G14.size()}, "
          f"min_deg={min(G14.degree())}")
    print(f"  D = G':    |V|={D.order()}, |E|={D.size()}")
    print(f"  H_1:       |V|={H1.order()}, |E|={H1.size()}")

    # Relabel H_1 in place so all vertex labels are comparable integers
    # (Sage's planar layout and face enumeration need comparable labels).
    # Translate coloring1 and named1 accordingly.
    H1_relabel_map = {v: i for i, v in enumerate(H1.vertex_iterator())}
    H1.relabel(perm=H1_relabel_map, inplace=True)
    vn1_int = H1_relabel_map['__v_n_1__']
    coloring1 = {frozenset(H1_relabel_map[u] for u in e): c
                 for e, c in coloring1.items()}
    named1 = {role: frozenset(H1_relabel_map[u] for u in e)
              for role, e in named1.items()}

    D.is_planar(set_embedding=True)
    D_layout = tutte_layout(D, avoid_verts=set(u for (u, v) in face))
    H1.is_planar(set_embedding=True)
    H1_layout = tutte_layout(H1, avoid_verts={vn1_int})

    boundary_face_verts = [u for (u, v) in face]
    fig, ax = plt.subplots(figsize=(8, 8))
    draw_graph(ax, D, D_layout, shade_vertices=boundary_face_verts)
    fig.savefig(os.path.join(OUT_DIR, 'fig_alg_step0.png'),
                dpi=170, bbox_inches='tight')
    plt.close(fig)
    print("Wrote fig_alg_step0.png")

    E = set(named1.values())
    fig, ax = plt.subplots(figsize=(8, 8))
    draw_graph(ax, H1, H1_layout, coloring=coloring1, protected=E,
               vn_labels={vn1_int: '$v_n^{(1)}$'})
    fig.savefig(os.path.join(OUT_DIR, 'fig_alg_step1.png'),
                dpi=170, bbox_inches='tight')
    plt.close(fig)
    print("Wrote fig_alg_step1.png")

    # ----- Step 2: try to continue -----
    H1.is_planar(set_embedding=True)
    chosen2 = None
    for face2 in H1.faces():
        if len(face2) != 5:
            continue
        boundary2 = [u for (u, v) in face2]
        boundary2_edges = [frozenset([u, v]) for (u, v) in face2]
        externals2 = []
        A2 = []
        valid_face = True
        for B_k in boundary2:
            outer = [w for w in H1.neighbor_iterator(B_k) if w not in boundary2]
            if len(outer) != 1:
                valid_face = False
                break
            externals2.append(frozenset([B_k, outer[0]]))
            A2.append(outer[0])
        if not valid_face:
            continue
        if any(e in E for e in boundary2_edges + externals2):
            continue
        # find valid i_t
        f_vec = [coloring1[e] for e in externals2]
        for i_t in range(5):
            if f_vec[(i_t + 3) % 5] != f_vec[(i_t + 4) % 5]:
                continue
            if len({f_vec[i_t], f_vec[(i_t + 1) % 5], f_vec[(i_t + 2) % 5]}) != 3:
                continue
            if A2[(i_t + 3) % 5] == A2[(i_t + 4) % 5]:
                continue
            chosen2 = (face2, i_t, boundary2, externals2, A2)
            break
        if chosen2 is not None:
            break

    if chosen2 is None:
        # algorithm terminates at H_1
        fig, ax = plt.subplots(figsize=(8, 8))
        ax.text(0.5, 0.5,
                "Algorithm terminates at $H_1$:\n"
                "no pentagonal face of $H_1$ has all\n"
                "ten incident edges outside $E$.",
                ha='center', va='center', fontsize=18, color=DARK,
                transform=ax.transAxes,
                bbox=dict(boxstyle='round,pad=0.6', fc=HIGHLIGHT,
                          ec=DARK, lw=1.0))
        ax.set_aspect('equal')
        ax.axis('off')
        fig.savefig(os.path.join(OUT_DIR, 'fig_alg_step2.png'),
                    dpi=170, bbox_inches='tight')
        plt.close(fig)
        print("Wrote fig_alg_step2.png (termination card)")
        print("  Algorithm terminates at H_1: no safe pentagonal face.")
        return

    face2, i_t, boundary2, externals2, A2 = chosen2
    print(f"Step 2: safe face found, i_t = {i_t}")

    H2 = H1.copy()
    for v in boundary2:
        H2.delete_vertex(v)
    # use a fresh int label for v_n^(2)
    v_n_2 = max(H1.vertices(sort=False)) + 1
    H2.add_vertex(v_n_2)
    side_0_2 = (v_n_2, A2[i_t])
    spike_2  = (v_n_2, A2[(i_t + 1) % 5])
    side_1_2 = (v_n_2, A2[(i_t + 2) % 5])
    merged_2 = (A2[(i_t + 3) % 5], A2[(i_t + 4) % 5])
    H2.add_edges([side_0_2, spike_2, side_1_2, merged_2])
    H2.is_planar(set_embedding=True)

    coloring2 = {e: c for e, c in coloring1.items()
                 if not any(u in boundary2 for u in e)}
    coloring2[frozenset(side_0_2)] = coloring1[externals2[i_t]]
    coloring2[frozenset(spike_2)]  = coloring1[externals2[(i_t + 1) % 5]]
    coloring2[frozenset(side_1_2)] = coloring1[externals2[(i_t + 2) % 5]]
    coloring2[frozenset(merged_2)] = coloring1[externals2[(i_t + 3) % 5]]

    E |= {frozenset(side_0_2), frozenset(spike_2),
          frozenset(side_1_2), frozenset(merged_2)}

    H2_layout = tutte_layout(H2, avoid_verts={vn1_int, v_n_2})
    fig, ax = plt.subplots(figsize=(8, 8))
    draw_graph(ax, H2, H2_layout, coloring=coloring2, protected=E,
               vn_labels={vn1_int: '$v_n^{(1)}$',
                          v_n_2: '$v_n^{(2)}$'})
    fig.savefig(os.path.join(OUT_DIR, 'fig_alg_step2.png'),
                dpi=170, bbox_inches='tight')
    plt.close(fig)
    print(f"Wrote fig_alg_step2.png: H_2 with |V|={H2.order()}, "
          f"|E|={H2.size()}, |protected|={len(E)}")

    # --- continue running to completion, checking Kempe condition each step --
    print()
    print("=" * 72)
    print("Running algorithm to completion, checking chord-apex + Kempe at "
          "each step.")
    print("=" * 72)
    # Step 1 status (by construction this is the matching coloring)
    cond1 = check_step_conditions(H1, coloring1, named1)
    print(f"  step t = 1: |V|={H1.order():>3}, |E_graph|={H1.size():>3}, "
          f"|E_prot|=  4 (initial)"
          f" | chord-apex: {cond1['chord_apex']}, "
          f"side_0-Kempe: {cond1['kc_side_0']}, "
          f"side_1-Kempe: {cond1['kc_side_1']}")
    run_to_completion_from(H2, coloring2, E,
                           {'spike': frozenset(spike_2),
                            'side_0': frozenset(side_0_2),
                            'side_1': frozenset(side_1_2),
                            'merged': frozenset(merged_2)},
                           start_t=2)


def check_step_conditions(H, coloring, named):
    """Given an H_t and the *just-added* spike/side_0/side_1/merged, check
    whether chord-apex and the two Kempe-cycle conditions hold."""
    edges = list(H.edges(labels=False))
    edges_fs = [frozenset((u, v)) for (u, v) in edges]
    col = [coloring[e] for e in edges_fs]
    idx = {role: edges_fs.index(e) for role, e in named.items()}
    c_spike  = col[idx['spike']]
    c_merged = col[idx['merged']]
    chord_apex = (c_spike == c_merged)
    if not chord_apex:
        return {'chord_apex': False, 'kc_side_0': False, 'kc_side_1': False}
    c_s0 = col[idx['side_0']]
    c_s1 = col[idx['side_1']]
    kc0 = kempe_cycle(edges, col, idx['spike'], (c_spike, c_s0))
    kc1 = kempe_cycle(edges, col, idx['spike'], (c_spike, c_s1))
    kc_side_0 = (idx['side_0'] in kc0 and idx['merged'] in kc0)
    kc_side_1 = (idx['side_1'] in kc1 and idx['merged'] in kc1)
    return {'chord_apex': True, 'kc_side_0': kc_side_0, 'kc_side_1': kc_side_1}


def find_safe_face(H, protected):
    """Return (face, externals, A) for some safe pentagonal face avoiding
    `protected`, or None."""
    for face in H.faces():
        if len(face) != 5:
            continue
        boundary = [u for (u, v) in face]
        boundary_edges = [frozenset([u, v]) for (u, v) in face]
        externals = []
        A = []
        valid = True
        for B_k in boundary:
            outer = [w for w in H.neighbor_iterator(B_k) if w not in boundary]
            if len(outer) != 1:
                valid = False
                break
            externals.append(frozenset([B_k, outer[0]]))
            A.append(outer[0])
        if not valid:
            continue
        if any(e in protected for e in boundary_edges + externals):
            continue
        return face, boundary, externals, A
    return None


def run_to_completion_from(H, coloring, E, last_named, start_t):
    """Continue iterating from H_{start_t}. The 'last_named' dict carries
    the spike/side/merged of step `start_t` so we can report its Kempe
    status. Print a row per step."""
    t = start_t
    print(f"  step t = {t}: |V|={H.order():>3}, |E_graph|={H.size():>3}, "
          f"|E_prot|={len(E):>3}", end='')
    cond = check_step_conditions(H, coloring, last_named)
    print(f"  | chord-apex: {cond['chord_apex']}, "
          f"side_0-Kempe: {cond['kc_side_0']}, "
          f"side_1-Kempe: {cond['kc_side_1']}")

    while True:
        H.is_planar(set_embedding=True)
        res = find_safe_face(H, E)
        if res is None:
            print(f"  step t = {t + 1}: no safe pentagonal face --> "
                  f"algorithm terminates at H_{t}.")
            return
        face, boundary, externals, A = res
        f_vec = [coloring[e] for e in externals]
        i_t = None
        for i in range(5):
            if f_vec[(i + 3) % 5] != f_vec[(i + 4) % 5]:
                continue
            if len({f_vec[i], f_vec[(i + 1) % 5], f_vec[(i + 2) % 5]}) != 3:
                continue
            if A[(i + 3) % 5] == A[(i + 4) % 5]:
                continue
            i_t = i
            break
        if i_t is None:
            print(f"  step t = {t + 1}: f = {f_vec}, no valid index --> "
                  f"terminate (Lemma 2.4 violation? Probably a parallel-edge "
                  f"or other degenerate case).")
            return
        t += 1
        v_n_new = max(H.vertices(sort=False)) + 1 if all(
            isinstance(v, int) for v in H.vertex_iterator()) else f'vn{t}'
        H_new = H.copy()
        for v in boundary:
            H_new.delete_vertex(v)
        H_new.add_vertex(v_n_new)
        side_0 = (v_n_new, A[i_t])
        spike  = (v_n_new, A[(i_t + 1) % 5])
        side_1 = (v_n_new, A[(i_t + 2) % 5])
        merged = (A[(i_t + 3) % 5], A[(i_t + 4) % 5])
        H_new.add_edges([side_0, spike, side_1, merged])
        H = H_new
        coloring = {e: c for e, c in coloring.items()
                    if not any(u in boundary for u in e)}
        coloring[frozenset(side_0)] = coloring[externals[i_t]] \
            if frozenset(externals[i_t]) in coloring else f_vec[i_t]
        # safer: directly use f_vec
        coloring[frozenset(side_0)] = f_vec[i_t]
        coloring[frozenset(spike)]  = f_vec[(i_t + 1) % 5]
        coloring[frozenset(side_1)] = f_vec[(i_t + 2) % 5]
        coloring[frozenset(merged)] = f_vec[(i_t + 3) % 5]
        named = {
            'spike':  frozenset(spike),
            'side_0': frozenset(side_0),
            'side_1': frozenset(side_1),
            'merged': frozenset(merged),
        }
        E |= set(named.values())
        cond = check_step_conditions(H, coloring, named)
        print(f"  step t = {t}: |V|={H.order():>3}, |E_graph|={H.size():>3}, "
              f"|E_prot|={len(E):>3}, i_t = {i_t}", end='')
        print(f"  | chord-apex: {cond['chord_apex']}, "
              f"side_0-Kempe: {cond['kc_side_0']}, "
              f"side_1-Kempe: {cond['kc_side_1']}")


if __name__ == '__main__':
    main()
