"""
Generate G → G' examples covering all 5 iso-classes on n=7,
each as source and as target. Uses balanced planar layout.
"""
import sys
sys.path.insert(0, '/home/claude')

import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
from collections import deque
from itertools import combinations, product as iproduct

from balanced_layout import balanced_planar_layout, _get_all_faces

# ── Helpers ──────────────────────────────────────────────────────────────────

def compute_levels(G, outer_face):
    levels, queue = {}, deque()
    for v in outer_face:
        levels[v] = 0; queue.append(v)
    while queue:
        v = queue.popleft()
        for w in G.neighbors(v):
            if w not in levels:
                levels[w] = levels[v] + 1; queue.append(w)
    return levels

def get_odd_level_cycles(emb, levels):
    odd = []
    for face in _get_all_faces(emb):
        lv = {levels.get(v) for v in face}
        if len(lv) == 1 and None not in lv and len(face) % 2 == 1:
            odd.append(face)
    return odd

def get_flip_candidates(G, emb, cycle):
    cands = []
    n = len(cycle)
    for i in range(n):
        u, v = cycle[i], cycle[(i+1) % n]
        if not G.has_edge(u, v): continue
        f1 = emb.traverse_face(u, v); f2 = emb.traverse_face(v, u)
        if len(f1) != 3 or len(f2) != 3: continue
        w = next(x for x in f1 if x != u and x != v)
        x = next(y for y in f2 if y != u and y != v)
        if w == x or G.has_edge(w, x): continue
        cands.append((u, v, w, x))
    return cands

def is_triangulation(G):
    n = G.number_of_nodes()
    if G.number_of_edges() != 3*n - 6: return False
    if not nx.is_connected(G): return False
    ip, emb = nx.check_planarity(G)
    if not ip: return False
    seen = set()
    for v in emb.nodes():
        for w in emb[v]:
            if (v, w) not in seen:
                f = emb.traverse_face(v, w)
                for i in range(len(f)):
                    seen.add((f[i], f[(i+1) % len(f)]))
                if len(f) != 3: return False
    return True

# ── Enumerate iso-class representatives for n=7 ──────────────────────────────

print("Enumerating triangulations on 7 vertices...")
nodes = list(range(7))
all_edges = list(combinations(nodes, 2))
reps = []
for sub in combinations(all_edges, 15):
    G = nx.Graph(); G.add_nodes_from(nodes); G.add_edges_from(sub)
    if is_triangulation(G):
        if all(not nx.is_isomorphic(G, r) for r in reps):
            reps.append(G)
        if len(reps) == 5: break
print(f"Found {len(reps)} reps.")
for i, G in enumerate(reps):
    deg_seq = sorted((d for _, d in G.degree()), reverse=True)
    print(f"  T{i}: degree {deg_seq}")

def iso_class(G):
    for i, r in enumerate(reps):
        if nx.is_isomorphic(G, r): return i
    return -1

# ── Find one (G_source, outer_face, flip, G_target) for each desired pair ───

desired_pairs = [
    (0, 2),  # T0 -> T2
    (1, 4),  # T1 -> T4
    (2, 0),  # T2 -> T0
    (3, 1),  # T3 -> T1
    (4, 3),  # T4 -> T3
]

def find_example(src_idx, tgt_idx):
    G = reps[src_idx]
    ip, emb = nx.check_planarity(G)
    for outer_face in _get_all_faces(emb):
        levels = compute_levels(G, outer_face)
        odd_cycles = get_odd_level_cycles(emb, levels)
        if not odd_cycles: continue
        candidate_lists = []
        ok = True
        for cycle in odd_cycles:
            cands = get_flip_candidates(G, emb, cycle)
            if not cands: ok = False; break
            candidate_lists.append(cands)
        if not ok: continue
        for choices in iproduct(*candidate_lists):
            Gp = G.copy()
            valid = True
            for u, v, w, x in choices:
                if not Gp.has_edge(u, v) or Gp.has_edge(w, x):
                    valid = False; break
                Gp.remove_edge(u, v); Gp.add_edge(w, x)
            if valid and is_triangulation(Gp) and iso_class(Gp) == tgt_idx:
                return G, outer_face, choices, Gp
    return None

print("\nFinding example flips for each pair:")
examples = []
for src, tgt in desired_pairs:
    ex = find_example(src, tgt)
    if ex is None:
        print(f"  T{src} -> T{tgt}: NO EXAMPLE FOUND")
        continue
    G, outer_face, choices, Gp = ex
    print(f"  T{src} -> T{tgt}: outer={outer_face}, flips={list(choices)}")
    examples.append((src, tgt, G, outer_face, choices, Gp))

# ── Layout and plot ──────────────────────────────────────────────────────────

print("\nComputing balanced layouts (this takes a minute)...")
fig, axes = plt.subplots(len(examples), 2, figsize=(13, 5.5 * len(examples)))
if len(examples) == 1:
    axes = np.array([axes])

palette = ['#1565C0', '#E65100', '#2E7D32']

