"""Figures for the 'outer triangle absorption' note.

Three panels:
  A. A closed PDS chain (concentric levels around source, ending at
     outer triangle).
  B. The outer-triangle constraint in G': dual vertex of degree 3
     forces the 3 incident edges to have distinct colours.
  C. The state-size trajectory across the chain showing the
     'absorption' to size 6.
"""
import math
import os
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon, FancyArrowPatch
import matplotlib.patches as patches

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


def ring(n, r, phase=math.pi / 2):
    return [(r * math.cos(phase + 2 * math.pi * i / n),
             r * math.sin(phase + 2 * math.pi * i / n))
            for i in range(n)]


# -----------------------------------------------------------------
# Panel A: the closed PDS chain (5, 6, 5, 3) — concentric layers
# -----------------------------------------------------------------
def panel_A(ax):
    radii = [0.0, 0.9, 1.7, 2.5, 3.4]
    sizes = [1, 5, 6, 5, 3]   # L_0, L_1, L_2, L_3, L_4
    rings = [None] + [ring(sizes[i], radii[i]) for i in range(1, 5)]

    # Source vertex
    ax.plot(0, 0, 'o', color='#333', markersize=10, zorder=5)
    ax.annotate("$L_0$ (source)", (0, 0), xytext=(0.3, 0.15),
                fontsize=9, color='#333')

    colors = ['#1f77b4', '#2ca02c', '#d62728', '#9467bd']
    labels = [r'$L_1$ ($|L_1|=5$)', r'$L_2$ ($|L_2|=6$)',
              r'$L_3$ ($|L_3|=5$)', r'$L_4$ (outer triangle, $|L_4|=3$)']

    # Draw the level cycles
    for i, (pts, color, label) in enumerate(zip(rings[1:], colors, labels)):
        for j in range(len(pts)):
            x1, y1 = pts[j]
            x2, y2 = pts[(j + 1) % len(pts)]
            ax.plot([x1, x2], [y1, y2], color=color, linewidth=2.5, zorder=3)
        for (x, y) in pts:
            ax.plot(x, y, 'o', color=color, markersize=7, zorder=4)
        ax.annotate(label,
                    xy=(0, radii[i + 1] + 0.05),
                    fontsize=9, color=color,
                    ha='center', va='bottom')

    # Draw the annular triangulation between consecutive levels
    # (simple even spacing — not geometrically faithful to a real PDS,
    # just illustrative).
    for i in range(4):
        inner_pts = rings[i] if rings[i] else [(0, 0)] * sizes[i + 1]
        # For L_0 the apex is at the origin — replicate per outer point
        outer_pts = rings[i + 1]
        m_in = len(inner_pts)
        m_out = len(outer_pts)
        if m_in == 1:
            # All outer points connected to apex
            for (x, y) in outer_pts:
                ax.plot([0, x], [0, y], color='#aaa', linewidth=0.6, zorder=1)
        else:
            # Zig-zag triangulation: each outer to its "nearest" inner
            # by angle.  This is rough; just for illustration.
            angles_in = [(math.atan2(y, x), j) for j, (x, y) in enumerate(inner_pts)]
            for (xo, yo) in outer_pts:
                ao = math.atan2(yo, xo)
                # Find 2 nearest inner points by angle
                diffs = sorted(((abs((a - ao + math.pi) % (2 * math.pi) - math.pi), j)
                                for a, j in angles_in))
                for _, j in diffs[:2]:
                    xi, yi = inner_pts[j]
                    ax.plot([xo, xi], [yo, yi], color='#aaa', linewidth=0.6, zorder=1)

    # Annotate tires (annular regions between consecutive levels)
    tire_labels = [
        ("$T_1$\n(apex)", 0.45, 0),
        ("$T_2$", 1.3, 0),
        ("$T_3$", 2.1, 0),
        ("$T_4$ (outer)", 2.95, 0),
    ]
    for lbl, x, y in tire_labels:
        ax.text(x, y, lbl, fontsize=8, color='#555',
                ha='center', va='center',
                bbox=dict(boxstyle='round,pad=0.18', facecolor='white',
                          edgecolor='none', alpha=0.7))

    ax.set_xlim(-3.8, 3.8)
    ax.set_ylim(-3.8, 3.8)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title(r"(A) Closed PDS chain: $L_0$ (source) $\to L_4$ (outer triangle)",
                 fontsize=10)


