from typing import Set, FrozenSet, Dict, Iterable, Tuple, List
from collections import defaultdict
from itertools import count

Node = int
Edge = FrozenSet[Node]
Color = int

edge = lambda u, v: frozenset((u, v))

def misra_gries(edges: Set[Edge]) -> Dict[Edge, Color]:
    r"""Return an edge-coloring of a simple connected graph."""
    colors = {}
    adjacencies = defaultdict(set)
    for u, v in edges:
        adjacencies[u].add(v)
        adjacencies[v].add(u)

    def adjacent_colors(u: Node) -> Iterable[Tuple[Node, Color]]:
        # yield all colored nodes adjacent to u, along with their color
        for n in adjacencies[u]:
            color = colors.get(edge(u, n), None)
            if color is not None:
                yield n, color

    def is_free(c: Color, u: Node) -> bool:
        # return whether node u touches an edge of color c
        return all(color != c for _, color in adjacent_colors(u))

    def free_color(u: Node) -> Color:
        # return the first free color of node u
        used_colors = {color for _, color in adjacent_colors(u)}
        for color in count():
            if color not in used_colors:
                return color
                
    def maximal_fan(u: Node, v: Node) -> List[Node]:
        # return a maximal on {u, v} at u
        fan = [v]
        remaining = set(adjacent_colors(u))
        while remaining:
            for n, c in remaining:
                if is_free(c, fan[-1]):
                    fan.append(n)
                    remaining.discard((n,c))
                    break
            else:
                break
        return fan
    
    def cdu_path(c: Color, d: Color, u: Node) -> Set[Edge]:
        # return a cd_u path
        path = set()
        endpoints = {u}
        while endpoints:
            e = endpoints.pop()
            for n, color in adjacent_colors(e):
                if edge(e, n) not in path and color in (c,d):
                    endpoints.add(n)
                    path.add(edge(e, n))
        return path 
    
    remaining = set(edges)
    while remaining:
        u, v = remaining.pop()
        f = maximal_fan(u, v)
        c = free_color(u)
        d = free_color(f[-1])
        
        # invert the cd_u path
        path = cdu_path(c, d, u)
        for e in path:
            colors[e] = c if colors[e] == d else d
        
        # rotate the v...w fan left
        w = next(n for n in f[::-1] if is_free(d, n))
        for a, b in zip(f, f[1:]):
            if a == w:
                break
            colors[edge(u, a)] = colors[edge(u, b)]
        colors[edge(u, w)] = d
        
    return colors

import numpy as np
from itertools import product
import matplotlib.pyplot as plt


def plot_colors(colors, all_edges=None):
    plt.scatter(*locs.T)
    for idx, (x, y) in enumerate(locs):
        plt.text(x, y, str(idx))
    remaining = set([] if all_edges is None else all_edges)
    for (a, b), c in colors.items():
        remaining.discard(edge(a, b))
        x0, y0 = locs[a]
        x1, y1 = locs[b]
        plt.plot([x0, x1], [y0, y1], color=plt.rcParams["axes.prop_cycle"].by_key()["color"][c])
    for a, b in remaining:
        x0, y0 = locs[a]
        x1, y1 = locs[b]
        plt.plot([x0, x1], [y0, y1], color="grey", alpha=0.3)


locs = [[0,2],[0,3],[0,4],[0,5],[0,6],[1,2],[1,6],[2,0],[2,1],[2,2],[2,3],[2,4],[2,5],[2,6],[2,7],[2,8],[3,0],[3,4],[3,8],[4,0],[4,1],[4,2],[4,3],[4,4],[4,5],[4,6],[4,7],[4,8],[5,2],[5,6],[6,0],[6,1],[6,2],[6,3],[6,4],[6,5],[6,6],[6,7],[6,8],[7,0],[7,4],[7,8],[8,0],[8,1],[8,2],[8,3],[8,4],[8,5],[8,6],[8,7],[8,8],[9,2],[9,6]]
locs = np.array(locs)[:, ::-1]
edges = set()
for (idx, (x0, y0)), (jdx, (x1, y1)) in product(enumerate(locs), repeat=2):
    if idx != jdx and (x0 - x1) ** 2 + (y0 - y1) ** 2 < 1.1:
        edges.add(edge(idx, jdx))

plot_colors(misra_gries(edges))