"""Worked example for fiber_decomposition.tex.
Three panels: C_4 with 4 pendant spokes, showing
  (a) sigma = (1,1,1,1) -- fiber count 2 (two alternating cycle colorings).
  (b) sigma = (1,2,1,2) -- infeasible, fiber count 0.
  (c) sigma = (1,1,2,2) -- fiber count 1 (unique extension).
"""
import math
import os
import matplotlib.pyplot as plt
from itertools import product


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


def square_pos(scale=1.0):
    # v_0 right, v_1 top, v_2 left, v_3 bottom -- a square
    return {
        0: (scale, 0),
        1: (0, scale),
        2: (-scale, 0),
        3: (0, -scale),
    }


def ext_pos(scale=1.0, radial=1.9):
    pos = square_pos(scale)
    out = {}
    for v, (x, y) in pos.items():
        nx, ny = x / math.hypot(x, y), y / math.hypot(x, y)
        out[v] = (radial * nx, radial * ny)
    return out


def enumerate_fiber(sigma):
    """Brute-force: enumerate all cycle-edge colorings (c_e0, c_e1, c_e2, c_e3)
    in [3]^4 satisfying:
      proper edge coloring at each v_i (its three incident edges all distinct)
      and the pendant at v_i is sigma[i].
    e_i goes from v_i to v_{i+1}.  v_i is incident to e_{i-1} (mod 4) and e_i.
    Returns the list of valid (c_e0, c_e1, c_e2, c_e3) tuples.
    """
    valid = []
    for ce in product([1, 2, 3], repeat=4):
        ok = True
        for i in range(4):
            ep = ce[(i - 1) % 4]
            en = ce[i]
            s = sigma[i]
            if len({ep, en, s}) != 3:
                ok = False
                break
        if ok:
            valid.append(ce)
    return valid


def draw_panel(ax, sigma, title):
    pos = square_pos()
    ext = ext_pos()
    # Cycle edges (drawn light first, overdrawn in colored if there is a unique
    # extension; if there are multiple extensions show grey-with-letters)
    fiber = enumerate_fiber(sigma)

    # Cycle edges: e_i from v_i to v_{i+1}
    for i in range(4):
        u, v = i, (i + 1) % 4
        x1, y1 = pos[u]; x2, y2 = pos[v]
        if len(fiber) == 1:
            ax.plot([x1, x2], [y1, y2], color=COLORS[fiber[0][i]],
                    linewidth=4.0, zorder=1)
        elif len(fiber) >= 2:
            # show all extensions as parallel offset lines
            colors_for_edge = sorted({f[i] for f in fiber})
            n = len(colors_for_edge)
            # tangent perpendicular to edge
            dx, dy = x2 - x1, y2 - y1
            ln = math.hypot(dx, dy)
            px, py = -dy / ln, dx / ln  # unit normal
            spacing = 0.08
            offsets = [(j - (n - 1) / 2) * spacing for j in range(n)]
            for off, c in zip(offsets, colors_for_edge):
                ax.plot([x1 + off * px, x2 + off * px],
                        [y1 + off * py, y2 + off * py],
                        color=COLORS[c], linewidth=3.0, zorder=1)
        else:
            # infeasible: dashed grey
            ax.plot([x1, x2], [y1, y2], color='#bbbbbb', linewidth=2.5,
                    linestyle=(0, (4, 3)), zorder=1)

    # Spoke edges: solid in sigma's color
    for i in range(4):
        x1, y1 = pos[i]; x2, y2 = ext[i]
        ax.plot([x1, x2], [y1, y2], color=COLORS[sigma[i]], linewidth=4.0,
                zorder=1)

    # Cycle vertices
    for v, (x, y) in pos.items():
        ax.plot(x, y, 'o', color='#1f3f70', markersize=18, zorder=2)
        ax.annotate(f"$v_{v}$", (x, y), color='white', ha='center',
                    va='center', fontsize=10, fontweight='bold', zorder=3)
    # Spoke vertices
    for v, (x, y) in ext.items():
        ax.plot(x, y, 's', color='#666666', markersize=15, zorder=2)
        ax.annotate(f"$u_{v}$", (x, y), color='white', ha='center',
                    va='center', fontsize=8, fontweight='bold', zorder=3)

    # Subtitle with fiber count
    n = len(fiber)
    status = (f"$N(\\sigma) = {n}$ "
              + ("(infeasible)" if n == 0
                 else "(unique extension)" if n == 1
                 else "(multiple extensions, shown as parallel strands)"))
    ax.set_xlim(-2.5, 2.5)
    ax.set_ylim(-2.5, 2.5)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title(title + "\n" + status, fontsize=10)


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    out = os.path.join(here, 'fig_fiber_c4_example.png')
    fig, axes = plt.subplots(1, 3, figsize=(13.5, 5.0))

    draw_panel(axes[0], (1, 1, 1, 1),
               r"$\sigma = (1,1,1,1)$ (all spokes red)")
    draw_panel(axes[1], (1, 2, 1, 2),
               r"$\sigma = (1,2,1,2)$ (alternating)")
    draw_panel(axes[2], (1, 1, 2, 2),
               r"$\sigma = (1,1,2,2)$")

    # legend
    legend_items = [
        plt.Line2D([], [], color=COLORS[1], lw=3,
                   label='color 1 (red)'),
        plt.Line2D([], [], color=COLORS[2], lw=3,
                   label='color 2 (blue)'),
        plt.Line2D([], [], color=COLORS[3], lw=3,
                   label='color 3 (green)'),
    ]
    fig.legend(handles=legend_items, loc='lower center', ncol=3,
               frameon=False, fontsize=10, bbox_to_anchor=(0.5, -0.02))
    plt.tight_layout()
    plt.savefig(out, dpi=160, bbox_inches='tight')
    plt.close()
    print(f"wrote {out}")

    # Also print the actual enumeration so we know our claims in the .tex
    # match the brute-force counts.
    for sigma in [(1, 1, 1, 1), (1, 2, 1, 2), (1, 1, 2, 2)]:
        f = enumerate_fiber(sigma)
        print(f"sigma = {sigma}: {len(f)} extensions: {f}")


if __name__ == '__main__':
    main()
