"""Visualise the 'rainbow orbit' intersection from
tire_fiber_step2.tex, Observation rainbow:

  k = 6, T_1 = (6, (0,3), SP) vs T_2 = (3, -, SR)
  |S_1 ∩ S_2| = 6
  The 6 elements are exactly {(a, b, c, b, c, a) : {a,b,c} = {1,2,3}}.

IMPORTANT INTERPRETATION:
  σ ∈ {1,2,3}^6 is the coloring of the SIX G'-EDGES THAT CROSS γ,
  equivalently a coloring of the 6 edges of the shared cycle γ ⊂ G.

  σ is NOT a proper 3-coloring of a 6-cycle: adjacent entries CAN
  coincide, because:
    - T_1's annular cycle has length m_1 + |γ| = 6 + 6 = 12.
    - T_2's annular cycle has length |γ| + k_2 = 6 + 3 = 9.
    - σ records 6 "γ-side" spoke positions interleaved with
      non-γ-side spoke positions on the longer cycle.

  Proper-edge-coloring constraints live on each tire's full annular
  cycle, not on σ directly.  σ is simply the boundary data, like a
  Birkhoff ring coloring.

Produces:
  fig_rainbow_orbit.png    -- all 6 orbit elements as colored γ-edges
  fig_rainbow_pattern.png  -- abstract pattern (a,b,c,b,c,a)
  fig_rainbow_setup.png    -- geometric setup, showing why σ is on
                              γ's EDGES, with the longer annular
                              cycles indicated.
"""
import math
import os
import matplotlib.pyplot as plt
import matplotlib.patches as patches


COLORS = {
    1: '#1f77b4',  # blue
    2: '#d62728',  # red
    3: '#2ca02c',  # green
}


def hex_positions(n=6, radius=1.0, angle_offset=math.pi / 2):
    return [
        (radius * math.cos(angle_offset - 2 * math.pi * i / n),
         radius * math.sin(angle_offset - 2 * math.pi * i / n))
        for i in range(n)
    ]


def draw_gamma_with_edge_colors(ax, sigma, *, radius=1.0, edge_width=5.0,
                                vertex_size=8, vertex_color='#222',
                                show_indices=True):
    """Draw the 6-cycle γ with edge i colored by sigma[i].  Vertices
    are small black dots (they carry no color)."""
    pos = hex_positions(6, radius=radius)
    # Edges
    for i in range(6):
        x1, y1 = pos[i]
        x2, y2 = pos[(i + 1) % 6]
        ax.plot([x1, x2], [y1, y2], color=COLORS[sigma[i]],
                linewidth=edge_width, solid_capstyle='round', zorder=1)
        # Edge index label, placed at midpoint, slightly inside
        if show_indices:
            mx, my = (x1 + x2) / 2, (y1 + y2) / 2
            # nudge toward origin
            ix, iy = mx * 0.78, my * 0.78
            ax.annotate(f"$e_{i}$", (ix, iy), color='#444',
                        ha='center', va='center', fontsize=8, zorder=4)
    # Vertices (small uncolored dots)
    for i in range(6):
        x, y = pos[i]
        ax.plot(x, y, 'o', color=vertex_color, markersize=vertex_size,
                zorder=2)
        if show_indices:
            # vertex label outside
            lx, ly = 1.28 * x, 1.28 * y
            ax.annotate(f"$v_{i}$", (lx, ly), color='#555',
                        ha='center', va='center', fontsize=8, zorder=3)


