r"""Draw diagrams illustrating the proof of Lemma 2.6 (chord-apex).

Three figures rendered on the dodecahedral G':
  step1: the assumed proper 3-edge-colouring of \widehat{G}'_{v,0}, with the
         v_n-edges receiving distinct colours X (red), Y (green, spike),
         Z (blue), and the chord coloured W with W != Y, forcing W in {X, Z}.
  step2: lift to G' in the case W = Z. The externals inherit
         psi(f) = (X, Y, Z, Z, Z) and Lemma 2.4 completes partial F_v.
  step3: lift to G' in the case W = X with psi(f) = (X, Y, Z, X, X).
"""
import os
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D

from reduced_dual import build_dual, apply_reduction


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

C1 = '#dc2626'      # red    -> X
C2 = '#16a34a'      # green  -> Y (spike)
C3 = '#2563eb'      # blue   -> Z
DARK = '#374151'
DEG2 = '#f59e0b'
GRAY = '#9ca3af'
LIGHT = '#e5e7eb'
APEX_FILL = '#fef3c7'


def base_canvas(title):
    fig, ax = plt.subplots(figsize=(8.5, 8.5))
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title(title, fontsize=12)
    return fig, ax


def draw_lift(ax, Gp, pos, A, B, other_surv, f_cols, e_cols, f_labels):
    """G' with f_k and partial F_v boundary edges coloured; rest in light gray."""
    pentagon_edges = set()
    for k in range(5):
        pentagon_edges.add(frozenset((A[k], B[k])))               # f_k
        pentagon_edges.add(frozenset((B[k], B[(k + 1) % 5])))     # boundary
    for u, v in Gp.edges():
        if frozenset((u, v)) in pentagon_edges:
            continue
        (x0, y0), (x1, y1) = pos[u], pos[v]
        ax.plot([x0, x1], [y0, y1], color=LIGHT, lw=1.2, zorder=1)
    other = [v for v in other_surv if v not in A and v not in B]
    xs = [pos[v][0] for v in other]
    ys = [pos[v][1] for v in other]
    ax.scatter(xs, ys, s=50, color=LIGHT, zorder=2)
    for k, v in enumerate(A):
        ax.scatter(*pos[v], s=240, color=DEG2, edgecolors='black',
                   linewidths=0.8, zorder=4)
        ax.annotate(f'$A_{k}$', (pos[v][0] * 1.18, pos[v][1] * 1.18),
                    ha='center', va='center', fontsize=12, color='#a16207',
                    zorder=5)
    for k, v in enumerate(B):
        ax.scatter(*pos[v], s=160, color='white', edgecolors=DARK,
                   linewidths=1.2, zorder=4)
        ax.annotate(f'$B_{k}$', (pos[v][0] * 0.62, pos[v][1] * 0.62),
                    ha='center', va='center', fontsize=11, color=DARK,
                    zorder=5)
    for k in range(5):
        (x0, y0), (x1, y1) = pos[B[k]], pos[A[k]]
        ax.plot([x0, x1], [y0, y1], color=f_cols[k], lw=3.2, zorder=3)
        mx, my = 0.55 * x0 + 0.45 * x1, 0.55 * y0 + 0.45 * y1
        ax.annotate(f_labels[k], (mx, my), fontsize=10, color=f_cols[k],
                    fontweight='bold', ha='center', va='center', zorder=7,
                    bbox=dict(boxstyle='round,pad=0.15', fc='white',
                              ec=f_cols[k], lw=0.5))
    for k in range(5):
        (x0, y0), (x1, y1) = pos[B[k]], pos[B[(k + 1) % 5]]
        ax.plot([x0, x1], [y0, y1], color=e_cols[k], lw=3.2, zorder=3)


