import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
from collections import deque

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_all_faces(emb):
    seen = set()
    faces = []
    for v in emb.nodes():
        for w in emb[v]:
            if (v, w) not in seen:
                face = emb.traverse_face(v, w)
                for i in range(len(face)):
                    seen.add((face[i], face[(i+1) % len(face)]))
                faces.append(tuple(face))
    return faces

G = nx.octahedral_graph()
_, emb = nx.check_planarity(G)
all_faces = get_all_faces(emb)

pos = nx.spring_layout(G, seed=7, iterations=200)

fig, axes = plt.subplots(2, 4, figsize=(16, 8))
axes = axes.flatten()

for idx, outer_face in enumerate(all_faces):
    ax = axes[idx]
    outer_face_set = set(outer_face)
    levels = compute_levels(G, outer_face)

    odd_cycle = None
    for face in all_faces:
        if set(face) == outer_face_set:
            continue
        face_levels = set(levels.get(v) for v in face)
        if len(face_levels) == 1 and None not in face_levels and len(face) % 2 == 1:
            odd_cycle = face
            break

    odd_edges = set()
    if odd_cycle:
        n = len(odd_cycle)
        for i in range(n):
            odd_edges.add(frozenset([odd_cycle[i], odd_cycle[(i+1) % n]]))

    outer_edges = set()
    n_of = len(outer_face)
    for i in range(n_of):
        outer_edges.add(frozenset([outer_face[i], outer_face[(i+1) % n_of]]))

    red_edges  = [e for e in G.edges() if frozenset(e) in odd_edges]
    blue_edges = [e for e in G.edges() if frozenset(e) in outer_edges]
    grey_edges = [e for e in G.edges()
                  if frozenset(e) not in odd_edges and frozenset(e) not in outer_edges]

    node_colors = ['#4a90d9' if levels[v] == 0 else '#f0a500' for v in G.nodes()]

    if odd_cycle:
        tri = np.array([pos[v] for v in odd_cycle])
        ax.add_patch(plt.Polygon(tri, closed=True, facecolor='#e03030', alpha=0.15, edgecolor='none'))

    out_arr = np.array([pos[v] for v in outer_face])
    ax.add_patch(plt.Polygon(out_arr, closed=True, facecolor='#4a90d9', alpha=0.15, edgecolor='none'))

    nx.draw_networkx_edges(G, pos, edgelist=grey_edges, ax=ax, edge_color='#aaaaaa', width=1.5)
    nx.draw_networkx_edges(G, pos, edgelist=blue_edges, ax=ax, edge_color='#4a90d9', width=2.5, style='dashed')
    nx.draw_networkx_edges(G, pos, edgelist=red_edges,  ax=ax, edge_color='#e03030', width=3.5)
    nx.draw_networkx_nodes(G, pos, ax=ax, node_color=node_colors, node_size=350)
    nx.draw_networkx_labels(G, pos, ax=ax, font_size=9, font_color='white', font_weight='bold')

    ax.set_title(f"outer={outer_face}\nodd cycle={odd_cycle}", fontsize=8, pad=4)
    ax.axis('off')

legend_elements = [
    mpatches.Patch(facecolor='#4a90d9', alpha=0.5, label='Level 0 (outer face)'),
    mpatches.Patch(facecolor='#f0a500', alpha=0.8, label='Level 1 node'),
    mpatches.Patch(facecolor='#e03030', alpha=0.5, label='Odd level cycle (level 1)'),
    plt.Line2D([0],[0], color='#4a90d9', linewidth=2.5, linestyle='dashed', label='Outer face edges'),
    plt.Line2D([0],[0], color='#e03030', linewidth=3.5, label='Odd cycle edges'),
    plt.Line2D([0],[0], color='#aaaaaa', linewidth=1.5, label='Other edges'),
]
fig.legend(handles=legend_elements, loc='lower center', ncol=3,
           fontsize=9, bbox_to_anchor=(0.5, -0.01))
fig.suptitle("Octahedron: odd level cycle (red) for each outer face",
             fontsize=13, fontweight='bold')
plt.tight_layout()
plt.savefig('/mnt/user-data/outputs/octahedron_level_cycles.png', dpi=150, bbox_inches='tight')
print("Saved.")
