import heapq
from collections.abc import Hashable

Node = Hashable
edge = tuple[Node, Node, float]
adjacency = dict[Node, list[tuple[Node, float]]]


def _collect_nodes_and_edges(graph: adjacency) -> tuple[list[Node], list[edge]]:
    nodes = set()
    edges: list[edge] = []
    for u, neighbors in graph.items():
        nodes.add(u)
        for v, w in neighbors:
            nodes.add(v)
            edges.append((u, v, w))
    return list(nodes), edges


def _bellman_ford(nodes: list[Node], edges: list[edge]) -> dict[Node, float]:
    """
    Bellman-Ford relaxation to compute potentials h[v] for all vertices.
    Raises ValueError if a negative weight cycle exists.
    """
    dist: dict[Node, float] = dict.fromkeys(nodes, 0.0)
    n = len(nodes)

    for _ in range(n - 1):
        updated = False
        for u, v, w in edges:
            if dist[u] + w < dist[v]:
                dist[v] = dist[u] + w
                updated = True
        if not updated:
            break
    else:
        for u, v, w in edges:
            if dist[u] + w < dist[v]:
                raise ValueError("Negative weight cycle detected")
    return dist


def _dijkstra(
    start: Node,
    nodes: list[Node],
    graph: adjacency,
    potentials: dict[Node, float],
) -> dict[Node, float]:
    """
    Dijkstra over reweighted graph, using potentials h to make weights non-negative.
    Returns distances from start in the reweighted space.
    """
    inf = float("inf")
    dist: dict[Node, float] = dict.fromkeys(nodes, inf)
    dist[start] = 0.0
    heap: list[tuple[float, Node]] = [(0.0, start)]

    while heap:
        d_u, u = heapq.heappop(heap)
        if d_u > dist[u]:
            continue
        for v, w in graph.get(u, []):
            w_prime = w + potentials[u] - potentials[v]
            if w_prime < 0:
                raise ValueError(
                    "Negative edge weight after reweighting: numeric error"
                )
            new_dist = d_u + w_prime
            if new_dist < dist[v]:
                dist[v] = new_dist
                heapq.heappush(heap, (new_dist, v))
    return dist


def johnson(graph: adjacency) -> dict[Node, dict[Node, float]]:
    """
    Compute all-pairs shortest paths using Johnson's algorithm.

    Reference:
        https://en.wikipedia.org/wiki/Johnson%27s_algorithm

    Args:
        graph: adjacency list {u: [(v, weight), ...], ...}

    Returns:
        dict of dicts: dist[u][v] = shortest distance from u to v

    Raises:
        ValueError: if a negative weight cycle is detected

    Example:
    >>> g = {
    ...     0: [(1, 3), (2, 8), (4, -4)],
    ...     1: [(3, 1), (4, 7)],
    ...     2: [(1, 4)],
    ...     3: [(0, 2), (2, -5)],
    ...     4: [(3, 6)],
    ... }
    >>> round(johnson(g)[0][3], 2)
    2.0
    """
    nodes, edges = _collect_nodes_and_edges(graph)
    potentials = _bellman_ford(nodes, edges)

    all_pairs: dict[Node, dict[Node, float]] = {}
    inf = float("inf")
    for s in nodes:
        dist_reweighted = _dijkstra(s, nodes, graph, potentials)
        dists_orig: dict[Node, float] = {}
        for v in nodes:
            d_prime = dist_reweighted[v]
            if d_prime < inf:
                dists_orig[v] = d_prime - potentials[s] + potentials[v]
            else:
                dists_orig[v] = inf
        all_pairs[s] = dists_orig

    return all_pairs