"""Draw an actual maximal planar graph exhibiting an (R2) violation.

Uses the first such example found by experiments/check_r2.py:
  n=10, source=4, depth=1, n_bdry=3.
"""
import os
import sys
from collections import defaultdict

from sage.all import Graph
from sage.graphs.graph_generators import graphs

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import matplotlib.pyplot as plt
import matplotlib.patches as patches

# Edges from the first n=10 R2-violating triangulation found.
EDGES = [(1, 2), (1, 3), (1, 4), (1, 5), (1, 6), (1, 7), (1, 8), (1, 9),
         (2, 3), (2, 8), (2, 9), (3, 4), (3, 5), (3, 6), (3, 7), (3, 8),
         (3, 10), (4, 5), (5, 6), (6, 7), (6, 10), (7, 8), (7, 10), (8, 9)]
SOURCE = 4
DEPTH = 1


def main():
    G = Graph(EDGES)
    G.is_planar(set_embedding=True)
    emb = G.get_embedding()
    dist = G.shortest_path_lengths(SOURCE)
    print(f"BFS distances from {SOURCE}: {dict(dist)}")

    faces = G.faces(embedding=emb)
    face_depths = []
    for face in faces:
        verts = set()
        for u, v in face:
            verts.add(u); verts.add(v)
        d = min(dist[v] for v in verts)
        face_depths.append((face, d, verts))

    # Find depth-DEPTH faces and their connected components
    target_faces = [(i, fd) for i, fd in enumerate(face_depths)
                    if fd[1] == DEPTH]
    print(f"Found {len(target_faces)} faces at depth {DEPTH}")

    edge_to_idxs = defaultdict(list)
    for k, (i, fd) in enumerate(target_faces):
        for u, v in fd[0]:
            edge_to_idxs[frozenset({u, v})].append(k)
    n = len(target_faces)
    adj = [[] for _ in range(n)]
    for idxs in edge_to_idxs.values():
        if len(idxs) >= 2:
            for a in idxs:
                for b in idxs:
                    if a != b:
                        adj[a].append(b)
    visited = [False] * n
    comps = []
    for s in range(n):
        if visited[s]:
            continue
        stack = [s]; comp = []
        while stack:
            x = stack.pop()
            if visited[x]: continue
            visited[x] = True; comp.append(target_faces[x])
            for y in adj[x]:
                if not visited[y]: stack.append(y)
        comps.append(comp)
    print(f"Connected components at depth {DEPTH}: {len(comps)} sizes "
          f"{[len(c) for c in comps]}")
    # Pick the largest one
    comp = max(comps, key=lambda c: len(c))

    # Compute boundary edges of comp
    edge_count = defaultdict(int)
    comp_face_edges = []
    for _, fd in comp:
        fe = {frozenset({u, v}) for u, v in fd[0]}
        comp_face_edges.append(fe)
        for e in fe:
            edge_count[e] += 1
    bdry_edges = {e for e, c in edge_count.items() if c == 1}

    # Boundary connected components
    bdry_graph = defaultdict(set)
    for e in bdry_edges:
        u, v = list(e)
        bdry_graph[u].add(v)
        bdry_graph[v].add(u)
    seen = set()
    bdry_cycles = []
    for v in bdry_graph:
        if v in seen: continue
        comp_v = set(); stack = [v]
        while stack:
            x = stack.pop()
            if x in seen: continue
            seen.add(x); comp_v.add(x)
            for y in bdry_graph[x]:
                if y not in seen: stack.append(y)
        bdry_cycles.append(comp_v)
    print(f"Boundary components: {len(bdry_cycles)}, vertex sets: {bdry_cycles}")

    # Compute layout using sage's planar layout
    pos = G.layout(layout='planar')

    # Plot
    fig, ax = plt.subplots(figsize=(9, 9))

    # Fill in depth-1 component faces
    for _, fd in comp:
        face = fd[0]
        cycle_verts = [edge[0] for edge in face]
        poly = plt.Polygon([pos[v] for v in cycle_verts],
                           facecolor='#fff3cd', edgecolor='none',
                           alpha=0.55, zorder=0)
        ax.add_patch(poly)

    # Draw all edges of G
    for u, v in G.edges(labels=False):
        x1, y1 = pos[u]; x2, y2 = pos[v]
        ax.plot([x1, x2], [y1, y2], color='lightgray', linewidth=0.8,
                zorder=1)

    # Highlight boundary edges of comp (in 3 different colors)
    colors = ['#d62728', '#1f77b4', '#2ca02c']
    cyclic_label = ['boundary 1', 'boundary 2', 'boundary 3']
    for i, bc_verts in enumerate(bdry_cycles):
        for e in bdry_edges:
            if e <= bc_verts:
                u, v = list(e)
                x1, y1 = pos[u]; x2, y2 = pos[v]
                ax.plot([x1, x2], [y1, y2], color=colors[i % 3],
                        linewidth=3.0, zorder=2)

    # Draw vertices, color-coded by BFS level
    level_colors = {0: 'black', 1: '#1f77b4', 2: '#d62728'}
    for v in G.vertices():
        x, y = pos[v]
        c = level_colors.get(dist[v], 'gray')
        ax.plot(x, y, 'o', color=c, markersize=22, zorder=3)
        ax.annotate(f"{v}", (x, y), color='white', ha='center', va='center',
                    fontsize=11, fontweight='bold', zorder=4)
        # Add depth annotation
        ax.annotate(f"$\\ell={dist[v]}$", (x, y),
                    xytext=(x + 0.05, y + 0.05), fontsize=9, zorder=4,
                    color='black',
                    bbox=dict(boxstyle='round,pad=0.15',
                              facecolor='white', edgecolor='none',
                              alpha=0.75))

    # Legend
    legend_items = [
        plt.Line2D([], [], marker='o', color='w', markerfacecolor='black',
                   markersize=12, label='source $v_0$ ($\\ell=0$)'),
        plt.Line2D([], [], marker='o', color='w', markerfacecolor='#1f77b4',
                   markersize=12, label='$\\ell=1$'),
        plt.Line2D([], [], marker='o', color='w', markerfacecolor='#d62728',
                   markersize=12, label='$\\ell=2$'),
        plt.Line2D([], [], color='#fff3cd', linewidth=12,
                   label=f'$R_{{C^\\prime}}$ (depth-{DEPTH} component, '
                         f'{len(comp)} faces)'),
        plt.Line2D([], [], color='#d62728', linewidth=3,
                   label='boundary cycle 1'),
        plt.Line2D([], [], color='#1f77b4', linewidth=3,
                   label='boundary cycle 2'),
        plt.Line2D([], [], color='#2ca02c', linewidth=3,
                   label='boundary cycle 3'),
    ]
    ax.legend(handles=legend_items, loc='upper right', fontsize=9,
              framealpha=0.95)

    ax.set_title(f"(R2) violation: depth-{DEPTH} component with "
                 f"3 boundary cycles\n(maximal planar graph, $n=10$, "
                 f"source $v_0={SOURCE}$)",
                 fontsize=12)
    ax.set_aspect('equal')
    ax.axis('off')

    out = os.path.join(HERE, 'r2_violation_real_example.png')
    plt.savefig(out, dpi=160, bbox_inches='tight')
    plt.close()
    print(f"wrote {out}")


if __name__ == '__main__':
    main()
