// Copyright 2026 Gregorio Momm
// Licensed under the Apache License, Version 2.0
//
// Ported from RustworkxCore (IBM Qiskit) — Apache License 2.0
// Source: https://github.com/Qiskit/rustworkx/tree/main/rustworkx-core/src/shortest_path

//! Additional shortest path algorithms ported from RustworkxCore.
//!
//! Complements the existing `path` module with:
//! - `dijkstra` — single-source shortest paths (Dijkstra)
//! - `all_shortest_paths` — all shortest paths between two nodes
//! - `single_source_all_shortest_paths` — all shortest paths from one source
//! - `k_shortest_path_lengths` — k shortest path lengths
//! - `distance_matrix` — all-pairs distance matrix as 2D Vec

use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque};
use std::cmp::Ordering;

use super::super::graph::{Graph, NodeId};

// ─── Min-heap state ───────────────────────────────────────────────────────────

#[derive(Clone)]
struct State {
    cost: f64,
    node: NodeId,
}

impl PartialEq for State {
    fn eq(&self, other: &Self) -> bool { self.cost == other.cost }
}
impl Eq for State {}
impl Ord for State {
    fn cmp(&self, other: &Self) -> Ordering {
        other.cost.partial_cmp(&self.cost).unwrap_or(Ordering::Equal)
    }
}
impl PartialOrd for State {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> { Some(self.cmp(other)) }
}

// ─── Dijkstra single source ───────────────────────────────────────────────────

/// Single-source shortest paths using Dijkstra's algorithm.
///
/// Returns a map `node → distance` for all nodes reachable from `source`.
/// Edges with no weight are treated as weight 1.0.
///
/// Runtime: O((n + m) log n)
///
/// Reference: RustworkxCore `dijkstra_shortest_paths`
pub fn dijkstra(
    graph: &Graph,
    source: NodeId,
    target: Option<NodeId>,
) -> HashMap<NodeId, f64> {
    let n = graph.upper_node_id_bound() as usize;
    let mut dist = vec![f64::INFINITY; n];
    let mut heap = BinaryHeap::new();

    dist[source as usize] = 0.0;
    heap.push(State { cost: 0.0, node: source });

    while let Some(State { cost, node }) = heap.pop() {
        if cost > dist[node as usize] { continue; }
        if target == Some(node) { break; }

        for e in graph.out_neighbors(node) {
            let next_cost = cost + e.weight.unwrap_or(1.0);
            if next_cost < dist[e.target as usize] {
                dist[e.target as usize] = next_cost;
                heap.push(State { cost: next_cost, node: e.target });
            }
        }
    }

    graph.nodes()
        .filter(|&n| dist[n as usize] < f64::INFINITY)
        .map(|n| (n, dist[n as usize]))
        .collect()
}

// ─── All shortest paths ───────────────────────────────────────────────────────

/// Find all shortest paths from `source` to `target` (unweighted BFS).
///
/// Returns all paths of minimum hop-count from source to target.
///
/// Runtime: O(n + m) for the BFS; path reconstruction may be exponential.
pub fn all_shortest_paths(
    graph: &Graph,
    source: NodeId,
    target: NodeId,
) -> Vec<Vec<NodeId>> {
    if source == target {
        return vec![vec![source]];
    }

    let n = graph.upper_node_id_bound() as usize;
    let mut dist = vec![u32::MAX; n];
    let mut predecessors: Vec<Vec<NodeId>> = vec![Vec::new(); n];
    let mut queue = VecDeque::new();

    dist[source as usize] = 0;
    queue.push_back(source);

    // BFS to compute distances and store all predecessors on shortest paths
    while let Some(node) = queue.pop_front() {
        let d = dist[node as usize];
        for nbr in graph.out_neighbors(node).iter().map(|e| e.target) {
            if dist[nbr as usize] == u32::MAX {
                dist[nbr as usize] = d + 1;
                predecessors[nbr as usize].push(node);
                queue.push_back(nbr);
            } else if dist[nbr as usize] == d + 1 {
                // Another shortest path predecessor
                predecessors[nbr as usize].push(node);
            }
        }
    }

    if dist[target as usize] == u32::MAX {
        return vec![];
    }

    // Reconstruct all paths backwards from target
    fn reconstruct(
        node: NodeId,
        source: NodeId,
        preds: &[Vec<NodeId>],
    ) -> Vec<Vec<NodeId>> {
        if node == source {
            return vec![vec![source]];
        }
        let mut paths = Vec::new();
        for &pred in &preds[node as usize] {
            for mut path in reconstruct(pred, source, preds) {
                path.push(node);
                paths.push(path);
            }
        }
        paths
    }

    reconstruct(target, source, &predecessors)
}

