"""Plot the depth distribution over time for the random n=40 trajectory
under the leaf-distance algorithm. Shows whether progress is being made
or whether the algorithm is just grinding."""
import os
import sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import matplotlib.pyplot as plt
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_full_tracking(faces, outer_set, max_steps=3000):
    """Run algorithm; record at each step: (#faces at each depth), and
    whether the step was BALANCED or x-preprocess."""
    history = []
    for step in range(max_steps):
        depth = compute_depths(faces, outer_set)
        d_max = max(depth.values())
        # Record current state
        d_dist = {}
        for d in depth.values():
            d_dist[d] = d_dist.get(d, 0) + 1
        history.append({'step': step, 'd_max': d_max, 'd_dist': dict(d_dist)})
        if d_max == 0:
            return history, 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:
            history[-1]['action'] = 'BALANCED'
            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]
            faces = apply_switch(faces, (u, v), (w, x))
            continue
        history[-1]['action'] = 'preprocess'
        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:
            return history, faces
        nb_choices.sort()
        _, e_chosen, nb_idx = nb_choices[0]
        Fnb = faces[nb_idx]
        u, v = tuple(e_chosen)
        w = [vert for vert in F if vert not in (u, v)][0]
        xv = [vert for vert in Fnb if vert not in (u, v)][0]
        faces = apply_switch(faces, (u, v), (w, xv))
    return history, faces


seed = 94476710
outer, _, faces = random_triangulation(40, seed=seed)
outer_set = {frozenset(e) for e in outer}
print(f'Running algorithm on n=40, seed={seed}...')
history, final_faces = run_with_full_tracking(faces, outer_set, max_steps=3000)
print(f'finished at step {len(history) - 1}, '
      f'final max depth = {history[-1]["d_max"]}')

# Plot count of depth-d faces for d = 1 and 2 over time, and mark
# balanced vs preprocessing steps
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 7), sharex=True)

steps = [h['step'] for h in history]
n_d1 = [h['d_dist'].get(1, 0) for h in history]
n_d2 = [h['d_dist'].get(2, 0) for h in history]
balanced_steps = [h['step'] for h in history
                  if h.get('action') == 'BALANCED']
preprocess_steps = [h['step'] for h in history
                    if h.get('action') == 'preprocess']

ax1.plot(steps, n_d1, color='#d97706', linewidth=1.2, label='depth-1 faces')
ax1.plot(steps, n_d2, color='#dc2626', linewidth=1.2, label='depth-2 faces')
ax1.set_ylabel('count of faces at depth')
ax1.legend(loc='upper right')
ax1.set_title(f'n=40 trajectory under leaf-distance algorithm (seed={seed}, '
              f'{len(history) - 1} steps)')
ax1.grid(alpha=0.3)

# Add markers showing balanced vs preprocess at bottom
ax2.scatter(balanced_steps, [1] * len(balanced_steps),
            color='#16a34a', s=8, label=f'BALANCED ({len(balanced_steps)} steps)')
ax2.scatter(preprocess_steps, [0] * len(preprocess_steps),
            color='#3b82f6', s=8, label=f'preprocess ({len(preprocess_steps)} steps)')
ax2.set_yticks([0, 1])
ax2.set_yticklabels(['preprocess', 'BALANCED'])
ax2.set_xlabel('step')
ax2.legend(loc='upper right')
ax2.grid(alpha=0.3)

fig.tight_layout()
out = os.path.join(OUT_DIR, 'fig_n40_trajectory.png')
fig.savefig(out, dpi=150, bbox_inches='tight')
plt.close(fig)
print(f'wrote {out}')

# Summary
print(f'\n#BALANCED switches: {len(balanced_steps)}')
print(f'#preprocess switches: {len(preprocess_steps)}')
print(f'Max depth-1 count seen: {max(n_d1)}')
print(f'Max depth-2 count seen: {max(n_d2)}')