"""Leaf-distance algorithm with the additional rule:

Maintain a `blocked` set of edges; an edge cannot be flipped while
the count x of depth-d faces (d = current max depth) is unchanged.
When x decreases (a depth-d face is removed via a balanced switch),
clear `blocked`. Also clear `blocked` when d itself decreases."""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
from matplotlib.backends.backend_pdf import PdfPages
import math
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 run_with_blocking(faces, outer_set, max_steps=10000):
    """Run with the blocked-edge rule."""
    log = []
    blocked = set()
    prev_x = None
    prev_d = None
    for step in range(max_steps):
        depth = compute_depths(faces, outer_set)
        d_max = max(depth.values())
        n_at_max = sum(1 for d in depth.values() if d == d_max)
        if prev_d is not None and (d_max < prev_d or
                                   (d_max == prev_d and n_at_max < prev_x)):
            blocked = set()  # reset: count of depth-d faces dropped
        prev_d = d_max
        prev_x = n_at_max

        if d_max == 0:
            log.append(f'TERMINATED at step {step}')
            return faces, log

        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:
            Fp = faces[fp_idx]
            u, v = tuple(e)
            w = [vert for vert in F if vert not in (u, v)][0]
            xv = [vert for vert in Fp if vert not in (u, v)][0]
            log.append(f'step {step}: BALANCED on ({u},{v}) (d={d_max},x={n_at_max})')
            blocked.add(frozenset((u, v)))
            faces = apply_switch(faces, (u, v), (w, xv))
            continue

        # Preprocessing: pick lowest-x unblocked neighbour
        x_vals, _ = compute_x(faces, outer_set)
        nb_choices = []
        for e_test in face_edges(F):
            if e_test in outer_set or e_test in blocked:
                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:
            log.append(f'step {step}: STUCK -- all interior edges blocked '
                       f'or none available; max depth={d_max}, x={n_at_max}')
            return faces, log
        nb_choices.sort()
        _, e_chosen, fp_idx = nb_choices[0]
        Fp = faces[fp_idx]
        u, v = tuple(e_chosen)
        w = [vert for vert in F if vert not in (u, v)][0]
        xv = [vert for vert in Fp if vert not in (u, v)][0]
        log.append(f'step {step}: preprocess ({u},{v}) (d={d_max},x={n_at_max}, '
                   f'#blocked={len(blocked)})')
        blocked.add(frozenset((u, v)))
        faces = apply_switch(faces, (u, v), (w, xv))

    log.append(f'hit max_steps')
    return faces, log


# ----- run on the stuck n=40 case -----
seed = 94476710
outer, _, faces = random_triangulation(40, seed=seed)
outer_set = {frozenset(e) for e in outer}

print(f'Running blocked algorithm on n=40, seed={seed}...')
faces_after, log = run_with_blocking(faces, outer_set, max_steps=5000)
print(f'Result: {log[-1]}')
print(f'Total log lines: {len(log)}')
print('First 25 lines:')
for line in log[:25]:
    print(f'  {line}')
print('Last 5 lines:')
for line in log[-5:]:
    print(f'  {line}')

# ----- also test on 9-vertex and 24-vertex -----
print('\n=== 9-vertex ===')
n9 = 9
outer9 = [(i, (i+1) % n9) for i in range(n9)]
outer_set9 = {frozenset(e) for e in outer9}
faces9 = [
    (0, 1, 2), (0, 2, 3), (3, 4, 5), (3, 5, 6),
    (6, 7, 8), (6, 8, 0), (0, 3, 6),
]
_, log9 = run_with_blocking(faces9, outer_set9)
for line in log9[:5]:
    print(f'  {line}')

print('\n=== 24-vertex doubly-lopsided ===')
n24 = 24
outer24 = [(i, (i+1) % n24) for i in range(n24)]
outer_set24 = {frozenset(e) for e in outer24}

def arm(a, b):
    return [
        (a, a+1, a+2), (a, a+2, b), (a+2, a+3, a+4),
        (a+2, a+4, b), (a+4, a+5, a+6), (a+4, a+6, b),
        (a+6, a+7, b),
    ]
faces24 = [(0, 8, 16)]
for (a, b) in [(0, 8), (8, 16), (16, 24)]:
    fs = arm(a, b)
    fs = [tuple(0 if v == 24 else v for v in vt) for vt in fs]
    faces24.extend(fs)
_, log24 = run_with_blocking(faces24, outer_set24)
print(f'  total: {len(log24) - 1} steps')
print(f'  final: {log24[-1]}')