"""Cut-and-depth-label procedure on a Holton-McKay graph.

Procedure:
  1. Take G' (we use a Holton-McKay non-Hamiltonian 38-vertex cubic
     plane graph, viewed as the dual of a 21-vertex triangulation G).
  2. Find 6 cut edges of G' whose removal disconnects it into two
     non-trivial cubic-minus-boundary pieces G'_0 and G'_1.
  3. In each G'_i:
       a. V_i = vertices of degree 2 (the 6 endpoints in G'_i of cut
          edges).
       b. For each v ∈ V_i, add a new pendant edge to a new vertex.
          Label these pendant edges depth 0.
       c. BFS-propagate: for d = 0, 1, ..., label every edge sharing
          a vertex with a depth-d edge (and not yet labelled) as
          depth d + 1.
       d. Stop when every edge has a depth label.
  4. Render G'_0 and G'_1 with edges coloured by depth.
"""
import os
import sys
from collections import deque, defaultdict
import math

import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np

from sage.all import Graph

HM_FILE = ('/Users/didericis/Code/math-research/papers/'
           'even_level_graph_generators/experiments/nonham38m4.pc')


def parse_planar_code(path):
    with open(path, 'rb') as f:
        data = f.read()
    header = b'>>planar_code<<'
    assert data.startswith(header), data[:20]
    pos = len(header)
    sage_graphs = []
    while pos < len(data):
        n = data[pos]; pos += 1
        edges = set()
        for v in range(n):
            while True:
                w = data[pos]; pos += 1
                if w == 0:
                    break
                edges.add(frozenset((v, w - 1)))
        G = Graph([tuple(e) for e in edges], multiedges=False, loops=False)
        G.is_planar(set_embedding=True)
        sage_graphs.append(G)
    return sage_graphs


def find_six_edge_cut(G, prefer_balanced=True, prefer_matching=True):
    """Find a 6-edge cut: a set of 6 edges whose removal yields two
    non-trivial components.

    For a cubic graph: edges between S and complement
    = 3|S| - 2 * e(S).  We want this to equal 6.

    Search strategy: BFS-grow a vertex set S from each possible start,
    track the boundary edge count.  Collect all sizes that yield a
    6-edge cut.  Return the best one according to:
      prefer_matching: cut is a "matching cut" (each boundary vertex
                       has exactly one cut edge).
      prefer_balanced: |S| closest to n/2.
    """
    n = G.order()
    candidates = []
    for start in G.vertices():
        S = {start}
        boundary = 3
        # BFS-grow; at each step pick the frontier vertex that adds the
        # most internal edges (reduces boundary).
        prev_size = -1
        while True:
            if boundary == 6 and 2 <= len(S) <= n - 2:
                cut = [(u, v) for u in S
                       for v in G.neighbors(u) if v not in S]
                if len(cut) == 6:
                    candidates.append((frozenset(S), cut))
            if len(S) == prev_size:
                break
            prev_size = len(S)
            # frontier
            frontier = []
            for w in G.vertices():
                if w in S: continue
                if any(nb in S for nb in G.neighbors(w)):
                    edges_into_S = sum(1 for nb in G.neighbors(w) if nb in S)
                    frontier.append((edges_into_S, w))
            if not frontier:
                break
            frontier.sort(reverse=True)
            _, w = frontier[0]
            edges_into_S = sum(1 for nb in G.neighbors(w) if nb in S)
            S.add(w)
            boundary = boundary + 3 - 2 * edges_into_S
            if len(S) >= n - 1:
                break
    if not candidates:
        return None, None
    # Pick best: prefer matching cuts (both sides have 6 distinct
    # boundary vertices = 6 cut edges form a matching).
    def score(item):
        S, cut = item
        endpoints_in_S = [u for (u, v) in cut if u in S] + \
                         [v for (u, v) in cut if v in S]
        endpoints_in_Sc = [u for (u, v) in cut if u not in S] + \
                          [v for (u, v) in cut if v not in S]
        is_matching = (len(set(endpoints_in_S)) == 6 and
                       len(set(endpoints_in_Sc)) == 6)
        # Distinctness within each side
        deg2_S = len(set(endpoints_in_S))
        deg2_Sc = len(set(endpoints_in_Sc))
        balance = abs(len(S) - n / 2)
        # Prefer maximum total V coverage, then matching, then balanced
        return (-(deg2_S + deg2_Sc), -int(is_matching), balance)
    candidates.sort(key=score)
    # Dedupe identical S sets
    seen = set()
    unique = []
    for S, cut in candidates:
        if S in seen: continue
        seen.add(S)
        unique.append((S, cut))
    print(f'  Found {len(unique)} distinct 6-edge cuts')
    return unique[0]