# ---------------------------------------------------------------------------
# Figure 1: abstract pattern abcbca, σ shown as EDGE colors.
# ---------------------------------------------------------------------------
def fig_pattern(outdir):
    fig, ax = plt.subplots(figsize=(7.0, 5.4))
    pos = hex_positions(6, radius=1.0)

    # Edges with class letters as text labels at midpoint
    classes = ['a', 'b', 'c', 'b', 'c', 'a']
    class_colors = {
        'a': '#bdbdbd',
        'b': '#ffe082',
        'c': '#a5d6a7',
    }
    for i in range(6):
        x1, y1 = pos[i]
        x2, y2 = pos[(i + 1) % 6]
        ax.plot([x1, x2], [y1, y2], color=class_colors[classes[i]],
                linewidth=8.0, solid_capstyle='round', zorder=1)
        mx, my = (x1 + x2) / 2, (y1 + y2) / 2
        ax.annotate(classes[i], (mx, my), color='black', ha='center',
                    va='center', fontsize=13, fontweight='bold', zorder=2)

    for i in range(6):
        x, y = pos[i]
        ax.plot(x, y, 'o', color='#222', markersize=11, zorder=3)
        lx, ly = 1.25 * x, 1.25 * y
        ax.annotate(f"$v_{i}$", (lx, ly), color='#444', ha='center',
                    va='center', fontsize=10, zorder=4)

    legend = (
        r"$\sigma$ colors the 6 edges of $\gamma$" + "\n"
        + r"(equivalently, the 6 dual edges of $G'$" + "\n"
        + r"  that cross $\gamma$)." + "\n\n"
        + r"Pattern: $(e_0,\dots,e_5) = (a,b,c,b,c,a)$" + "\n"
        + r"  $a$ on $\{e_0, e_5\}$" + "\n"
        + r"  $b$ on $\{e_1, e_3\}$" + "\n"
        + r"  $c$ on $\{e_2, e_4\}$" + "\n\n"
        + r"Adjacent edges *may* share a color:" + "\n"
        + r"  $e_5$ and $e_0$ both $= a$ here." + "\n"
        + r"  This is fine because proper-coloring" + "\n"
        + r"  constraints live on each tire's" + "\n"
        + r"  longer annular cycle, not on $\sigma$."
    )
    ax.text(1.85, 0.0, legend, fontsize=9, color='#222',
            verticalalignment='center',
            bbox=dict(boxstyle='round,pad=0.5', facecolor='#f8f8f8',
                      edgecolor='#bbb'))

    ax.set_xlim(-1.6, 5.0)
    ax.set_ylim(-1.7, 1.7)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title(r"Rainbow pattern $(a,b,c,b,c,a)$ as an edge coloring of "
                 + r"$\gamma$  ($|\gamma| = 6$)",
                 fontsize=11)

    out = os.path.join(outdir, 'fig_rainbow_pattern.png')
    plt.savefig(out, dpi=160, bbox_inches='tight')
    plt.close()
    print(f"wrote {out}")


