#!/usr/bin/python3
from random import randint, random
import networkx as nx
import sys

DEBUG = False

class NodeInformation:
    def __init__(self, color_amt):
        self.available_colors = [i for i in range(0, color_amt)]
        self.neighbor_colors = []
        self.is_done = False
        self.assigned_color = None

color_amt = None
ni = {}
g = None

def run_tick():
    global color_amt, ni, g
    if DEBUG: print("===== RUN TICK =====")
    # run for each node
    for n in g.nodes:
        if ni[n].is_done:
            continue
        neighbors = list(g.adj[n])

        # each uncolored node selects random candidate color
        available_colors = ni[n].available_colors
        ni[n].assigned_color = ni[n].available_colors[randint(0, len(available_colors)-1)]
        if DEBUG: print(f"{n} | assigned_color: {ni[n].assigned_color}")
        if DEBUG: print(f"{n} | neighbors: {neighbors}")

        # nodes exchange candidate colors
        for neigh in neighbors:
            if not ni[neigh].is_done:
                ni[neigh].neighbor_colors.append(ni[n].assigned_color)

    # run again, but only after neighbor_colors assigned for all nodes
    for n in g.nodes:
        if ni[n].is_done:
            continue
        neighbors = list(g.adj[n])

        if DEBUG: print(f"{n} | neighbor_colors: {ni[n].neighbor_colors}, assigned_color: {ni[n].assigned_color}")

        # check if assigned color is used by neighbor
        if ni[n].assigned_color in ni[n].neighbor_colors:
            # reset
            ni[n].assigned_color = None
            ni[n].neighbor_colors = []
        else:
            # get permantly colored
            ni[n].is_done = True
            # remove color from available colors of neighbor
            for neigh in neighbors:
                if not ni[neigh].is_done:
                    if ni[n].assigned_color in ni[neigh].available_colors:
                        ni[neigh].available_colors.remove(ni[n].assigned_color)

    is_done = True
    for n in g.nodes:
        is_done &= ni[n].is_done
    
    return 0 if is_done else 1

def check_correctness():
    global color_amt, ni, g
    total_colors = []
    for n in g.nodes:
        if ni[n].assigned_color not in total_colors:
            total_colors.append(ni[n].assigned_color)
        neighbors = list(g.adj[n])
        for neigh in neighbors:
            if ni[n].assigned_color == ni[neigh].assigned_color:
                print(f"FAIL: {n} and {neigh} have same color {ni[n].assigned_color}")
                assert False
        
    if len(total_colors) > color_amt:
        print(f"FAIL: color_amt not correct: {len(total_colors)} <= {color_amt}")
        assert False

def main():
    global color_amt, ni, g
    amt_nodes = randint(100, 300)
    if len(sys.argv) == 2:
        amt_nodes = int(sys.argv[1])
    
    g = nx.erdos_renyi_graph(amt_nodes, random())

    # find out max degree
    max_degree = 0
    for n in g.nodes:
        max_degree = max(max_degree, g.degree[n])

    color_amt = max_degree + 1
    print(f"Amount of Colors: {color_amt}")

    for n in g.nodes:
        ni[n] = NodeInformation(color_amt)    

    ticks = 0
    while True:
        ret = run_tick()
        ticks += 1
        if ret == 0:
            # print info
            print("Final colors:")
            for n in g.nodes:
                print(f"{n}[{ni[n].assigned_color}] ", end="")
            print()
            print(f"Amount of Ticks: {ticks}")

            # assert if correct
            check_correctness()
            exit(0)

if __name__ == '__main__':
    main()