"""Render the 21-vertex d=2 example and its post-preprocessing state."""
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_balanced_existence import (
    POS, n, OUTER_EDGES, outer_set, compute_depths,
    FACES as FACES0, CHORDS as CHORDS0
)
from d2_preprocessing import apply_switch

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.4, 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=10,
                color=edge_pal.get(d, '#333'), fontweight='bold')
    for (a, b) in OUTER_EDGES + chords:
        color = '#333'; lw = 1.2
        if highlight_edges and ((a, b) in highlight_edges or
                                 (b, a) in highlight_edges):
            color = '#dc2626'; lw = 3.0
        if green_edges and ((a, b) in green_edges or
                            (b, a) in green_edges):
            color = '#16a34a'; lw = 3.0
        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=240, c='#1f2937', edgecolors='black',
                   linewidths=0.8, zorder=2)
        ax.text(x, y, str(i), ha='center', va='center',
                fontsize=8, color='white', fontweight='bold', zorder=3)
    ax.set_aspect('equal'); ax.axis('off')
    ax.set_xlim(-1.25, 1.25); ax.set_ylim(-1.25, 1.25)
    ax.set_title(title, fontsize=11)


depth0, _ = compute_depths(FACES0, outer_set)
faces1, chords1 = apply_switch(FACES0, CHORDS0, (0, 7), (14, 3))
depth1, _ = compute_depths(faces1, outer_set)

fig, axes = plt.subplots(1, 2, figsize=(14, 7))
draw(axes[0], FACES0, CHORDS0, depth0,
     'Before: F=(0,7,14) depth 2, all three depth-1 neighbours lopsided',
     highlight_edges=[(0, 7)])
draw(axes[1], faces1, chords1, depth1,
     'After preprocessing on (0,7): new F=(3,7,14) admits balanced switch on (3,7)',
     green_edges=[(3, 14)], highlight_edges=[(3, 7)])
fig.tight_layout()
out = os.path.join(OUT_DIR, 'fig_d2_preprocessing.png')
fig.savefig(out, dpi=180, bbox_inches='tight')
plt.close(fig)
print(f'wrote {out}')