def apply_procedure(G, S, cut, base_pos, side_label='0',
                    pendant_start_id=None):
    """Build G'_i from G[S] plus pendant edges at degree-2 vertices.
    Reuses positions from base_pos (a dict v -> (x, y) for G'); pendant
    vertices are placed near their boundary vertex, offset in the
    direction of the cut edge's other endpoint."""
    induced_edges = [(u, v) for (u, v) in G.edges(labels=False)
                     if u in S and v in S]
    H = Graph(induced_edges, multiedges=False, loops=False)
    for v in S:
        H.add_vertex(v)
    induced_deg = {v: H.degree(v) for v in S}
    V_deg2 = sorted([v for v in S if induced_deg[v] == 2])
    high_loss = sorted([v for v in S if induced_deg[v] < 2])
    pendant_edges = []
    if pendant_start_id is None:
        pendant_start_id = max(G.vertices()) + 1
    next_pendant_id = pendant_start_id
    pendant_to_boundary = {}
    # Position pendants: for each boundary vertex v in V_deg2, find
    # its cut-edge neighbor across the cut and offset the pendant in
    # that direction.
    cut_neighbor = {}
    for (u, w) in cut:
        if u in S and w not in S:
            cut_neighbor.setdefault(u, []).append(w)
        if w in S and u not in S:
            cut_neighbor.setdefault(w, []).append(u)
    pos = dict(base_pos)
    for v in V_deg2:
        H.add_edge(v, next_pendant_id)
        pendant_edges.append((min(v, next_pendant_id),
                              max(v, next_pendant_id)))
        pendant_to_boundary[next_pendant_id] = v
        # Position: pendant goes in direction of cut neighbor
        nbrs = cut_neighbor.get(v, [])
        if nbrs:
            nx_, ny_ = base_pos[nbrs[0]]
            vx, vy = base_pos[v]
            dx, dy = nx_ - vx, ny_ - vy
            norm = (dx * dx + dy * dy) ** 0.5 or 1.0
            pos[next_pendant_id] = (vx + 0.35 * dx / norm,
                                    vy + 0.35 * dy / norm)
        else:
            vx, vy = base_pos[v]
            pos[next_pendant_id] = (vx + 0.3, vy)
        next_pendant_id += 1
    # BFS-label edges by depth
    edge_depth = {}
    queue = deque()
    for e in pendant_edges:
        edge_depth[e] = 0
        queue.append((e, 0))
    while queue:
        (u, v), d = queue.popleft()
        # adjacent edges share vertex u or v
        for endpoint in (u, v):
            for nb in H.neighbors(endpoint):
                e_nb = (min(endpoint, nb), max(endpoint, nb))
                if e_nb == (u, v): continue
                if e_nb not in edge_depth:
                    edge_depth[e_nb] = d + 1
                    queue.append((e_nb, d + 1))
    return H, pos, edge_depth, V_deg2, pendant_to_boundary, high_loss


