"""
For each n=7 example (G → G' covering all 5 iso-classes), apply the level
4-coloring to G' using levels from G, and visualize.
"""
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
from four_color import four_color_via_levels

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

# ── Reps ─────────────────────────────────────────────────────────────────────
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

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

# Same pairs as before
desired_pairs = [(0, 2), (1, 4), (2, 0), (3, 1), (4, 3)]

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

examples = [find_example(s, t) for s, t in desired_pairs]
print("Examples found:", sum(1 for e in examples if e is not None))

# ── 4-color and visualize ────────────────────────────────────────────────────

color_hex = {'red':'#e53935','blue':'#1e88e5',
             'yellow':'#fdd835','green':'#43a047'}

fig, axes = plt.subplots(len(examples), 2, figsize=(14, 5.5 * len(examples)))
if len(examples) == 1:
    axes = np.array([axes])

print("\nColoring each G':")
for row, ex in enumerate(examples):
    if ex is None: continue
    G, outer_face, choices, Gp = ex
    src, tgt = desired_pairs[row]

    levels_G = compute_levels(G, outer_face)
    coloring, bad = four_color_via_levels(Gp, levels_G)
    valid = not any(bad.values())
    print(f"  T{src}→T{tgt}: valid 4-coloring = {valid}")
    if not valid:
        print(f"    bad even: {bad['even']}, bad odd: {bad['odd']}")

    # Layout for G and G'
    _, emb_G = nx.check_planarity(G)
    _, emb_p = nx.check_planarity(Gp)
    Gp_faces = _get_all_faces(emb_p)
    outer_set = set(outer_face)
    Gp_outer = max(Gp_faces,
                   key=lambda f: (len(set(f) & outer_set), -len(f)))
    pos_G  = balanced_planar_layout(G,  outer_face, n_explore=4000,
                                    n_refine=2000, seed=42 + src)
    pos_Gp = balanced_planar_layout(Gp, Gp_outer,  n_explore=4000,
                                    n_refine=2000, seed=42 + tgt)

    # ── Draw G on the left (no 4-coloring; just structure) ────────────────
    ax_G = axes[row, 0]
    odd_cycles_G = get_odd_level_cycles(emb_G, levels_G)
    ax_G.add_patch(plt.Polygon([pos_G[v] for v in outer_face], closed=True,
                               facecolor='#E3F2FD', alpha=0.5,
                               edgecolor='none', zorder=0))
    for cycle in odd_cycles_G:
        if set(cycle) == outer_set: continue
        ax_G.add_patch(plt.Polygon([pos_G[v] for v in cycle], closed=True,
                                   facecolor='#FFF9C4', alpha=0.6,
                                   edgecolor='none', zorder=0))
    flip_edges_G = set(frozenset([u, v]) for u, v, _, _ in choices)
    for u, v in G.edges():
        if frozenset([u, v]) in flip_edges_G:
            nx.draw_networkx_edges(G, pos_G, edgelist=[(u, v)], ax=ax_G,
                                   edge_color='#e53935', width=3.5,
                                   style='dashed')
        elif levels_G[u] == levels_G[v]:
            nx.draw_networkx_edges(G, pos_G, edgelist=[(u, v)], ax=ax_G,
                                   edge_color='#2196F3', width=2.5)
        else:
            nx.draw_networkx_edges(G, pos_G, edgelist=[(u, v)], ax=ax_G,
                                   edge_color='#bdbdbd', width=1.5)
    # Nodes by level (un-coloured)
    level_palette = ['#1565C0', '#E65100', '#2E7D32']
    node_lv_cols = [level_palette[min(levels_G[v], len(level_palette)-1)]
                    for v in G.nodes()]
    nx.draw_networkx_nodes(G, pos_G, ax=ax_G, node_color=node_lv_cols,
                           node_size=550)
    nx.draw_networkx_labels(G, pos_G, ax=ax_G, font_size=11,
                            font_color='white', font_weight='bold')
    for v, p in pos_G.items():
        ax_G.annotate(f'L{levels_G[v]}', xy=p + np.array([0.07, 0.07]),
                      fontsize=8, color='#333', zorder=5)
    flip_str = ", ".join(f"({u},{v})→({w},{x})" for u, v, w, x in choices)
    ax_G.set_title(f"G  (T{src})  |  outer face {outer_face}\n"
                   f"flips: {flip_str}",
                   fontsize=9, fontweight='bold', pad=8)
    ax_G.set_xlim(-1.5, 1.5); ax_G.set_ylim(-1.5, 1.5)
    ax_G.set_aspect('equal'); ax_G.axis('off')

    # ── Draw G' on the right (4-coloured) ─────────────────────────────────
    ax = axes[row, 1]
    ax.add_patch(plt.Polygon([pos_Gp[v] for v in Gp_outer], closed=True,
                             facecolor='#ECEFF1', alpha=0.6,
                             edgecolor='none', zorder=0))

    bad_edge_set = set()
    for e in bad['even'] + bad['odd'] + bad['cross']:
        bad_edge_set.add(frozenset(e))

    for u, v in Gp.edges():
        key = frozenset([u, v])
        same_level = levels_G[u] == levels_G[v]
        if key in bad_edge_set:
            nx.draw_networkx_edges(Gp, pos_Gp, edgelist=[(u, v)], ax=ax,
                                   edge_color='#000', width=4.0)
        elif same_level:
            nx.draw_networkx_edges(Gp, pos_Gp, edgelist=[(u, v)], ax=ax,
                                   edge_color='#666', width=2.5)
        else:
            nx.draw_networkx_edges(Gp, pos_Gp, edgelist=[(u, v)], ax=ax,
                                   edge_color='#bdbdbd', width=1.5)

    node_colors_4 = [color_hex[coloring[v]] for v in Gp.nodes()]
    nx.draw_networkx_nodes(Gp, pos_Gp, ax=ax,
                           node_color=node_colors_4,
                           edgecolors='black', linewidths=2,
                           node_size=600)
    nx.draw_networkx_labels(Gp, pos_Gp, ax=ax, font_size=11,
                            font_color='white', font_weight='bold')
    for v, p in pos_Gp.items():
        ax.annotate(f'L{levels_G[v]}', xy=p + np.array([0.07, 0.07]),
                    fontsize=8, color='#333', zorder=5)

    title = (f"G'  (T{tgt})  |  layout outer {Gp_outer}\n"
             f"valid 4-coloring: {valid}")
    if not valid:
        title += f"\nconflicts: even={bad['even']} odd={bad['odd']}"
    ax.set_title(title, fontsize=9, 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')

legend_elements = [
    mpatches.Patch(facecolor='#E3F2FD', edgecolor='none',
                   label="G outer face"),
    mpatches.Patch(facecolor='#FFF9C4', edgecolor='none',
                   label='G odd level cycle'),
    plt.Line2D([0],[0], color='#e53935', lw=3.0, ls='dashed',
               label='Edge flipped out (in G)'),
    plt.Line2D([0],[0], color='#2196F3', lw=2.5, label='Level edge'),
    mpatches.Patch(facecolor=color_hex['red'],    edgecolor='black',
                   label="G' red (even-level)"),
    mpatches.Patch(facecolor=color_hex['blue'],   edgecolor='black',
                   label="G' blue (even-level)"),
    mpatches.Patch(facecolor=color_hex['yellow'], edgecolor='black',
                   label="G' yellow (odd-level)"),
    mpatches.Patch(facecolor=color_hex['green'],  edgecolor='black',
                   label="G' green (odd-level)"),
    plt.Line2D([0],[0], color='#000',    lw=4.0,
               label='Coloring conflict edge'),
]
fig.legend(handles=legend_elements, loc='lower center', ncol=4,
           fontsize=10, bbox_to_anchor=(0.5, -0.002))
plt.suptitle("4-coloring of G' using levels of G  "
             "(even=red/blue, odd=yellow/green)",
             fontsize=13, fontweight='bold', y=1.0)
plt.tight_layout()
plt.savefig('/mnt/user-data/outputs/four_color.png',
            dpi=130, bbox_inches='tight')
print("\nSaved.")
