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

/// Adjacency-list graph: `graph[u]` contains `(v, weight)` edges from `u` to `v`.
pub type Graph = Vec<Vec<(usize, u64)>>;

/// Returns the shortest distance from `source` to every vertex.
///
/// Unreachable vertices are returned as `None`.
pub fn dijkstra(graph: &[Vec<(usize, u64)>], source: usize) -> Vec<Option<u64>> {
    assert!(source < graph.len(), "source vertex out of bounds");

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

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

    while let Some(Reverse((cost, u))) = heap.pop() {
        // Skip stale heap entries created before a better path was found.
        if dist[u] != Some(cost) {
            continue;
        }

        for &(v, weight) in &graph[u] {
            assert!(v < graph.len(), "edge target out of bounds");

            let next = cost.checked_add(weight).expect("path cost overflowed u64");

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

    dist
}

/// Returns the shortest path and distance from `source` to `target`, if reachable.
pub fn shortest_path(
    graph: &[Vec<(usize, u64)>],
    source: usize,
    target: usize,
) -> Option<(u64, Vec<usize>)> {
    assert!(source < graph.len(), "source vertex out of bounds");
    assert!(target < graph.len(), "target vertex out of bounds");

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

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

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

        if u == target {
            break;
        }

        for &(v, weight) in &graph[u] {
            assert!(v < graph.len(), "edge target out of bounds");

            let next = cost.checked_add(weight).expect("path cost overflowed u64");

            if dist[v].map_or(true, |best| next < best) {
                dist[v] = Some(next);
                prev[v] = Some(u);
                heap.push(Reverse((next, v)));
            }
        }
    }

    let total = dist[target]?;
    let mut path = Vec::new();
    let mut current = target;

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

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

fn main() {}

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

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

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

    #[test]
    fn reconstructs_shortest_path() {
        let graph: Graph = vec![
            vec![(1, 10), (2, 3)],
            vec![(3, 2)],
            vec![(1, 4), (3, 8), (4, 2)],
            vec![(4, 1)],
            vec![(3, 9)],
        ];

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

    #[test]
    fn reports_unreachable_target() {
        let graph: Graph = vec![vec![(1, 7)], vec![], vec![]];

        assert_eq!(shortest_path(&graph, 0, 2), None);
    }
}