def draw_labeled_graph(ax, H, pos, edge_depth, boundary_vertices,
                        pendant_to_boundary, title):
    H_verts = set(H.vertices())
    max_depth = max(edge_depth.values()) if edge_depth else 0
    cmap = plt.get_cmap('viridis', max_depth + 2)
    legend_handles = []
    for d in range(max_depth + 1):
        color = cmap((d + 0.5) / (max_depth + 1))
        for e, ed in edge_depth.items():
            if ed != d: continue
            u, v = e
            (x1, y1) = pos[u]
            (x2, y2) = pos[v]
            ax.plot([x1, x2], [y1, y2], color=color,
                    linewidth=2.3,
                    linestyle=('--' if d == 0 else '-'),
                    zorder=1)
        legend_handles.append(
            plt.Line2D([], [], color=color, linewidth=2.5,
                       linestyle='--' if d == 0 else '-',
                       label=f'depth {d}')
        )
    for v in H_verts:
        x, y = pos[v]
        if v in pendant_to_boundary:
            ax.plot(x, y, 's', color='#ffaa66', markersize=9, zorder=2,
                    markeredgecolor='#aa5500')
        elif v in boundary_vertices:
            ax.plot(x, y, 'o', color='#ff5555', markersize=11, zorder=3,
                    markeredgecolor='#aa0000', markeredgewidth=1.0)
        else:
            ax.plot(x, y, 'o', color='#666', markersize=7, zorder=2)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title(title, fontsize=11)
    ax.legend(handles=legend_handles, loc='upper left',
              bbox_to_anchor=(1.02, 1.0), fontsize=9, frameon=False)