# ---------------------------------------------------------------------------
# Figure 2: the 6 orbit elements as edge colorings of γ.
# ---------------------------------------------------------------------------
def fig_orbit(outdir):
    perms = [
        (1, 2, 3),
        (1, 3, 2),
        (2, 1, 3),
        (2, 3, 1),
        (3, 1, 2),
        (3, 2, 1),
    ]

    fig, axes = plt.subplots(2, 3, figsize=(11.0, 7.6))

    for idx, (a, b, c) in enumerate(perms):
        ax = axes[idx // 3, idx % 3]
        sigma = (a, b, c, b, c, a)
        draw_gamma_with_edge_colors(ax, sigma)
        ax.set_xlim(-1.6, 1.6)
        ax.set_ylim(-1.6, 1.6)
        ax.set_aspect('equal')
        ax.axis('off')
        title = (f"$(a,b,c) = ({a},{b},{c})$\n"
                 + r"$\sigma = ($" + ",".join(str(x) for x in sigma) + r"$)$")
        ax.set_title(title, fontsize=10)

    fig.suptitle(r"The 6-element rainbow orbit, drawn as edge colorings of "
                 + r"the shared cycle $\gamma$" + "\n"
                 + r"$\sigma$ is the coloring of the 6 dual edges of $G'$ "
                 + r"crossing $\gamma$, equivalently the 6 edges of $\gamma$"
                 + "\n"
                 + r"Case: $T_1{=}(m{=}6,\, \mathrm{chord}{=}(0,3),\, "
                 + r"\mathrm{SP})$  vs.  $T_2{=}(k{=}3,\, \mathrm{SR})$,  "
                 + r"$|S_1 \cap S_2| = 6$",
                 fontsize=11, y=1.00)
    plt.tight_layout()
    out = os.path.join(outdir, 'fig_rainbow_orbit.png')
    plt.savefig(out, dpi=160, bbox_inches='tight')
    plt.close()
    print(f"wrote {out}")


# ---------------------------------------------------------------------------
# Figure 3: geometric setup.  Shows γ between T_1 and T_2; σ as colored
# G'-edges across γ; each tire's annular cycle indicated with the correct
# number of vertices and interleaved non-γ spokes.
# ---------------------------------------------------------------------------
def fig_setup(outdir):
    fig, ax = plt.subplots(figsize=(11.5, 7.2))

    # γ at radius R_gamma; we draw it as a hexagon.
    R_gamma = 1.5
    gamma_pos = hex_positions(6, radius=R_gamma)

    # T_1's annular cycle is a 12-cycle in G'.  In Π_G it sits between
    # γ and a larger outer boundary; vertices of T_1's annular cycle are
    # DUAL vertices of T_1's annular faces.  6 of them are on the
    # γ-side, 6 on the m_1 side.  We draw them as 12 points spread
    # outside γ.
    R_T1_inner = 2.0  # γ-side of T_1's annular cycle vertices
    R_T1_outer = 3.2  # m_1-side of T_1's annular cycle vertices
    # Position T_1 dual vertices: alternate inner / outer as one goes around.
    # In a balanced spoke-only triangulation, the 12 dual vertices alternate
    # between γ-side (I-move) and outer-side (O-move).
    T1_inner_pts = []
    T1_outer_pts = []
    for i in range(6):
        # γ-side dual vertex: just outside γ-edge i, halfway between v_i and v_{i+1}
        x1, y1 = gamma_pos[i]
        x2, y2 = gamma_pos[(i + 1) % 6]
        mx, my = (x1 + x2) / 2, (y1 + y2) / 2
        r_m = math.hypot(mx, my)
        ux, uy = mx / r_m, my / r_m
        T1_inner_pts.append((R_T1_inner * ux, R_T1_inner * uy))
        # m_1-side dual vertex: just past v_i, outside.
        r_v = math.hypot(x1, y1)
        ux, uy = x1 / r_v, y1 / r_v
        T1_outer_pts.append((R_T1_outer * ux, R_T1_outer * uy))

    # T_1 outer boundary B_out (length m_1 = 6) and the 12-cycle of
    # annular face dual vertices: alternate T1_inner_pts and T1_outer_pts
    # going around.  Draw the 12-cycle edges as light gray.
    T1_cycle = []
    for i in range(6):
        T1_cycle.append(T1_inner_pts[i])
        T1_cycle.append(T1_outer_pts[(i + 1) % 6])
    for i in range(12):
        x1, y1 = T1_cycle[i]
        x2, y2 = T1_cycle[(i + 1) % 12]
        ax.plot([x1, x2], [y1, y2], color='#90a4ae', linewidth=1.0,
                linestyle=':', zorder=1)
    for (x, y) in T1_inner_pts:
        ax.plot(x, y, 'o', color='#607d8b', markersize=7, zorder=3)
    for (x, y) in T1_outer_pts:
        ax.plot(x, y, 'o', color='#90a4ae', markersize=5, zorder=3)

    # T_1's chord (in O_1, antipodal) -- draw outside γ as an arc.
    chord_p1 = T1_outer_pts[0]
    chord_p2 = T1_outer_pts[3]
    chord_arc = patches.FancyArrowPatch(chord_p1, chord_p2,
                                        connectionstyle='arc3,rad=0.45',
                                        arrowstyle='-', color='#9c27b0',
                                        linewidth=1.8, linestyle=(0, (3, 2)),
                                        zorder=2)
    ax.add_patch(chord_arc)

    # T_2's annular cycle (length |γ| + k_2 = 6 + 3 = 9) inside γ.
    # 6 γ-side dual vertices and 3 k_2-side dual vertices.
    R_T2_inner = 0.45  # k_2-side
    R_T2_outer = 1.05  # γ-side
    T2_outer_pts = []
    for i in range(6):
        x1, y1 = gamma_pos[i]
        x2, y2 = gamma_pos[(i + 1) % 6]
        mx, my = (x1 + x2) / 2, (y1 + y2) / 2
        r_m = math.hypot(mx, my)
        ux, uy = mx / r_m, my / r_m
        T2_outer_pts.append((R_T2_outer * ux, R_T2_outer * uy))
    T2_inner_pts = []
    for i in range(3):
        angle = math.pi / 2 + 2 * math.pi * i / 3
        T2_inner_pts.append((R_T2_inner * math.cos(angle),
                             R_T2_inner * math.sin(angle)))
    # T_2's 9-cycle: 2 γ-spokes, 1 k_2-spoke, ... let's just draw the cycle
    # cleanly without committing to an exact interleaving.
    T2_cycle = []
    for i in range(6):
        T2_cycle.append(T2_outer_pts[i])
        if i % 2 == 1:
            T2_cycle.append(T2_inner_pts[i // 2])
    for i in range(9):
        x1, y1 = T2_cycle[i]
        x2, y2 = T2_cycle[(i + 1) % 9]
        ax.plot([x1, x2], [y1, y2], color='#ffb74d', linewidth=1.0,
                linestyle=':', zorder=1)
    for (x, y) in T2_outer_pts:
        ax.plot(x, y, 'o', color='#f57c00', markersize=7, zorder=3)
    for (x, y) in T2_inner_pts:
        ax.plot(x, y, 'o', color='#ffb74d', markersize=5, zorder=3)

    # Now draw γ with edges colored by σ = (1, 2, 3, 2, 3, 1).
    sigma = (1, 2, 3, 2, 3, 1)
    for i in range(6):
        x1, y1 = gamma_pos[i]
        x2, y2 = gamma_pos[(i + 1) % 6]
        ax.plot([x1, x2], [y1, y2], color=COLORS[sigma[i]], linewidth=5.0,
                solid_capstyle='round', zorder=4)
        mx, my = (x1 + x2) / 2, (y1 + y2) / 2
        # Also draw the dual G'-edges crossing γ at midpoint -- short segments
        # from T_2-side dual vertex to T_1-side dual vertex.
        sx1, sy1 = T2_outer_pts[i]
        sx2, sy2 = T1_inner_pts[i]
        ax.plot([sx1, sx2], [sy1, sy2], color=COLORS[sigma[i]],
                linewidth=2.5, linestyle=(0, (4, 2)), zorder=2)
    # γ vertices
    for i in range(6):
        x, y = gamma_pos[i]
        ax.plot(x, y, 'o', color='#222', markersize=10, zorder=5)
        lx, ly = R_gamma * 1.12 * x / math.hypot(x, y), R_gamma * 1.12 * y / math.hypot(x, y)
        ax.annotate(f"$v_{i}$", (lx + 0.05 * x, ly + 0.05 * y),
                    color='#222', ha='center', va='center', fontsize=9,
                    fontweight='bold', zorder=6)

    # Labels with arrows
    ax.annotate(r'$T_1$ annular cycle, length $|γ|+m_1 = 12$' + '\n'
                + r'γ-side dual vertices (dark): 6' + '\n'
                + r'$m_1$-side dual vertices (light): 6',
                xy=(R_T1_inner * 1.0, 0.0), xytext=(4.0, 1.0),
                fontsize=9, color='#37474f',
                ha='left', va='center',
                arrowprops=dict(arrowstyle='->', color='#37474f', lw=0.8))
    ax.annotate(r'$T_2$ annular cycle, length $|γ|+k_2 = 9$' + '\n'
                + r'γ-side dual vertices (orange): 6' + '\n'
                + r'$k_2$-side dual vertices (light): 3',
                xy=(R_T2_outer * 0.7, -0.7), xytext=(4.0, -1.0),
                fontsize=9, color='#e65100',
                ha='left', va='center',
                arrowprops=dict(arrowstyle='->', color='#e65100', lw=0.8))
    ax.annotate(r'antipodal chord' + '\n' + r'$(v_0, v_3) \in O_1$',
                xy=(0, R_T1_outer * 1.0 - 0.1), xytext=(-4.6, 2.4),
                fontsize=9, color='#9c27b0',
                ha='left', va='center',
                arrowprops=dict(arrowstyle='->', color='#9c27b0', lw=0.8))
    ax.annotate(r'shared cycle $\gamma$' + '\n'
                + r'$|γ|=6$, drawn with' + '\n'
                + r'$\sigma=(1,2,3,2,3,1)$',
                xy=(R_gamma * 0.7, 0.7), xytext=(-5.2, 0.0),
                fontsize=10, color='#222',
                ha='left', va='center',
                arrowprops=dict(arrowstyle='->', color='#222', lw=0.8))
    ax.annotate(r'$\sigma$ is the coloring of the 6 dual' + '\n'
                + r'$G\,\!\!\!\!^{\,\prime}$-edges crossing $γ$' + '\n'
                + r'(dashed, same colors as $γ$-edges)',
                xy=(R_gamma * 1.05, -R_gamma * 0.05), xytext=(-5.2, -1.4),
                fontsize=9, color='#333',
                ha='left', va='center',
                arrowprops=dict(arrowstyle='->', color='#333', lw=0.8))

    ax.set_xlim(-5.5, 7.5)
    ax.set_ylim(-3.8, 3.8)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title(r"Geometric setup: $\sigma = (1,2,3,2,3,1)$ is one of the "
                 + r"6 orbit elements. " + "\n"
                 + r"It colors the 6 $G'$-edges that cross $\gamma$.  "
                 + r"Each tire's full annular cycle is longer than 6, with "
                 + r"interleaved non-$\gamma$ spokes.",
                 fontsize=11)

    out = os.path.join(outdir, 'fig_rainbow_setup.png')
    plt.savefig(out, dpi=160, bbox_inches='tight')
    plt.close()
    print(f"wrote {out}")


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    notes = os.path.normpath(os.path.join(here, '..', 'notes'))
    fig_pattern(notes)
    fig_orbit(notes)
    fig_setup(notes)


if __name__ == '__main__':
    main()