# -----------------------------------------------------------------
# Panel B: the outer-triangle constraint in G'
# -----------------------------------------------------------------
def panel_B(ax):
    # Outer triangle vertices in G
    pts = [(math.cos(math.pi/2 + 2*math.pi*i/3),
            math.sin(math.pi/2 + 2*math.pi*i/3)) for i in range(3)]
    # Outer face's dual vertex (outside the triangle)
    outer_dual = (0, 1.9)
    # Each outer triangle edge has a midpoint that connects to outer_dual
    # in G'.

    # Draw the outer triangle (in G)
    for i in range(3):
        x1, y1 = pts[i]
        x2, y2 = pts[(i + 1) % 3]
        ax.plot([x1, x2], [y1, y2], color='#9467bd', linewidth=2.5, zorder=3)
    for (x, y) in pts:
        ax.plot(x, y, 'o', color='#9467bd', markersize=9, zorder=4)

    # Edge midpoints (these are 'duals' of the edges, in some sense)
    mids = []
    for i in range(3):
        x1, y1 = pts[i]
        x2, y2 = pts[(i + 1) % 3]
        mids.append(((x1 + x2) / 2, (y1 + y2) / 2))

    # Inside the triangle: the annular face dual of T_n on each edge
    # (a single G'-vertex per outer edge, one of T_n's annular face duals).
    inside_pts = [(0.5 * mx, 0.5 * my - 0.05) for (mx, my) in mids]

    # Outer face dual vertex (above the triangle, on G' side)
    ax.plot(outer_dual[0], outer_dual[1], 's', color='#d62728',
            markersize=18, zorder=4)
    ax.annotate("outer face\ndual vertex", outer_dual,
                xytext=(outer_dual[0] - 1.7, outer_dual[1] + 0.1),
                fontsize=9, color='#d62728',
                arrowprops=dict(arrowstyle='->', color='#d62728', lw=1.0))

    # Three G'-edges from outer face dual to inside (each crossing an outer-triangle edge)
    g_prime_colors = ['#d62728', '#1f77b4', '#2ca02c']  # R, B, G — the three colors
    for k, (mx, my) in enumerate(mids):
        # Curve from outer_dual through this midpoint to inside_pts[k]
        ix, iy = inside_pts[k]
        ax.plot([outer_dual[0], mx], [outer_dual[1], my],
                color=g_prime_colors[k], linewidth=2.2, zorder=2,
                linestyle='--')
        ax.plot([mx, ix], [my, iy],
                color=g_prime_colors[k], linewidth=2.2, zorder=2,
                linestyle='--')
        ax.plot(ix, iy, 'o', color='#888', markersize=8, zorder=3)

    # Annotate the inside vertices (T_n annular face duals)
    ax.text(0, -0.55, r"$T_n$'s annular face duals", fontsize=8, color='#555',
            ha='center')

    # Annotate the constraint
    ax.text(0, 2.6, "outer face dual has 3 incident\n"
                   "$G'$-edges → all 3 colours distinct",
            fontsize=9, ha='center', va='bottom',
            color='#d62728')

    ax.set_xlim(-2.3, 2.3)
    ax.set_ylim(-1.3, 3.1)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title(r"(B) Outer-triangle dual constraint forces $\sigma|_{L_n}$ to be a permutation",
                 fontsize=10)


# -----------------------------------------------------------------
# Panel C: state-size trajectory
# -----------------------------------------------------------------
def panel_C(ax):
    # Example trajectory from the experiment (7-tire chain)
    chain_labels = [r"$L_1$", r"$L_2$", r"$L_3$", r"$L_4$",
                    r"$L_5$", r"$L_6$", r"$L_7$"]
    traj = [63, 783, 9993, 14643, 1641, 60, 6]
    xs = list(range(len(traj)))
    ax.plot(xs, traj, 'o-', color='#1f3f70', linewidth=2, markersize=8)
    # Annotate each value
    for x, v in zip(xs, traj):
        ax.annotate(str(v), (x, v),
                    textcoords="offset points", xytext=(0, 9),
                    ha='center', fontsize=8, color='#1f3f70')
    ax.set_yscale('log')
    ax.set_xticks(xs)
    ax.set_xticklabels(chain_labels)
    ax.set_xlabel("shared cycle (level)")
    ax.set_ylabel("state size (log scale)")
    ax.grid(True, alpha=0.3, axis='y')
    ax.set_title(r"(C) State trajectory: chain (6,1)$|$(8,6)$|$(10,8)$|$(10,10)$|$(8,10)$|$(5,8)$|$(3,5)",
                 fontsize=10)
    # Annotate the "absorption" at the last step
    ax.annotate("absorption to\n6 permutations\nof $\\{1,2,3\\}$",
                xy=(6, 6), xytext=(4.5, 30),
                fontsize=9, color='#d62728',
                arrowprops=dict(arrowstyle='->', color='#d62728', lw=1.2))


def main():
    fig = plt.figure(figsize=(15.5, 5.5))
    gs = fig.add_gridspec(1, 3, width_ratios=[1.0, 1.0, 1.2])
    ax_A = fig.add_subplot(gs[0, 0])
    ax_B = fig.add_subplot(gs[0, 1])
    ax_C = fig.add_subplot(gs[0, 2])
    panel_A(ax_A)
    panel_B(ax_B)
    panel_C(ax_C)
    plt.tight_layout()
    out = os.path.join(HERE, 'fig_outer_triangle_absorption.png')
    plt.savefig(out, dpi=160, bbox_inches='tight')
    plt.close()
    print(f"wrote {out}")


if __name__ == '__main__':
    main()
