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

def dijkstra_with_trace_and_visualization(graph, start):
    # 距離を無限大で初期化
    distances = {node: float('inf') for node in graph}
    distances[start] = 0
    priority_queue = [(0, start)]
    visited = set()
    trace_steps = []

    while priority_queue:
        current_distance, current_node = heapq.heappop(priority_queue)
        trace_steps.append((current_node, current_distance))
        visited.add(current_node)

        for neighbor, weight in graph[current_node].items():
            distance = current_distance + weight
            if distance < distances[neighbor]:
                distances[neighbor] = distance
                heapq.heappush(priority_queue, (distance, neighbor))

    return distances, trace_steps

def visualize_graph(graph, trace_steps):
    G = nx.Graph()
    for node, edges in graph.items():
        for neighbor, weight in edges.items():
            G.add_edge(node, neighbor, weight=weight)
    
    pos = nx.spring_layout(G)  # グラフの配置を決定
    nx.draw(G, pos, with_labels=True, node_size=700, node_color="skyblue", font_weight="bold")
    edge_labels = nx.get_edge_attributes(G, 'weight')
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)

    # 探索経路をハイライト
    for step in trace_steps:
        node = step[0]
        nx.draw_networkx_nodes(G, pos, nodelist=[node], node_color="yellow", node_size=800)

    plt.title("Dijkstra Algorithm Visualization")
    plt.show()

# グラフの定義
graph = {
    'A': {'B': 1, 'C': 4},
    'B': {'A': 1, 'C': 2, 'D': 5},
    'C': {'A': 4, 'B': 2, 'D': 1},
    'D': {'B': 5, 'C': 1}
}

# 実行
start_node = 'A'
distances, trace_steps = dijkstra_with_trace_and_visualization(graph, start_node)
print(f"Start node: {start_node}, Shortest distances: {distances}")

visualize_graph(graph, trace_steps)