use std::collections::BinaryHeap;
use std::hash::Hash;
use hashbrown::HashMap;
use petgraph::visit::{
Data, EdgeRef, IntoEdges, NodeCount, NodeIndexable, Visitable,
};
use pyo3::prelude::*;
use crate::astar::MinScored;
pub fn k_shortest_path<G, F>(
graph: G,
start: G::NodeId,
goal: Option<G::NodeId>,
k: usize,
mut edge_cost: F,
) -> PyResult<HashMap<G::NodeId, f64>>
where
G: IntoEdges + Visitable + NodeCount + NodeIndexable,
G: Data<NodeWeight = PyObject, EdgeWeight = PyObject>,
G::NodeId: Eq + Hash,
F: FnMut(&PyObject) -> PyResult<f64>,
{
let mut counter: Vec<usize> = vec![0; graph.node_count()];
let mut scores = HashMap::with_capacity(graph.node_count());
let mut visit_next = BinaryHeap::new();
let zero_score = 0.0;
visit_next.push(MinScored(zero_score, start));
while let Some(MinScored(node_score, node)) = visit_next.pop() {
counter[graph.to_index(node)] += 1;
let current_counter = counter[graph.to_index(node)];
if current_counter > k {
continue;
}
if current_counter == k {
scores.insert(node, node_score);
}
if goal.as_ref() == Some(&node) && current_counter == k {
break;
}
for edge in graph.edges(node) {
visit_next.push(MinScored(
node_score + edge_cost(edge.weight())?,
edge.target(),
));
}
}
Ok(scores)
}