use std::cmp::Reverse;
use std::collections::BinaryHeap;

/// Adjacency list graph: graph[u] contains (v, weight) edges.
/// Edge weights must be non-negative, which is guaranteed by `u64`.
pub type Graph = Vec<Vec<(usize, u64)>>;

/// Returns the shortest distance from `start` to every node.
/// `None` means the node is unreachable.
pub fn dijkstra(graph: &[Vec<(usize, u64)>], start: usize) -> Vec<Option<u64>> {
    assert!(start < graph.len(), "start node is out of bounds");

    let mut dist = vec![None; graph.len()];
    let mut heap = BinaryHeap::new();

    dist[start] = Some(0);
    heap.push(Reverse((0_u64, start)));

    while let Some(Reverse((cost, node))) = heap.pop() {
        // Ignore stale heap entries left behind after a better path was found.
        if dist[node] != Some(cost) {
            continue;
        }

        for &(next, weight) in &graph[node] {
            assert!(next < graph.len(), "edge points to an out-of-bounds node");

            let next_cost = cost
                .checked_add(weight)
                .expect("shortest path distance overflowed u64");

            if dist[next].map_or(true, |current| next_cost < current) {
                dist[next] = Some(next_cost);
                heap.push(Reverse((next_cost, next)));
            }
        }
    }

    dist
}

/// Returns one shortest path from `start` to `goal`, if reachable.
pub fn shortest_path(graph: &[Vec<(usize, u64)>], start: usize, goal: usize) -> Option<(u64, Vec<usize>)> {
    assert!(start < graph.len(), "start node is out of bounds");
    assert!(goal < graph.len(), "goal node is out of bounds");

    let mut dist = vec![None; graph.len()];
    let mut prev = vec![None; graph.len()];
    let mut heap = BinaryHeap::new();

    dist[start] = Some(0);
    heap.push(Reverse((0_u64, start)));

    while let Some(Reverse((cost, node))) = heap.pop() {
        if dist[node] != Some(cost) {
            continue;
        }

        if node == goal {
            break;
        }

        for &(next, weight) in &graph[node] {
            assert!(next < graph.len(), "edge points to an out-of-bounds node");

            let next_cost = cost
                .checked_add(weight)
                .expect("shortest path distance overflowed u64");

            if dist[next].map_or(true, |current| next_cost < current) {
                dist[next] = Some(next_cost);
                prev[next] = Some(node);
                heap.push(Reverse((next_cost, next)));
            }
        }
    }

    let total = dist[goal]?;
    let mut path = Vec::new();
    let mut node = goal;

    loop {
        path.push(node);
        if node == start {
            break;
        }
        node = prev[node]?;
    }

    path.reverse();
    Some((total, path))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn computes_shortest_distances() {
        let graph = vec![
            vec![(1, 4), (2, 1)],
            vec![(3, 1)],
            vec![(1, 2), (3, 5)],
            vec![],
            vec![(0, 10)],
        ];

        assert_eq!(dijkstra(&graph, 0), vec![Some(0), Some(3), Some(1), Some(4), None]);
    }

    #[test]
    fn reconstructs_shortest_path() {
        let graph = vec![
            vec![(1, 4), (2, 1)],
            vec![(3, 1)],
            vec![(1, 2), (3, 5)],
            vec![],
        ];

        assert_eq!(shortest_path(&graph, 0, 3), Some((4, vec![0, 2, 1, 3])));
        assert_eq!(shortest_path(&graph, 3, 0), None);
    }
}