"""
4-color G' using the level structure from G:
  - Even-level subgraph: 2-color with RED/BLUE via BFS
  - Odd-level subgraph:  2-color with YELLOW/GREEN via BFS

This succeeds iff each parity subgraph is bipartite — which is the goal of
the level resolution. If a parity subgraph contains an odd cycle, BFS will
find a conflict and we report which edge violates the 2-coloring.
"""
import networkx as nx
import numpy as np
from collections import deque


def two_color_subgraph(G_sub, color_a, color_b):
    """
    Two-color the (possibly disconnected) subgraph via BFS, alternating
    color_a / color_b by BFS distance parity from a root in each component.

    Returns:
        coloring: dict[node, color]
        bad_edges: list of edges where adjacent vertices got the same colour
                   (empty iff subgraph is bipartite)
    """
    coloring = {}
    bad_edges = []
    for start in G_sub.nodes():
        if start in coloring:
            continue
        coloring[start] = color_a
        queue = deque([start])
        while queue:
            v = queue.popleft()
            for w in G_sub.neighbors(v):
                if w not in coloring:
                    coloring[w] = color_b if coloring[v] == color_a else color_a
                    queue.append(w)
                elif coloring[w] == coloring[v]:
                    e = tuple(sorted([v, w]))
                    if e not in bad_edges:
                        bad_edges.append(e)
    return coloring, bad_edges


def four_color_via_levels(G_prime, levels):
    """
    4-color G' using level labels from G.
    Even-level vertices get RED/BLUE; odd-level get YELLOW/GREEN.

    Returns:
        coloring: dict[node, str]
        bad_edges: dict with keys 'even', 'odd', 'cross' for violations
                   within each parity subgraph and between them
                   (the cross list should always be empty by construction)
    """
    even_nodes = [v for v in G_prime.nodes() if levels[v] % 2 == 0]
    odd_nodes  = [v for v in G_prime.nodes() if levels[v] % 2 == 1]

    even_sub = G_prime.subgraph(even_nodes).copy()
    odd_sub  = G_prime.subgraph(odd_nodes).copy()

    coloring_even, bad_even = two_color_subgraph(even_sub, 'red',    'blue')
    coloring_odd,  bad_odd  = two_color_subgraph(odd_sub,  'yellow', 'green')

    coloring = {**coloring_even, **coloring_odd}

    cross_bad = []
    for u, v in G_prime.edges():
        if coloring[u] == coloring[v]:
            cross_bad.append(tuple(sorted([u, v])))

    return coloring, {
        'even':  bad_even,
        'odd':   bad_odd,
        'cross': [e for e in cross_bad
                  if tuple(sorted(e)) not in bad_even
                  and tuple(sorted(e)) not in bad_odd],
    }


if __name__ == "__main__":
    # Quick demo on the n=7 graph used earlier
    edges = [(0,1),(0,2),(0,3),(0,4),(0,5),(0,6),
             (1,2),(1,3),(1,4),(1,5),(1,6),
             (2,3),(2,4),(3,5),(4,6)]
    G = nx.Graph(); G.add_nodes_from(range(7)); G.add_edges_from(edges)
    Gp = G.copy(); Gp.remove_edge(1, 2); Gp.add_edge(4, 3)

    levels = {0:0, 3:0, 5:0, 1:1, 2:1, 4:1, 6:1}
    coloring, bad = four_color_via_levels(Gp, levels)
    print("Coloring:")
    for v in sorted(coloring):
        print(f"  vertex {v} (level {levels[v]}): {coloring[v]}")
    print(f"\nBad edges in even subgraph:  {bad['even']}")
    print(f"Bad edges in odd subgraph:   {bad['odd']}")
    print(f"Cross-parity bad edges:      {bad['cross']}")
    print(f"\nValid 4-coloring: "
          f"{not any(bad.values())}")
