import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import heapq
from ACO import crea_grafo_random

# Crea grafo pesato
G = nx.Graph()
edges = [
    ('A', 'B', 1), ('B', 'C', 2), ('C', 'D', 1),
    ('D', 'A', 3), ('C', 'E', 4), ('D', 'F', 2),
    ('E', 'F', 1)
]
G.add_weighted_edges_from(edges)

# Posizioni fisse per visualizzazione
pos = {
    'A': (0, 0), 'B': (1, 1), 'C': (2, 0),
    'D': (1, -1), 'E': (3, 1), 'F': (3, -1)
}

G=crea_grafo_random(4,6)
pos = nx.spring_layout(G)
# Dizionario delle distanze minime e dei predecessori
dist = {node: float('inf') for node in G.nodes}
prev = {node: None for node in G.nodes}

# Nodo di partenza
start = '11'
dist[start] = 0

# Coda con priorità
heap = [(0, start)]
visited = set()

# Salva gli stati per animazione
steps = []

# Dijkstra passo passo
while heap:
    d, current = heapq.heappop(heap)
    if current in visited:
        continue
    visited.add(current)

    # Salva stato attuale
    steps.append((current, dist.copy(), list(visited)))

    for neighbor in G.neighbors(current):
        weight = G[current][neighbor]['weight']
        if dist[current] + weight < dist[neighbor]:
            dist[neighbor] = dist[current] + weight
            prev[neighbor] = current
            heapq.heappush(heap, (dist[neighbor], neighbor))

# Animazione
fig, ax = plt.subplots(figsize=(10, 8))

def update(i):
    ax.clear()
    current, dist_now, visited_now = steps[i]

    # Colora nodi visitati
    node_colors = []
    for node in G.nodes:
        if node == current:
            node_colors.append('red')  # "persona" che si muove
        elif node in visited_now:
            node_colors.append('lightgreen')
        else:
            node_colors.append('lightgray')
        if node == start:
            node_colors[-1] = 'blue'  # Nodo iniziale

    # Etichette con distanza attuale
    labels = {node: f"{node}\n{dist_now[node] if dist_now[node] != float('inf') else '∞'}" for node in G.nodes}

    # Disegna grafo
    nx.draw(G, pos, ax=ax, with_labels=False,
            node_color=node_colors, edge_color='gray', node_size=600)
    nx.draw_networkx_labels(G, pos, labels, font_size=6, font_weight='bold')

    # Disegna i pesi sugli archi
    edge_labels = nx.get_edge_attributes(G, 'weight')
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)

    ax.set_title(f"Dijkstra - Step {i+1}: Visito {current}", fontsize=14)

ani = animation.FuncAnimation(fig, update, frames=len(steps), interval=900, repeat=False)
plt.tight_layout()
plt.show()