// ─── Single source all shortest paths ────────────────────────────────────────

/// All shortest paths from `source` to every other node (unweighted).
///
/// Returns a map `target → Vec<path>`.
///
/// Runtime: O(n * (n + m)) worst case
pub fn single_source_all_shortest_paths(
    graph: &Graph,
    source: NodeId,
) -> HashMap<NodeId, Vec<Vec<NodeId>>> {
    let n = graph.upper_node_id_bound() as usize;
    let mut dist = vec![u32::MAX; n];
    let mut predecessors: Vec<Vec<NodeId>> = vec![Vec::new(); n];
    let mut queue = VecDeque::new();

    dist[source as usize] = 0;
    queue.push_back(source);

    while let Some(node) = queue.pop_front() {
        let d = dist[node as usize];
        for nbr in graph.out_neighbors(node).iter().map(|e| e.target) {
            if dist[nbr as usize] == u32::MAX {
                dist[nbr as usize] = d + 1;
                predecessors[nbr as usize].push(node);
                queue.push_back(nbr);
            } else if dist[nbr as usize] == d + 1 {
                predecessors[nbr as usize].push(node);
            }
        }
    }

    fn reconstruct(
        node: NodeId,
        source: NodeId,
        preds: &[Vec<NodeId>],
    ) -> Vec<Vec<NodeId>> {
        if node == source {
            return vec![vec![source]];
        }
        let mut paths = Vec::new();
        for &pred in &preds[node as usize] {
            for mut path in reconstruct(pred, source, preds) {
                path.push(node);
                paths.push(path);
            }
        }
        paths
    }

    graph.nodes()
        .filter(|&t| t != source && dist[t as usize] < u32::MAX)
        .map(|t| (t, reconstruct(t, source, &predecessors)))
        .collect()
}

// ─── K shortest path lengths ──────────────────────────────────────────────────

/// Find the k shortest path lengths from `source` to `target`.
///
/// Uses Yen's algorithm for k-shortest loop-less paths.
/// Returns at most k (distance, path) pairs sorted by distance.
///
/// Runtime: O(k * n * (m + n log n))
pub fn k_shortest_paths(
    graph: &Graph,
    source: NodeId,
    target: NodeId,
    k: usize,
) -> Vec<(f64, Vec<NodeId>)> {
    // Dijkstra helper returning (dist_map, prev_map)
    fn dijkstra_with_path(
        graph: &Graph,
        source: NodeId,
        ignore_edges: &HashSet<(NodeId, NodeId)>,
        ignore_nodes: &HashSet<NodeId>,
    ) -> (HashMap<NodeId, f64>, HashMap<NodeId, NodeId>) {
        let n = graph.upper_node_id_bound() as usize;
        let mut dist = vec![f64::INFINITY; n];
        let mut prev: HashMap<NodeId, NodeId> = HashMap::new();
        let mut heap = BinaryHeap::new();

        dist[source as usize] = 0.0;
        heap.push(State { cost: 0.0, node: source });

        while let Some(State { cost, node }) = heap.pop() {
            if ignore_nodes.contains(&node) { continue; }
            if cost > dist[node as usize] { continue; }

            for e in graph.out_neighbors(node) {
                if ignore_edges.contains(&(node, e.target)) { continue; }
                if ignore_nodes.contains(&e.target) { continue; }
                let nc = cost + e.weight.unwrap_or(1.0);
                if nc < dist[e.target as usize] {
                    dist[e.target as usize] = nc;
                    prev.insert(e.target, node);
                    heap.push(State { cost: nc, node: e.target });
                }
            }
        }

        let dist_map = graph.nodes()
            .filter(|&n| dist[n as usize] < f64::INFINITY)
            .map(|n| (n, dist[n as usize]))
            .collect();
        (dist_map, prev)
    }

    fn extract_path(prev: &HashMap<NodeId, NodeId>, src: NodeId, tgt: NodeId) -> Option<Vec<NodeId>> {
        if !prev.contains_key(&tgt) && src != tgt { return None; }
        let mut path = vec![tgt];
        let mut cur = tgt;
        while cur != src {
            match prev.get(&cur) {
                Some(&p) => { path.push(p); cur = p; }
                None => return None,
            }
        }
        path.reverse();
        Some(path)
    }

    // A: confirmed shortest paths
    let mut a: Vec<(f64, Vec<NodeId>)> = Vec::new();
    // B: candidate paths (heap of (cost, path))
    // Use a manual newtype to make f64 orderable for BinaryHeap
    #[derive(PartialEq)]
    struct OrdF64(f64);
    impl Eq for OrdF64 {}
    impl PartialOrd for OrdF64 {
        fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { Some(self.cmp(other)) }
    }
    impl Ord for OrdF64 {
        fn cmp(&self, other: &Self) -> std::cmp::Ordering {
            self.0.partial_cmp(&other.0).unwrap_or(std::cmp::Ordering::Equal)
        }
    }

    let mut b: BinaryHeap<(OrdF64, Vec<NodeId>)> = BinaryHeap::new();

    // Find first shortest path
    let (dist_map, prev) = dijkstra_with_path(graph, source, &HashSet::new(), &HashSet::new());
    if let Some(d) = dist_map.get(&target) {
        if let Some(path) = extract_path(&prev, source, target) {
            a.push((*d, path));
        }
    }

    if a.is_empty() {
        return vec![];
    }

    for ki in 1..k {
        if a.len() <= ki - 1 { break; }
        let prev_path = &a[ki - 1].1;

        for i in 0..(prev_path.len().saturating_sub(1)) {
            let spur_node = prev_path[i];
            let root_path = &prev_path[..=i];

            let mut ignore_edges: HashSet<(NodeId, NodeId)> = HashSet::new();
            let mut ignore_nodes: HashSet<NodeId> = HashSet::new();

            // Remove edges already used in paths sharing root_path
            for (_, p) in &a {
                if p.len() > i && &p[..=i] == root_path {
                    if i + 1 < p.len() {
                        ignore_edges.insert((p[i], p[i + 1]));
                    }
                }
            }
            // Remove root path nodes (except spur node)
            for &n in &root_path[..i] {
                ignore_nodes.insert(n);
            }

            let (spur_dist, spur_prev) = dijkstra_with_path(graph, spur_node, &ignore_edges, &ignore_nodes);
            if let Some(&spur_d) = spur_dist.get(&target) {
                if let Some(spur_path) = extract_path(&spur_prev, spur_node, target) {
                    let mut total_path: Vec<NodeId> = root_path.to_vec();
                    total_path.extend_from_slice(&spur_path[1..]);

                    // Compute root cost
                    let root_cost: f64 = root_path.windows(2)
                        .map(|w| {
                            graph.out_neighbors(w[0])
                                .iter().find(|e| e.target == w[1])
                                .map(|e| e.weight.unwrap_or(1.0))
                                .unwrap_or(f64::INFINITY)
                        })
                        .sum();

                    let total_cost = root_cost + spur_d;
                    if !a.iter().any(|(_, p)| p == &total_path) {
                        b.push((OrdF64(-total_cost), total_path));
                    }
                }
            }
        }

        if b.is_empty() { break; }
        let (neg_cost, best_path) = b.pop().unwrap();
        a.push((-neg_cost.0, best_path));
    }

    a
}

