"""Generate a multi-page PDF showing the L_k state after each edge
switch, on the n=40 stuck example."""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import math
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from matplotlib.backends.backend_pdf import PdfPages
from leaf_distance_algorithm import compute_x
from stress_test_termination import (
    compute_depths, apply_switch, check_balanced, face_edges,
    random_triangulation
)

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


def neighbour_at_edge(faces, F_idx, e):
    for j, fj in enumerate(faces):
        if j != F_idx and e in face_edges(fj):
            return j
    return None


def positions(n):
    return {i: (math.cos(math.radians(90 - i * 360 / n)),
                math.sin(math.radians(90 - i * 360 / n)))
            for i in range(n)}


def draw_state(ax, faces, n, outer_set, switched_edge=None,
               action=None, step=None):
    POS = positions(n)
    depth = compute_depths(faces, outer_set)
    palette = {0: '#86efac', 1: '#fde68a', 2: '#fca5a5'}
    edge_pal = {0: '#16a34a', 1: '#d97706', 2: '#dc2626'}
    for f in faces:
        d = depth.get(faces.index(f), 0)
        # Find depth by iterating
        for i, ff in enumerate(faces):
            if ff == f:
                d = depth[i]
                break
        poly = Polygon([POS[v] for v in f], closed=True,
                       facecolor=palette.get(d, '#ddd'),
                       edgecolor=edge_pal.get(d, '#333'),
                       linewidth=0.7, alpha=0.6, zorder=0)
        ax.add_patch(poly)
    # Draw all edges
    all_edges = set()
    for f in faces:
        all_edges.update(face_edges(f))
    outer_edges = [tuple(e) for e in all_edges if e in outer_set]
    chord_edges = [tuple(e) for e in all_edges if e not in outer_set]
    for (a, b) in outer_edges + chord_edges:
        color = '#333'; lw = 0.5
        if switched_edge and {a, b} == set(switched_edge):
            color = '#dc2626' if action == 'preprocess' else '#16a34a'
            lw = 2.5
        ax.plot([POS[a][0], POS[b][0]], [POS[a][1], POS[b][1]],
                color=color, linewidth=lw, zorder=1)
    # Vertices
    for i, (x, y) in POS.items():
        ax.scatter([x], [y], s=70, c='#1f2937', edgecolors='black',
                   linewidths=0.3, zorder=2)
        ax.text(x, y, str(i), ha='center', va='center',
                fontsize=5, 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)
    d_max = max(depth.values())
    n_d1 = sum(1 for d in depth.values() if d == 1)
    n_d2 = sum(1 for d in depth.values() if d == 2)
    title = f'step {step}'
    if action and switched_edge:
        title += f': {action} on {tuple(sorted(switched_edge))}'
    title += f'  (max={d_max}, #d1={n_d1}, #d2={n_d2})'
    ax.set_title(title, fontsize=9)


def run_and_record(faces, outer_set, n, max_steps=80):
    """Run algorithm, capturing the state and the action at each step.
    Returns list of (faces_at_start_of_step, action, edge)."""
    records = []
    for step in range(max_steps):
        depth = compute_depths(faces, outer_set)
        d_max = max(depth.values())
        if d_max == 0:
            records.append((list(faces), None, None))
            return records, faces

        max_d_faces = [i for i, d in depth.items() if d == d_max]
        F_idx = max_d_faces[0]
        F = faces[F_idx]
        ok, _, fp_idx, e = check_balanced(F_idx, faces, depth, outer_set)
        if ok:
            action = 'BALANCED'
        else:
            x_vals, _ = compute_x(faces, outer_set)
            nb_choices = []
            for e_test in face_edges(F):
                if e_test in outer_set:
                    continue
                nb_idx = neighbour_at_edge(faces, F_idx, e_test)
                if nb_idx is None:
                    continue
                nb_choices.append((x_vals[nb_idx], e_test, nb_idx))
            if not nb_choices:
                records.append((list(faces), 'STUCK', None))
                return records, faces
            nb_choices.sort()
            _, e, fp_idx = nb_choices[0]
            action = 'preprocess'

        Fp = faces[fp_idx]
        u, v = tuple(e)
        w = [vert for vert in F if vert not in (u, v)][0]
        x = [vert for vert in Fp if vert not in (u, v)][0]
        records.append((list(faces), action, (u, v)))
        faces = apply_switch(faces, (u, v), (w, x))
    return records, faces


seed = 94476710
outer, _, faces = random_triangulation(40, seed=seed)
outer_set = {frozenset(e) for e in outer}

print('recording trajectory...')
records, _ = run_and_record(faces, outer_set, 40, max_steps=80)
print(f'recorded {len(records)} steps')

out_pdf = os.path.join(OUT_DIR, 'fig_n40_every_step.pdf')
with PdfPages(out_pdf) as pdf:
    for step, (state_faces, action, edge) in enumerate(records):
        fig, ax = plt.subplots(figsize=(7, 7))
        draw_state(ax, state_faces, 40, outer_set,
                   switched_edge=edge, action=action, step=step)
        pdf.savefig(fig, dpi=120, bbox_inches='tight')
        plt.close(fig)
print(f'wrote {out_pdf}')