"""Draw the iterated reduction algorithm's trace on the dodecahedron.

Produces three figures:
  fig_alg_step0.png -- G' (dodecahedron) with F_v (inner pentagon) shaded.
  fig_alg_step1.png -- H_1 (post step 1), 3-edge-coloured; 4 protected edges.
  fig_alg_step2.png -- H_2 (post step 2), 3-edge-coloured; 8 protected edges;
                       algorithm terminates.

Run with: sage experiments/draw_iterated_reduction.py
"""
from sage.all import Graph
from sage.graphs.graph_coloring import edge_coloring
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']     # proper-edge-colour palette
GRAY = '#9ca3af'
DARK = '#374151'
HIGHLIGHT = '#fef3c7'


def dodecahedron_positions():
    pos = {}
    R = {'a': 1.0, 'b': 2.2, 'c': 3.6, 'd': 4.8}
    for i in range(5):
        for fam in ('a', 'b'):
            th = math.radians(90 - 72 * i)
            pos[(fam, i)] = (R[fam] * math.cos(th), R[fam] * math.sin(th))
        for fam in ('c', 'd'):
            th = math.radians(90 - 72 * i - 36)
            pos[(fam, i)] = (R[fam] * math.cos(th), R[fam] * math.sin(th))
    return pos


def build_dodecahedron():
    edges = []
    for i in range(5):
        edges.append((('a', i), ('a', (i + 1) % 5)))
        edges.append((('a', i), ('b', i)))
        edges.append((('b', i), ('c', i)))
        edges.append((('b', i), ('c', (i - 1) % 5)))
        edges.append((('c', i), ('d', i)))
        edges.append((('d', i), ('d', (i + 1) % 5)))
    G = Graph(edges, multiedges=False, loops=False)
    G.is_planar(set_embedding=True)
    return G


def find_safe_pentagonal_face(G, protected):
    for face in G.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 = []
        for B_k in boundary:
            outer = [w for w in G.neighbor_iterator(B_k) if w not in boundary]
            if len(outer) != 1:
                break
            externals.append(frozenset([B_k, outer[0]]))
            A.append(outer[0])
        else:
            if not any(e in protected for e in boundary_edges + externals):
                return boundary, externals, A
    return None


def valid_indices(f_vec):
    out = []
    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:
            out.append(i)
    return out


