"""Draw a Conjecture 3.6 witness: on H_1 with its chord-apex+Kempe coloring,
find a face with two green edges that lie (with the merged edge) on a common
{green, blue}-Kempe cycle. Subdivide both green edges with new vertices and
join the two new vertices by a new red edge.

Run with: sage experiments/draw_step1_conj36.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

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

C = ['#dc2626', '#16a34a', '#2563eb']    # 0=red 1=green 2=blue
GRAY = '#9ca3af'
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 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,
        'named': {
            'spike': frozenset(spike),
            'side_0': frozenset(side_0),
            'side_1': frozenset(side_1),
            'merged': frozenset(merged),
        },
    }


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(edges, coloring, start_idx, color_pair):
    a, b = color_pair
    if coloring[start_idx] not in (a, b): return set()
    in_sub = set(i for i in range(len(edges)) if coloring[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 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(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():
    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, face, i_red, H, named, coloring_dict
    return None


def tutte_layout(G_sage, avoid_verts=None, iterations=300):
    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


def find_conj_witness(H, edges, col_list, named):
    """Find a face F of H with two distinct green edges e1, e2, NEITHER equal
    to the merged edge, such that e1, e2, merged all lie on the
    {green, blue}-Kempe cycle through merged."""
    GREEN, BLUE = 1, 2
    merged_idx = edge_idx(edges, named['merged'])
    kc_gb = kempe_cycle(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("Searching for the first n=14 chord-apex+Kempe match ...")
    result = find_first_match()
    G14, D, face_chosen, i_red, H, named, coloring = result
    print(f"  Found: i_red = {i_red}")

    H_relabel_map = {v: i for i, v in enumerate(H.vertex_iterator())}
    H.relabel(perm=H_relabel_map, inplace=True)
    vn = H_relabel_map['__v_n_1__']
    coloring = {frozenset(H_relabel_map[u] for u in e): c
                for e, c in coloring.items()}
    named = {role: frozenset(H_relabel_map[u] for u in e)
             for role, e in named.items()}

    H.is_planar(set_embedding=True)
    pos = tutte_layout(H, avoid_verts={vn})
    E_protected = set(named.values())

    # Build (edges, coloring) in list/tuple form to use kempe helpers
    edges = list(H.edges(labels=False))
    col_list = [coloring[frozenset((u, v))] for (u, v) in edges]

    witness = find_conj_witness(H, edges, col_list, named)
    if witness is None:
        print("ERROR: no witness found.")
        return
    face_w, e1, e2, kc_gb = witness
    e1_uv = tuple(edges[e1]); e2_uv = tuple(edges[e2])
    print(f"  Witness face has {len(face_w)} edges.")
    print(f"  e1 = {e1_uv}, e2 = {e2_uv}")
    print(f"  {{green, blue}}-Kempe cycle through merged: {len(kc_gb)} edges.")

    # Midpoints in the layout
    mp1 = ((pos[e1_uv[0]][0] + pos[e1_uv[1]][0]) / 2,
           (pos[e1_uv[0]][1] + pos[e1_uv[1]][1]) / 2)
    mp2 = ((pos[e2_uv[0]][0] + pos[e2_uv[1]][0]) / 2,
           (pos[e2_uv[0]][1] + pos[e2_uv[1]][1]) / 2)

    # Draw
    fig, ax = plt.subplots(figsize=(8, 8))
    for u, v, _ in H.edges():
        e = frozenset([u, v])
        c = C[coloring[e]]
        lw = 3.8 if e in E_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 H.vertices(sort=False):
        x, y = pos[v]
        if v == vn:
            ax.scatter(x, y, s=320, color=HIGHLIGHT, marker='s',
                       edgecolors='black', linewidths=1.2, zorder=4)
            ax.annotate('$v_n^{(1)}$', (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=70, color=DARK, zorder=3)

    # New red edge between midpoints
    ax.plot([mp1[0], mp2[0]], [mp1[1], mp2[1]],
            color=C[0], lw=4.0, zorder=5)
    # New vertices
    for (mx, my) in (mp1, mp2):
        ax.scatter(mx, my, s=130, color=DARK, edgecolors='white',
                   linewidths=1.6, zorder=6)

    ax.set_aspect('equal')
    ax.axis('off')
    out_path = os.path.join(OUT_DIR, 'fig_alg_step1_conj36.png')
    fig.savefig(out_path, dpi=170, bbox_inches='tight')
    plt.close(fig)
    print(f"Wrote {out_path}")


if __name__ == '__main__':
    main()