import networkx as nx
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# Création du graphe
G = nx.DiGraph()
edges = [
    (0, 1, 4), (0, 2, 1),
    (2, 1, 2), (1, 3, 1),
    (2, 3, 5), (3, 4, 3)
]
G.add_weighted_edges_from(edges)

# Position des nœuds pour affichage
pos = nx.spring_layout(G, seed=42)
labels = {n: str(n) for n in G.nodes()}

# Initialisation de la figure
fig, ax = plt.subplots(figsize=(6, 6))

# Algorithme de Dijkstra avec animation
def dijkstra_animation(start_node=0):
    distances = {node: float('inf') for node in G.nodes()}
    distances[start_node] = 0
    visited = set()
    steps = []
    
    while len(visited) < len(G.nodes()):
        min_node = min((node for node in distances if node not in visited), key=lambda x: distances[x])
        visited.add(min_node)
        
        for neighbor in G.neighbors(min_node):
            new_distance = distances[min_node] + G[min_node][neighbor]["weight"]
            if new_distance < distances[neighbor]:
                distances[neighbor] = new_distance
        
        steps.append((visited.copy(), [(min_node, v) for v in G.neighbors(min_node)]))
    
    return steps

steps = dijkstra_animation()

def update(frame):
    ax.clear()
    visited, edges = steps[frame]
    node_colors = ["green" if n in visited else "lightblue" for n in G.nodes()]
    edge_colors = ["red" if (u, v) in edges else "black" for u, v in G.edges()]
    
    nx.draw(G, pos, with_labels=True, labels=labels, node_size=800, 
            node_color=node_colors, edge_color=edge_colors, font_size=14, font_weight="bold", ax=ax)
    edge_labels = {(u, v): G[u][v]["weight"] for u, v in G.edges()}
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=12, ax=ax)

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