from random import choice

from networkx import (
    Graph, bipartite_layout, draw_networkx_nodes, draw_networkx_edges, set_node_attributes,
    get_node_attributes, get_edge_attributes, circular_layout
)
from networkx.algorithms.bipartite import random_graph
from networkx.exception import AmbiguousSolution

from altbacken.external.annealing import SimpleSimulatedAnnealing
from altbacken.external.neighbourhood.graph import GraphColoringNeighbourhood
from altbacken.external.temperature import LogarithmicCooling, ExponentialCooling
from altbacken.external.neighbourhood.graph import color_adjustments_for_limited_colors

def show_graph(graph: Graph):
    try:
        pos = bipartite_layout(graph)
    except AmbiguousSolution:
        pos = circular_layout(graph)
    for edge in graph.edges:
        if graph.nodes[edge[0]]["color"] == graph.nodes[edge[1]]["color"]:
            graph.edges[edge]["color"] = "red"
        else:
            graph.edges[edge]["color"] = "white"
    draw_networkx_nodes(graph, pos, node_color=get_node_attributes(graph, "color").values())
    draw_networkx_edges(graph, pos, edge_color=get_edge_attributes(graph, "color").values())

COLORS: list[str] = [
    "#E69F00",  # orange
    "#56B4E9",  # sky blue
    "#009E73",  # bluish green
    "#F0E442",  # yellow
    "#0072B2",  # blue
    "#D55E00",  # vermilion
    "#CC79A7",  # reddish purple
    "#999999",  # gray
]

def fitness(solution: Graph) -> float:
    color_count: int = len(set(get_node_attributes(solution, "color").values()))
    invalid_edges: int = sum(
        1
        for source, target in solution.edges
        if solution.nodes[source]["color"] == solution.nodes[target]["color"]
    )
    return (color_count - 2)**2 + invalid_edges

def create_random_graph(n: int=10, m: int=10, p: float=0.2) -> Graph:
    graph: Graph = random_graph(n, m, p)
    set_node_attributes(graph, COLORS[0], "color")
    return graph

graph: Graph = create_random_graph()

annealing: SimpleSimulatedAnnealing[Graph] = SimpleSimulatedAnnealing(
    fitness,
    GraphColoringNeighbourhood(color_adjustments_for_limited_colors(COLORS)),
    ExponentialCooling(10000)
)
annealing_iterations: list[int] = []
annealing.tracer = lambda x: annealing_iterations.append(x.iteration)

annealing_best, annealing_best_value = annealing.simulate(graph.copy())


print(f"Best solution value: {annealing_best_value}")
print(f"Iterations: {len(annealing_iterations)}")
print(f"Colors used: {len(set(get_node_attributes(annealing_best, 'color').values()))}")

Best solution value: 1
Iterations: 1000
Colors used: 3

show_graph(annealing_best)

def random_coloring(graph: Graph, iterations: int) -> tuple[Graph, float]:
    best: Graph = graph.copy()
    best_value: float = fitness(graph)
    for _ in range(iterations):
        new_graph: Graph = graph.copy()
        for node in new_graph.nodes:
            new_graph.nodes[node]["color"] = choice(COLORS)
        new_value: float = fitness(new_graph)
        if new_value < best_value:
            best = new_graph.copy()
            best_value = new_value
    return best, best_value

random_best, random_best_value = random_coloring(graph, len(annealing_iterations))
print(f"Best solution value (random): {random_best_value}")
print(f"Iterations (random): {len(annealing_iterations)}")
print(f"Colors used (random): {len(set(get_node_attributes(random_best, 'color').values()))}")

Best solution value (random): 7
Iterations (random): 1000
Colors used (random): 4

show_graph(random_best)