"""Visualize the 24-vertex doubly-lopsided d=2 example and the
preprocessing trajectory."""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from d2_recursive_lopsided import (
    POS, n, OUTER_EDGES, outer_set, compute_depths, apply_switch,
    FACES as FACES0, CHORDS as CHORDS0
)

OUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir)


def draw(ax, faces, chords, depth, title,
         highlight_edges=None, green_edges=None):
    palette = {0: '#86efac', 1: '#fde68a', 2: '#fca5a5'}
    edge_pal = {0: '#16a34a', 1: '#d97706', 2: '#dc2626'}
    for i, f in enumerate(faces):
        d = depth[i]
        poly = Polygon([POS[v] for v in f], closed=True,
                       facecolor=palette.get(d, '#ddd'),
                       edgecolor=edge_pal.get(d, '#333'),
                       linewidth=1.2, alpha=0.65, zorder=0)
        ax.add_patch(poly)
        cx = sum(POS[v][0] for v in f) / 3
        cy = sum(POS[v][1] for v in f) / 3
        ax.text(cx, cy, str(d), ha='center', va='center', fontsize=9,
                color=edge_pal.get(d, '#333'), fontweight='bold')
    for (a, b) in OUTER_EDGES + chords:
        color = '#333'; lw = 1.1
        if highlight_edges and ((a, b) in highlight_edges or
                                 (b, a) in highlight_edges):
            color = '#dc2626'; lw = 2.8
        if green_edges and ((a, b) in green_edges or
                            (b, a) in green_edges):
            color = '#16a34a'; lw = 2.8
        ax.plot([POS[a][0], POS[b][0]], [POS[a][1], POS[b][1]],
                color=color, linewidth=lw, zorder=1)
    for i, (x, y) in POS.items():
        ax.scatter([x], [y], s=200, c='#1f2937', edgecolors='black',
                   linewidths=0.6, zorder=2)
        ax.text(x, y, str(i), ha='center', va='center',
                fontsize=7, color='white', fontweight='bold', zorder=3)
    ax.set_aspect('equal'); ax.axis('off')
    ax.set_xlim(-1.2, 1.2); ax.set_ylim(-1.2, 1.2)
    ax.set_title(title, fontsize=10)


depth0 = compute_depths(FACES0)
faces1, chords1 = apply_switch(FACES0, CHORDS0, (0, 8), (16, 2))
depth1 = compute_depths(faces1)
faces2, chords2 = apply_switch(faces1, chords1, (8, 2), (16, 4))
depth2 = compute_depths(faces2)

fig, axes = plt.subplots(1, 3, figsize=(18, 6.5))
draw(axes[0], FACES0, CHORDS0, depth0,
     'Start: F=(0,8,16) depth 2, all arms doubly-lopsided',
     highlight_edges=[(0, 8)])
draw(axes[1], faces1, chords1, depth1,
     'After preprocess 1: F=(2,8,16) still depth 2, still no balanced switch',
     green_edges=[(2, 16)], highlight_edges=[(8, 2)])
draw(axes[2], faces2, chords2, depth2,
     'After preprocess 2: F=(4,8,16) admits balanced switch on (4,8)',
     green_edges=[(4, 16)], highlight_edges=[(4, 8)])
fig.tight_layout()
out = os.path.join(OUT_DIR, 'fig_d2_recursive.png')
fig.savefig(out, dpi=170, bbox_inches='tight')
plt.close(fig)
print(f'wrote {out}')