for row, (src, tgt, G, outer_face, choices, Gp) in enumerate(examples):
    levels = compute_levels(G, outer_face)
    _, emb = nx.check_planarity(G)
    _, emb_p = nx.check_planarity(Gp)
    odd_cycles_G = get_odd_level_cycles(emb, levels)
    # Re-compute levels for Gp using the same outer face (same level-0 vertices)
    levels_p = compute_levels(Gp, outer_face)
    odd_cycles_Gp = get_odd_level_cycles(emb_p, levels_p)

    pos_G  = balanced_planar_layout(G,  outer_face, n_explore=4000,
                                    n_refine=2000, seed=42 + src)

    # G' may have a different planar embedding; the original outer face
    # might no longer be a face. Pick a face of G' that maximally overlaps
    # the original outer face vertices.
    Gp_faces = _get_all_faces(emb_p)
    outer_set_orig = set(outer_face)
    Gp_outer = max(Gp_faces,
                   key=lambda f: (len(set(f) & outer_set_orig), -len(f)))
    pos_Gp = balanced_planar_layout(Gp, Gp_outer, n_explore=4000,
                                    n_refine=2000, seed=42 + tgt)
    print(f"  T{src}→T{tgt}: layouts done")

    flip_edges = set(frozenset([u, v]) for u, v, _, _ in choices)
    new_edges  = set(frozenset([w, x]) for _, _, w, x in choices)

    def classify_G(e):
        if frozenset(e) in flip_edges: return 'flip'
        eu, ev = e
        if levels.get(eu) == levels.get(ev): return 'level'
        return 'normal'

    def classify_Gp(e):
        if frozenset(e) in new_edges: return 'new'
        eu, ev = e
        if levels_p.get(eu) == levels_p.get(ev): return 'level'
        return 'normal'

    outer_set = set(outer_face)

    def draw(ax, Gd, pos, levels_d, odd_cycles_d, classifier, of, title):
        # Outer face
        ax.add_patch(plt.Polygon([pos[nd] for nd in of],
                                 closed=True, facecolor='#E3F2FD',
                                 alpha=0.5, edgecolor='none', zorder=0))
        # Interior odd cycles
        of_set = set(of)
        for cycle in odd_cycles_d:
            if set(cycle) == of_set: continue
            ax.add_patch(plt.Polygon([pos[nd] for nd in cycle],
                                     closed=True, facecolor='#FFF9C4',
                                     alpha=0.6, edgecolor='none', zorder=0))
        # Edges
        cats = {'normal': [], 'level': [], 'flip': [], 'new': []}
        for e in Gd.edges():
            cats[classifier(e)].append(e)
        style = {
            'normal': dict(edge_color='#bdbdbd', width=1.5),
            'level':  dict(edge_color='#2196F3', width=2.5),
            'flip':   dict(edge_color='#e53935', width=3.5, style='dashed'),
            'new':    dict(edge_color='#43a047', width=3.5),
        }
        for cat, eds in cats.items():
            if eds:
                nx.draw_networkx_edges(Gd, pos, edgelist=eds, ax=ax,
                                       **style[cat])
        # Nodes coloured by level
        node_colors = [palette[min(levels_d.get(nd, 0), len(palette)-1)]
                       for nd in Gd.nodes()]
        nx.draw_networkx_nodes(Gd, pos, ax=ax, node_color=node_colors,
                               node_size=450)
        nx.draw_networkx_labels(Gd, pos, ax=ax, font_size=11,
                                font_color='white', font_weight='bold')
        for nd, p in pos.items():
            lv = levels_d.get(nd, '?')
            ax.annotate(f'L{lv}', xy=p + np.array([0.06, 0.06]),
                        fontsize=7, color='#444', zorder=5)
        ax.set_title(title, fontsize=10, fontweight='bold', pad=8)
        ax.set_xlim(-1.5, 1.5); ax.set_ylim(-1.5, 1.5)
        ax.set_aspect('equal'); ax.axis('off')

    flip_str = ", ".join(f"({u},{v})→({w},{x})" for u,v,w,x in choices)
    draw(axes[row, 0], G,  pos_G,  levels,   odd_cycles_G,  classify_G,
         outer_face,
         f"G  (T{src})  |  outer face {outer_face}\nflip: {flip_str}")
    draw(axes[row, 1], Gp, pos_Gp, levels_p, odd_cycles_Gp, classify_Gp,
         Gp_outer,
         f"G'  (T{tgt})  |  outer face {Gp_outer}")

legend_elements = [
    mpatches.Patch(facecolor='#E3F2FD', edgecolor='none', label='Outer face'),
    mpatches.Patch(facecolor='#FFF9C4', edgecolor='none', label='Odd level cycle'),
    plt.Line2D([0],[0], color='#bdbdbd', lw=1.5, label='Normal edge'),
    plt.Line2D([0],[0], color='#2196F3', lw=2.5, label='Level edge'),
    plt.Line2D([0],[0], color='#e53935', lw=3.0, ls='dashed',
               label='Removed edge'),
    plt.Line2D([0],[0], color='#43a047', lw=3.0, label='Added edge'),
]
fig.legend(handles=legend_elements, loc='lower center', ncol=6,
           fontsize=10, bbox_to_anchor=(0.5, -0.005))
plt.suptitle(
    "Level resolutions on n=7  |  every iso-class appears as source and target",
    fontsize=14, fontweight='bold', y=1.0)
plt.tight_layout()
plt.savefig('/mnt/user-data/outputs/n7_all_classes.png',
            dpi=130, bbox_inches='tight')
print("\nSaved.")