// ─── Distance matrix ─────────────────────────────────────────────────────────

/// Compute the all-pairs distance matrix.
///
/// Returns a 2D Vec where `matrix[i][j]` is the shortest path distance
/// from node i to node j (f64::INFINITY if unreachable).
///
/// Uses BFS for unweighted graphs, Dijkstra for weighted.
///
/// Runtime: O(n * (n + m) log n) with Dijkstra
pub fn distance_matrix(graph: &Graph, use_weights: bool) -> Vec<Vec<f64>> {
    let nodes: Vec<NodeId> = graph.nodes().collect();
    let n = nodes.len();
    let bound = graph.upper_node_id_bound() as usize;

    let mut matrix = vec![vec![f64::INFINITY; n]; n];

    for (i, &src) in nodes.iter().enumerate() {
        matrix[i][i] = 0.0;

        if use_weights {
            let dist = dijkstra(graph, src, None);
            for (j, &tgt) in nodes.iter().enumerate() {
                if let Some(&d) = dist.get(&tgt) {
                    matrix[i][j] = d;
                }
            }
        } else {
            // Unweighted BFS
            let mut dist_bfs = vec![u32::MAX; bound];
            dist_bfs[src as usize] = 0;
            let mut queue = VecDeque::new();
            queue.push_back(src);
            while let Some(node) = queue.pop_front() {
                let d = dist_bfs[node as usize];
                for nbr in graph.out_neighbors(node).iter().map(|e| e.target) {
                    if dist_bfs[nbr as usize] == u32::MAX {
                        dist_bfs[nbr as usize] = d + 1;
                        queue.push_back(nbr);
                    }
                }
            }
            for (j, &tgt) in nodes.iter().enumerate() {
                if dist_bfs[tgt as usize] < u32::MAX {
                    matrix[i][j] = dist_bfs[tgt as usize] as f64;
                }
            }
        }
    }

    matrix
}