def draw(ax, G, pos, *, coloring=None, protected=None,
         shade_face=None):
    if shade_face:
        poly = [pos[v] for v in shade_face]
        ax.add_patch(Polygon(poly, closed=True, facecolor=HIGHLIGHT,
                             edgecolor='none', zorder=0))
    protected = protected or set()
    for u, v in G.edges(labels=False):
        e = frozenset([u, v])
        c = C[coloring[e]] if coloring is not None 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 isinstance(v, tuple) and v[0] == 'v_n':
            t = v[1]
            ax.scatter(x, y, s=320, color=HIGHLIGHT, marker='s',
                       edgecolors='black', linewidths=1.2, zorder=4)
            ax.annotate(f'$v_n^{{({t})}}$', (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)
    ax.set_aspect('equal')
    ax.axis('off')


def main():
    G = build_dodecahedron()
    pos = dodecahedron_positions()
    F_v = [('a', i) for i in range(5)]

    # ----- Step 0: G' with F_v shaded -----
    fig, ax = plt.subplots(figsize=(8, 8))
    draw(ax, G, pos, shade_face=F_v)
    fig.savefig(os.path.join(OUT_DIR, 'fig_alg_step0.png'),
                dpi=170, bbox_inches='tight')
    plt.close(fig)

    # ----- Step 1: Definition 2.1 at F_v with i_1 = 0 -----
    safe = find_safe_pentagonal_face(G, set())
    boundary_1, externals_1, A_1 = safe
    G1 = G.copy()
    for v in boundary_1:
        G1.delete_vertex(v)
    v_n_1 = ('v_n', 1)
    G1.add_vertex(v_n_1)
    G1.add_edge(v_n_1, A_1[0])
    G1.add_edge(v_n_1, A_1[1])
    G1.add_edge(v_n_1, A_1[2])
    G1.add_edge(A_1[3], A_1[4])
    G1.is_planar(set_embedding=True)
    pos1 = {v: p for v, p in pos.items() if v not in boundary_1}
    cx = (pos[A_1[0]][0] + pos[A_1[1]][0] + pos[A_1[2]][0]) / 3
    cy = (pos[A_1[0]][1] + pos[A_1[1]][1] + pos[A_1[2]][1]) / 3
    pos1[v_n_1] = (cx * 0.55, cy * 0.55)

    cols = edge_coloring(G1, value_only=False)
    coloring = {}
    for k, edge_list in enumerate(cols):
        for u, v in edge_list:
            coloring[frozenset([u, v])] = k

    E = {
        frozenset([v_n_1, A_1[1]]),    # spike
        frozenset([v_n_1, A_1[0]]),    # side-0
        frozenset([v_n_1, A_1[2]]),    # side-1
        frozenset([A_1[3], A_1[4]]),   # merged
    }

    fig, ax = plt.subplots(figsize=(8, 8))
    draw(ax, G1, pos1, coloring=coloring, protected=E)
    fig.savefig(os.path.join(OUT_DIR, 'fig_alg_step1.png'),
                dpi=170, bbox_inches='tight')
    plt.close(fig)

    # ----- Step 2: reduce at the only remaining safe face (outer pentagon) -----
    safe = find_safe_pentagonal_face(G1, E)
    if safe is None:
        print("ERROR: expected an outer pentagonal face but none found.")
        return
    boundary_2, externals_2, A_2 = safe
    f_vec = [coloring[e] for e in externals_2]
    choices = valid_indices(f_vec)
    if not choices:
        print(f"ERROR: f-vector {f_vec} has no valid index.")
        return
    i_t = choices[0]

    G2 = G1.copy()
    for v in boundary_2:
        G2.delete_vertex(v)
    v_n_2 = ('v_n', 2)
    G2.add_vertex(v_n_2)
    G2.add_edge(v_n_2, A_2[i_t])
    G2.add_edge(v_n_2, A_2[(i_t + 1) % 5])
    G2.add_edge(v_n_2, A_2[(i_t + 2) % 5])
    G2.add_edge(A_2[(i_t + 3) % 5], A_2[(i_t + 4) % 5])
    G2.is_planar(set_embedding=True)

    coloring2 = {e: c for e, c in coloring.items()
                 if not any(u in boundary_2 for u in e)}
    side_0_2 = frozenset([v_n_2, A_2[i_t]])
    spike_2  = frozenset([v_n_2, A_2[(i_t + 1) % 5]])
    side_1_2 = frozenset([v_n_2, A_2[(i_t + 2) % 5]])
    merged_2 = frozenset([A_2[(i_t + 3) % 5], A_2[(i_t + 4) % 5]])
    coloring2[side_0_2] = coloring[externals_2[i_t]]
    coloring2[spike_2]  = coloring[externals_2[(i_t + 1) % 5]]
    coloring2[side_1_2] = coloring[externals_2[(i_t + 2) % 5]]
    coloring2[merged_2] = coloring[externals_2[(i_t + 3) % 5]]

    pos2 = {v: p for v, p in pos1.items() if v not in boundary_2}
    nbrs = [A_2[i_t], A_2[(i_t + 1) % 5], A_2[(i_t + 2) % 5]]
    cx = sum(pos2[a][0] for a in nbrs) / 3
    cy = sum(pos2[a][1] for a in nbrs) / 3
    r = math.hypot(cx, cy)
    # v_n^{(2)} lies outside the surviving graph (the deleted d's were outermost)
    target_r = 5.0
    pos2[v_n_2] = (cx * target_r / r, cy * target_r / r)

    E |= {side_0_2, spike_2, side_1_2, merged_2}

    fig, ax = plt.subplots(figsize=(8, 8))
    draw(ax, G2, pos2, coloring=coloring2, protected=E)
    fig.savefig(os.path.join(OUT_DIR, 'fig_alg_step2.png'),
                dpi=170, bbox_inches='tight')
    plt.close(fig)

    print(f"Wrote fig_alg_step{{0,1,2}}.png to {OUT_DIR}")


if __name__ == '__main__':
    main()