use std::cmp::Ordering;
use std::collections::BinaryHeap;

#[derive(Copy, Clone, Eq, PartialEq)]
struct State {
    cost: u64,
    position: usize,
}

// BinaryHeap is a max-heap by default, so reverse the ordering to pop
// the smallest distance first.
impl Ord for State {
    fn cmp(&self, other: &Self) -> Ordering {
        other
            .cost
            .cmp(&self.cost)
            .then_with(|| self.position.cmp(&other.position))
    }
}

impl PartialOrd for State {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

/// Computes shortest paths from `start` over an adjacency list.
///
/// `graph[u]` contains `(v, weight)` edges from `u` to `v`.
/// All weights must be non-negative, represented here as `u64`.
/// Returns `(distances, predecessors)`, where unreachable nodes have
/// distance `None`.
pub fn dijkstra(graph: &[Vec<(usize, u64)>], start: usize) -> (Vec<Option<u64>>, Vec<Option<usize>>) {
    assert!(start < graph.len(), "start index out of bounds");

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

    dist[start] = 0;
    heap.push(State {
        cost: 0,
        position: start,
    });

    while let Some(State { cost, position }) = heap.pop() {
        if cost > dist[position] {
            continue;
        }

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

            let Some(next_cost) = cost.checked_add(weight) else {
                continue;
            };

            if next_cost < dist[next] {
                dist[next] = next_cost;
                prev[next] = Some(position);
                heap.push(State {
                    cost: next_cost,
                    position: next,
                });
            }
        }
    }

    let distances = dist
        .into_iter()
        .map(|d| if d == u64::MAX { None } else { Some(d) })
        .collect();

    (distances, prev)
}

/// Reconstructs a path from a predecessor list returned by `dijkstra`.
pub fn shortest_path(prev: &[Option<usize>], start: usize, goal: usize) -> Option<Vec<usize>> {
    if start >= prev.len() || goal >= prev.len() {
        return None;
    }

    let mut path = Vec::new();
    let mut current = goal;

    loop {
        path.push(current);

        if current == start {
            path.reverse();
            return Some(path);
        }

        current = prev[current]?;
    }
}

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

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

        let (dist, prev) = dijkstra(&graph, 0);

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

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

        let (dist, prev) = dijkstra(&graph, 0);

        assert_eq!(dist, vec![Some(0), Some(7), None]);
        assert_eq!(shortest_path(&prev, 0, 2), None);
    }
}