factorio_blueprint/src/graph/wheighted_graph/shortest_path.rs

use std::{collections::HashMap, fmt::Debug, hash::Hash, hash::Hasher};

use crate::priority_queue::PriorityQueue;

use super::WheightedGraph;

#[derive(Debug, Clone, Copy)]
pub struct QueueObject<N> {
    node: N,
    score: f64,
}

impl<N> QueueObject<N> {
    fn new(node: N, score: f64) -> Self {
        Self { node, score }
    }
}

impl<N> PartialOrd for QueueObject<N> {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        self.score.partial_cmp(&other.score)
    }
}

impl<N> PartialEq for QueueObject<N> {
    fn eq(&self, other: &Self) -> bool {
        self.score == other.score
    }
}

impl<N> Hash for QueueObject<N>
where
    N: Hash,
{
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.node.hash(state);
    }
}

#[derive(Debug)]
struct MapObject<N, H> {
    parent: N,
    score: f64,
    key: Option<H>,
}

impl<N, H> MapObject<N, H> {
    fn new(parent: N, score: f64, key: H) -> Self {
        Self {
            parent,
            score,
            key: Some(key),
        }
    }
}

pub fn dijkstra<G, P>(graph: &G, start: G::Node, end: G::Node) -> Option<Vec<G::Node>>
where
    P: PriorityQueue<QueueObject<G::Node>> + Debug,
    P::Handle: Debug,
    G::Node: Eq + Hash + Clone + Debug,
    G: WheightedGraph,
{
    if start == end {
        return Some(vec![start]);
    }

    let mut map: HashMap<G::Node, MapObject<G::Node, P::Handle>> = HashMap::new();

    let mut q = P::new();
    q.insert(QueueObject::new(start.clone(), 0.0));

    while let Some(o) = q.pop_min() {
        if let Some(m) = map.get_mut(&o.node) {
            m.key = None;
        }

        if o.node == end {
            break;
        }

        for (node, wheight) in graph.edge_iter(&o.node) {
            let score = o.score + wheight;
            if let Some(n) = map.get_mut(&node) {
                if let Some(h) = &n.key {
                    if score < n.score {
                        n.parent = o.node.clone();
                        n.score = score;
                        q.decrease_key(h, |i| i.score = score);
                    }
                }
            } else {
                let h = q.insert(QueueObject::new(node.clone(), o.score + wheight));
                map.insert(node, MapObject::new(o.node.clone(), o.score + wheight, h));
            }
        }
        // dbg!(&q);
    }

    map.get(&end)?;

    let mut result = vec![end];

    // dbg!(&map);

    loop {
        let parent = map.get(result.last().expect("last")).expect("get");
        result.push(parent.parent.clone());
        if parent.parent == start {
            break;
        }
    }

    result.reverse();

    Some(result)
}

#[cfg(test)]
mod test {

    use crate::{
        graph::wheighted_graph::{
            adjacency_list::WheightedAdjacencyList, shortest_path::dijkstra, WheightedGraph,
        },
        priority_queue::{fibonacci_heap::FibonacciHeap, BinaryHeap},
    };

    #[test]
    fn trivial() {
        let a = WheightedAdjacencyList {
            nodes: vec![0, 1],
            edges: vec![(1, 1.0)],
        };

        assert_eq!(
            dijkstra::<WheightedAdjacencyList, BinaryHeap<_>>(&a, 0, 1),
            Some(vec![0, 1])
        );
    }

    #[test]
    fn simple() {
        let a = WheightedAdjacencyList {
            nodes: vec![0, 2, 3, 5, 5, 7, 10],
            edges: vec![
                (1, 2.0),
                (4, 10.0),
                (2, 3.0),
                (3, 2.0),
                (5, 1.0),
                (0, 4.0),
                (2, 5.0),
                (2, 9.0),
                (3, 8.0),
                (4, 0.0),
                (5, 7.0),
            ],
        };

        assert_eq!(
            dijkstra::<WheightedAdjacencyList, BinaryHeap<_>>(&a, 0, 4),
            Some(vec![0, 1, 2, 5, 4])
        );
        assert_eq!(
            dijkstra::<WheightedAdjacencyList, BinaryHeap<_>>(&a, 0, 6),
            None
        );
        assert_eq!(
            dijkstra::<WheightedAdjacencyList, FibonacciHeap<_>>(&a, 0, 4),
            Some(vec![0, 1, 2, 5, 4])
        );
        assert_eq!(
            dijkstra::<WheightedAdjacencyList, FibonacciHeap<_>>(&a, 0, 6),
            None
        );
    }

    struct Grid {
        width: usize,
        height: usize,
    }

    impl WheightedGraph for Grid {
        type Node = (usize, usize);

        fn num_edges(&self, node: &Self::Node) -> usize {
            let mut c = 0;

            if node.0 > 0 {
                c += 1
            }

            if node.0 < self.width - 1 {
                c += 1
            }

            if node.1 > 0 {
                c += 1
            }

            if node.1 < self.height - 1 {
                c += 1
            }

            c
        }

        fn edge(&self, node: &Self::Node, num: usize) -> Option<(Self::Node, f64)> {
            dbg!(&node);
            let edges = [
                (node.0 > 0).then_some((node.0.saturating_sub(1), node.1)),
                (node.0 < self.width - 1).then_some((node.0 + 1, node.1)),
                (node.1 > 0).then_some((node.0, node.1.saturating_sub(1))),
                (node.1 < self.height - 1).then_some((node.0, node.1 + 1)),
            ];

            edges.iter().flatten().nth(num).map(|&p| (p, 1.0))
        }
    }

    #[test]
    fn grid() {
        let g = Grid {
            width: 600,
            height: 600,
        };

        dbg!(dijkstra::<Grid, BinaryHeap<_>>(
            &g,
            (0, 0),
            (g.width - 1, g.height - 1)
        ));
    }
}