// ─── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::graph::backends::networkit_rust::graph::GraphConfig;

    fn path_directed(n: usize) -> Graph {
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..n { g.add_node(); }
        for i in 0..(n as NodeId - 1) { g.add_edge(i, i + 1, Some(1.0)); }
        g
    }

    // ── dijkstra ──────────────────────────────────────────────────────────────

    #[test]
    fn test_dijkstra_path() {
        let g = path_directed(4); // 0→1→2→3
        let dist = dijkstra(&g, 0, None);
        assert_eq!(dist[&0], 0.0);
        assert_eq!(dist[&1], 1.0);
        assert_eq!(dist[&2], 2.0);
        assert_eq!(dist[&3], 3.0);
    }

    #[test]
    fn test_dijkstra_unreachable() {
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..3 { g.add_node(); }
        g.add_edge(0, 1, Some(1.0));
        // Node 2 is unreachable from 0
        let dist = dijkstra(&g, 0, None);
        assert!(!dist.contains_key(&2));
    }

    #[test]
    fn test_dijkstra_weighted() {
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..3 { g.add_node(); }
        g.add_edge(0, 1, Some(1.0));
        g.add_edge(1, 2, Some(4.0));
        g.add_edge(0, 2, Some(10.0));
        let dist = dijkstra(&g, 0, None);
        assert_eq!(dist[&2], 5.0); // via 0→1→2
    }

    // ── all_shortest_paths ────────────────────────────────────────────────────

    #[test]
    fn test_all_shortest_paths_single() {
        let g = path_directed(4);
        let paths = all_shortest_paths(&g, 0, 3);
        assert_eq!(paths.len(), 1);
        assert_eq!(paths[0], vec![0, 1, 2, 3]);
    }

    #[test]
    fn test_all_shortest_paths_diamond() {
        // Diamond: 0→1→3 and 0→2→3
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..4 { g.add_node(); }
        g.add_edge(0, 1, None);
        g.add_edge(0, 2, None);
        g.add_edge(1, 3, None);
        g.add_edge(2, 3, None);
        let mut paths = all_shortest_paths(&g, 0, 3);
        paths.sort();
        assert_eq!(paths.len(), 2);
        assert_eq!(paths[0], vec![0, 1, 3]);
        assert_eq!(paths[1], vec![0, 2, 3]);
    }

    #[test]
    fn test_all_shortest_paths_unreachable() {
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..2 { g.add_node(); }
        // No edge from 0 to 1
        let paths = all_shortest_paths(&g, 0, 1);
        assert!(paths.is_empty());
    }

    // ── single_source_all_shortest_paths ──────────────────────────────────────

    #[test]
    fn test_single_source_all_shortest_paths() {
        let g = path_directed(4);
        let paths = single_source_all_shortest_paths(&g, 0);
        assert_eq!(paths[&1], vec![vec![0, 1]]);
        assert_eq!(paths[&3], vec![vec![0, 1, 2, 3]]);
    }

    // ── k_shortest_paths ──────────────────────────────────────────────────────

    #[test]
    fn test_k_shortest_paths_basic() {
        // Two paths from 0 to 3: 0→1→3 (cost 2) and 0→2→3 (cost 2)
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..4 { g.add_node(); }
        g.add_edge(0, 1, Some(1.0));
        g.add_edge(0, 2, Some(1.0));
        g.add_edge(1, 3, Some(1.0));
        g.add_edge(2, 3, Some(1.0));
        let result = k_shortest_paths(&g, 0, 3, 2);
        assert_eq!(result.len(), 2);
        assert!((result[0].0 - 2.0).abs() < 1e-9);
        assert!((result[1].0 - 2.0).abs() < 1e-9);
    }

    #[test]
    fn test_k_shortest_paths_single() {
        let g = path_directed(4);
        let result = k_shortest_paths(&g, 0, 3, 3);
        // Only one simple path exists
        assert_eq!(result.len(), 1);
        assert_eq!(result[0].1, vec![0, 1, 2, 3]);
    }

    // ── distance_matrix ───────────────────────────────────────────────────────

    #[test]
    fn test_distance_matrix_path() {
        let g = path_directed(3); // 0→1→2
        let m = distance_matrix(&g, false);
        assert_eq!(m[0][0], 0.0);
        assert_eq!(m[0][1], 1.0);
        assert_eq!(m[0][2], 2.0);
        assert_eq!(m[1][2], 1.0);
        assert_eq!(m[2][0], f64::INFINITY); // no reverse edge
    }

    #[test]
    fn test_distance_matrix_weighted() {
        let mut g = Graph::new(GraphConfig::directed());
        for _ in 0..2 { g.add_node(); }
        g.add_edge(0, 1, Some(3.0));
        let m = distance_matrix(&g, true);
        assert_eq!(m[0][1], 3.0);
        assert_eq!(m[1][0], f64::INFINITY);
    }
}