def main():
    Gp, pos, Fv = build_dual()
    res = apply_reduction(Gp, pos, Fv, i=0)
    npos, A = res['pos'], res['A']
    v_n, apex_nbrs, chord = res['v_n'], res['apex_nbrs'], res['chord']

    a_pair = {}
    for u in A:
        for nbr in Gp.neighbors(u):
            if nbr in Fv:
                a_pair[u] = nbr
                break
    B = [a_pair[A[k]] for k in range(5)]

    survivors = [v for v in Gp if v not in Fv]
    surv_set = set(survivors)
    surv_edges = [(u, v) for u, v in Gp.edges()
                  if u in surv_set and v in surv_set]
    other_surv = [v for v in survivors if v not in A]

    # ----- Step 1: assumed colouring of the reduced dual -----
    fig, ax = base_canvas(
        "Step 1: $\\varphi$ on $\\widehat{G}'_{v,0}$ assigns distinct colours "
        "$X, Y, Z$ to the $v_n$-edges (propriety at $v_n$);\n"
        "by hypothesis $W \\neq Y$, forcing $W \\in \\{X, Z\\}$.")
    for u, v in surv_edges:
        (x0, y0), (x1, y1) = pos[u], pos[v]
        ax.plot([x0, x1], [y0, y1], color=GRAY, lw=1.4, zorder=1)
    xs = [pos[v][0] for v in other_surv]
    ys = [pos[v][1] for v in other_surv]
    ax.scatter(xs, ys, s=80, color=DARK, zorder=3)
    for k, v in enumerate(A):
        ax.scatter(*pos[v], s=240, color=DEG2, edgecolors='black',
                   linewidths=0.8, zorder=4)
        ax.annotate(f'$A_{k}$', (pos[v][0] * 1.18, pos[v][1] * 1.18),
                    ha='center', va='center', fontsize=12, color='#a16207',
                    zorder=5)
    for k, (color, lbl) in enumerate([(C1, '$X$'), (C2, '$Y$'), (C3, '$Z$')]):
        u = A[k]
        (x0, y0), (x1, y1) = npos[v_n], pos[u]
        ax.plot([x0, x1], [y0, y1], color=color, lw=3.4, zorder=5)
        mx, my = 0.5 * x1 + 0.5 * x0, 0.5 * y1 + 0.5 * y0
        ax.annotate(lbl, (mx, my), fontsize=12, color=color, fontweight='bold',
                    ha='center', va='center', zorder=7,
                    bbox=dict(boxstyle='round,pad=0.18', fc='white', ec=color,
                              lw=0.6))
    ax.scatter(*npos[v_n], s=300, color=APEX_FILL, marker='s',
               edgecolors='black', linewidths=1.0, zorder=6)
    ax.annotate('$v_n$', npos[v_n], textcoords='offset points', xytext=(0, 14),
                ha='center', fontsize=12, fontweight='bold', color=DARK,
                zorder=7)
    (x0, y0), (x1, y1) = pos[chord[0]], pos[chord[1]]
    ax.plot([x0, x1], [y0, y1], color=DARK, lw=2.4, ls='--', zorder=5)
    mx, my = (x0 + x1) / 2, (y0 + y1) / 2
    ax.annotate('$W$', (mx, my), fontsize=12, color=DARK, fontweight='bold',
                ha='center', va='center', zorder=7,
                bbox=dict(boxstyle='round,pad=0.18', fc='white', ec=DARK,
                          lw=0.6))
    ax.legend(handles=[
        Line2D([0], [0], color=C1, lw=3, label='$X$ (side)'),
        Line2D([0], [0], color=C2, lw=3, label='$Y$ (spike)'),
        Line2D([0], [0], color=C3, lw=3, label='$Z$ (side)'),
        Line2D([0], [0], color=DARK, lw=2, ls='--',
               label='$W$ (merged), $W \\neq Y$'),
    ], loc='upper left', fontsize=10)
    fig.savefig(os.path.join(OUT_DIR, 'fig_chord_apex_step1.png'),
                dpi=170, bbox_inches='tight')
    plt.close(fig)

    cases = [
        dict(
            title=("Step 2: lift to $G'$ when $W = Z$. The externals inherit "
                   "$\\psi(f) = (X, Y, Z, Z, Z)$;\n"
                   "Lemma 2.4 colours the five edges of $\\partial F_v$."),
            f_cols=[C1, C2, C3, C3, C3],
            e_cols=[C3, C1, C2, C1, C2],
            f_labels=['$X$', '$Y$', '$Z$', '$Z$', '$Z$'],
            out='fig_chord_apex_step2.png',
        ),
        dict(
            title=("Step 3: lift to $G'$ when $W = X$. The externals inherit "
                   "$\\psi(f) = (X, Y, Z, X, X)$;\n"
                   "Lemma 2.4 colours the five edges of $\\partial F_v$."),
            f_cols=[C1, C2, C3, C1, C1],
            e_cols=[C3, C1, C2, C3, C2],
            f_labels=['$X$', '$Y$', '$Z$', '$X$', '$X$'],
            out='fig_chord_apex_step3.png',
        ),
    ]

    for case in cases:
        fig, ax = base_canvas(case['title'])
        draw_lift(ax, Gp, pos, A, B, other_surv,
                  case['f_cols'], case['e_cols'], case['f_labels'])
        fig.savefig(os.path.join(OUT_DIR, case['out']),
                    dpi=170, bbox_inches='tight')
        plt.close(fig)

    print(f"wrote fig_chord_apex_step{{1,2,3}}.png to {OUT_DIR}")


if __name__ == '__main__':
    main()