def compute_nice_layout(G):
    """Compute a nice layout for G.  Tries several methods, picks the
    best by a heuristic that penalizes long/crossing edges."""
    # First try Sage's planar layout from the planar embedding.
    G.is_planar(set_embedding=True)
    pos_planar = G.layout(layout='planar')
    # Spring layout as fallback / alternative.
    try:
        pos_spring = G.layout(layout='spring', iterations=200)
    except Exception:
        pos_spring = None
    # Use networkx's planar layout if available.
    try:
        import networkx as nx
        nxg = nx.Graph()
        for (u, v) in G.edges(labels=False):
            nxg.add_edge(u, v)
        pos_nx_planar = nx.planar_layout(nxg)
        pos_nx_planar = {k: tuple(v) for k, v in pos_nx_planar.items()}
    except Exception:
        pos_nx_planar = None

    def edge_length_variance(pos):
        if pos is None:
            return float('inf')
        lengths = []
        for (u, v) in G.edges(labels=False):
            x1, y1 = pos[u]; x2, y2 = pos[v]
            lengths.append(((x2 - x1)**2 + (y2 - y1)**2) ** 0.5)
        if not lengths:
            return float('inf')
        mean = sum(lengths) / len(lengths)
        var = sum((l - mean)**2 for l in lengths) / len(lengths)
        return var / (mean * mean + 1e-9)  # coefficient of variation^2

    options = [
        ('sage-planar', pos_planar),
        ('sage-spring', pos_spring),
        ('nx-planar', pos_nx_planar),
    ]
    scored = [(edge_length_variance(p), name, p) for name, p in options if p]
    scored.sort()
    name, _, pos = scored[0][1], scored[0][0], scored[0][2]
    print(f'Selected layout: {name} '
          f'(edge-length CV^2 = {scored[0][0]:.4f})')
    return pos


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    out = os.path.normpath(os.path.join(here, '..', 'notes',
                                        'fig_cut_depth_label.png'))

    gs = parse_planar_code(HM_FILE)
    print(f'Loaded {len(gs)} Holton-McKay graphs')

    G = gs[0]
    print(f'Graph 0: {G.order()} vertices, {G.size()} edges, '
          f'cubic={all(d == 3 for d in G.degree())}, '
          f'planar={G.is_planar()}')

    S, cut = find_six_edge_cut(G)
    if S is None:
        print('No 6-edge cut found by greedy search.')
        return
    print(f'Found 6-edge cut: |S| = {len(S)}, |cut| = {len(cut)}')
    print(f'Cut edges: {cut}')

    # Compute a single nice layout for G' once
    base_pos = compute_nice_layout(G)

    S0 = frozenset(S)
    S1 = frozenset(G.vertices()) - S0

    H0, pos0, ed0, bv0, pmap0, hl0 = apply_procedure(
        G, S0, cut, base_pos, side_label='0',
        pendant_start_id=max(G.vertices()) + 1)
    H1, pos1, ed1, bv1, pmap1, hl1 = apply_procedure(
        G, S1, cut, base_pos, side_label='1',
        pendant_start_id=max(G.vertices()) + 1 + len(bv0))

    print(f"G'_0: {H0.order()} vertices, {H0.size()} edges, "
          f"|V|={len(bv0)}, max depth={max(ed0.values())}")
    print(f"G'_1: {H1.order()} vertices, {H1.size()} edges, "
          f"|V|={len(bv1)}, max depth={max(ed1.values())}")

    # Three-panel figure: G' (full), G'_0, G'_1, all with same positions
    # Choose figure size based on aspect ratio of the layout.
    xs = [p[0] for p in base_pos.values()]
    ys = [p[1] for p in base_pos.values()]
    xrange = max(xs) - min(xs)
    yrange = max(ys) - min(ys)
    aspect = xrange / yrange if yrange > 0 else 1.0
    panel_w = 6.5
    panel_h = panel_w / aspect
    fig, axes = plt.subplots(3, 1, figsize=(panel_w * 1.4, 3 * panel_h * 1.05))

    # Panel 0: G' (full) with cut edges highlighted
    ax = axes[0]
    cut_set = set(frozenset(e) for e in cut)
    for (u, v) in G.edges(labels=False):
        x1, y1 = base_pos[u]; x2, y2 = base_pos[v]
        if frozenset((u, v)) in cut_set:
            ax.plot([x1, x2], [y1, y2], color='#d62728',
                    linewidth=2.5, linestyle='--', zorder=2)
        else:
            ax.plot([x1, x2], [y1, y2], color='#bbbbbb',
                    linewidth=1.2, zorder=1)
    for v, (x, y) in base_pos.items():
        if v in S0:
            ax.plot(x, y, 'o', color='#4c72b0', markersize=8, zorder=3)
        else:
            ax.plot(x, y, 'o', color='#dd8452', markersize=8, zorder=3)
    ax.set_title(f"$G'$ = Holton-McKay #0 with 6-edge cut highlighted\n"
                 f"Blue = $S$ ($|S|={len(S0)}$);  "
                 f"orange = $V \\setminus S$ ($|V\\setminus S|={len(S1)}$);  "
                 f"red dashed = cut",
                 fontsize=10)
    ax.set_aspect('equal'); ax.axis('off')

    # Panels 1 and 2: G'_0 and G'_1 in the same layout
    draw_labeled_graph(axes[1], H0, pos0, ed0, bv0, pmap0,
                       f"$G'_0$ (|S| = {len(S0)}, |V| = {len(bv0)}, "
                       f"max depth = {max(ed0.values())})")
    draw_labeled_graph(axes[2], H1, pos1, ed1, bv1, pmap1,
                       f"$G'_1$ (|S| = {len(S1)}, |V| = {len(bv1)}, "
                       f"max depth = {max(ed1.values())})")

    # Force same axis limits for the three panels (using base_pos)
    xs = [p[0] for p in base_pos.values()]
    ys = [p[1] for p in base_pos.values()]
    xmin, xmax = min(xs) - 0.5, max(xs) + 0.5
    ymin, ymax = min(ys) - 0.5, max(ys) + 0.5
    for ax in axes:
        ax.set_xlim(xmin, xmax)
        ax.set_ylim(ymin, ymax)

    fig.suptitle("Cut-and-depth-label procedure on Holton-McKay graph #0\n"
                 "Same vertex positions used across all three panels.",
                 fontsize=12, y=1.00)
    plt.tight_layout()
    plt.savefig(out, dpi=160, bbox_inches='tight')
    plt.close()
    print(f'Wrote {out}')


if __name__ == '__main__':
    main()
