// 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

//! Steiner tree approximation algorithms.
//!
//! - `metric_closure` — complete weighted graph of shortest-path distances
//! - `steiner_tree` — approximate minimum Steiner tree (Kou et al. 2-approximation)
//!
//! Reference:
//! Kou, L., Markowsky, G., & Berman, L. (1981).
//! "A fast algorithm for Steiner trees."
//! Acta Informatica, 15(2), 141-145.

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

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

// ─── OrdF64 (ordered float helper) ──────────────────────────────────────────

#[derive(Clone, PartialEq)]
struct OrdF64(f64);
impl Eq for OrdF64 {}
impl Ord for OrdF64 {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.0.partial_cmp(&other.0).unwrap_or(std::cmp::Ordering::Equal)
    }
}
impl PartialOrd for OrdF64 {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

// ─── Dijkstra helper ─────────────────────────────────────────────────────────

/// Single-source Dijkstra returning (distances, predecessors).
fn dijkstra_sssp(
    graph: &Graph,
    source: NodeId,
) -> (HashMap<NodeId, f64>, HashMap<NodeId, NodeId>) {
    let mut dist: HashMap<NodeId, f64> = HashMap::new();
    let mut pred: HashMap<NodeId, NodeId> = HashMap::new();
    // (Reverse(dist), node)
    let mut heap: BinaryHeap<(Reverse<OrdF64>, NodeId)> = BinaryHeap::new();

    dist.insert(source, 0.0);
    heap.push((Reverse(OrdF64(0.0)), source));

    while let Some((Reverse(OrdF64(d)), u)) = heap.pop() {
        if dist.get(&u).copied().unwrap_or(f64::INFINITY) < d { continue; }
        for e in graph.out_neighbors(u).iter() {
            let w = e.weight.unwrap_or(1.0);
            let new_d = d + w;
            if new_d < dist.get(&e.target).copied().unwrap_or(f64::INFINITY) {
                dist.insert(e.target, new_d);
                pred.insert(e.target, u);
                heap.push((Reverse(OrdF64(new_d)), e.target));
            }
        }
    }

    (dist, pred)
}

/// Reconstruct path from source to target using predecessor map.
fn reconstruct_path(pred: &HashMap<NodeId, NodeId>, source: NodeId, target: NodeId) -> Vec<NodeId> {
    let mut path = Vec::new();
    let mut cur = target;
    while cur != source {
        path.push(cur);
        match pred.get(&cur) {
            Some(&p) => cur = p,
            None => return vec![], // disconnected
        }
    }
    path.push(source);
    path.reverse();
    path
}

// ─── Metric Closure ──────────────────────────────────────────────────────────

/// Result of metric closure computation.
pub struct MetricClosure {
    /// The complete graph on the terminal nodes.
    /// Node IDs in this graph correspond to the original graph's node IDs.
    pub graph: Graph,
    /// Distances between all pairs of terminals.
    /// `distances[u][v]` = shortest path distance from u to v in original graph.
    pub distances: HashMap<NodeId, HashMap<NodeId, f64>>,
    /// Predecessors for path reconstruction.
    pub predecessors: HashMap<NodeId, HashMap<NodeId, NodeId>>,
}

/// Compute the metric closure of a graph restricted to a set of terminal nodes.
///
/// The metric closure is a complete weighted graph on the terminal nodes where
/// the weight of edge (u, v) is the shortest-path distance from u to v in the
/// original graph.
///
/// Runtime: O(|terminals| × (m + n log n)) using Dijkstra from each terminal.
pub fn metric_closure(graph: &Graph, terminals: &[NodeId]) -> Option<MetricClosure> {
    let mut distances: HashMap<NodeId, HashMap<NodeId, f64>> = HashMap::new();
    let mut predecessors: HashMap<NodeId, HashMap<NodeId, NodeId>> = HashMap::new();

    for &t in terminals {
        let (dist, pred) = dijkstra_sssp(graph, t);
        // Check all terminals are reachable
        for &other in terminals {
            if other != t && !dist.contains_key(&other) {
                return None; // terminals disconnected
            }
        }
        distances.insert(t, dist);
        predecessors.insert(t, pred);
    }

    // Build the metric closure graph
    let mut closure = Graph::new(GraphConfig::directed()); // weighted complete graph on terminals
    let terminal_set: HashSet<NodeId> = terminals.iter().copied().collect();

    // Ensure all terminal nodes exist in closure (add them in order)
    let mut sorted_terminals: Vec<NodeId> = terminals.to_vec();
    sorted_terminals.sort_unstable();

    for &t in &sorted_terminals {
        // We store each terminal as a node (node id matches original)
        // Use a local mapping since Graph auto-assigns ids
        let _ = t; // will use edge weights to embed distances
    }

    // Add all nodes first
    let mut node_map: HashMap<NodeId, NodeId> = HashMap::new();
    for &t in &sorted_terminals {
        let local_id = closure.add_node();
        node_map.insert(t, local_id);
    }

    // Add edges between all pairs
    for &u in &sorted_terminals {
        for &v in &sorted_terminals {
            if u == v { continue; }
            if !terminal_set.contains(&v) { continue; }
            let d = distances[&u].get(&v).copied().unwrap_or(f64::INFINITY);
            closure.add_edge(node_map[&u], node_map[&v], Some(d));
        }
    }

    Some(MetricClosure { graph: closure, distances, predecessors })
}

// ─── Steiner Tree ─────────────────────────────────────────────────────────────

/// Approximate minimum Steiner tree connecting a set of terminal nodes.
///
/// Uses the Kou et al. 2-approximation algorithm:
/// 1. Compute metric closure on terminals.
/// 2. Find MST of the metric closure.
/// 3. Expand each edge (s, t) in the MST to the actual shortest path s→t.
/// 4. Extract an MST of the expanded subgraph (remove any cycles).
///
/// The result spans all terminal nodes and is at most 2×OPT in weight.
///
/// # Returns
/// A subgraph of `graph` that is a tree spanning all terminals,
/// or `None` if the terminals are not all reachable from each other.
pub fn steiner_tree(graph: &Graph, terminals: &[NodeId]) -> Option<Graph> {
    if terminals.is_empty() { return Some(Graph::new(GraphConfig::simple())); }
    if terminals.len() == 1 {
        let mut g = Graph::new(GraphConfig::simple());
        g.add_node();
        return Some(g);
    }

    // Step 1: Compute metric closure
    let mut sorted_terminals: Vec<NodeId> = terminals.to_vec();
    sorted_terminals.sort_unstable();
    let terminal_set: HashSet<NodeId> = sorted_terminals.iter().copied().collect();

    // Compute all-pairs shortest paths between terminals
    let mut sp_dist: HashMap<NodeId, HashMap<NodeId, f64>> = HashMap::new();
    let mut sp_pred: HashMap<NodeId, HashMap<NodeId, NodeId>> = HashMap::new();

    for &t in &sorted_terminals {
        let (dist, pred) = dijkstra_sssp(graph, t);
        for &other in &sorted_terminals {
            if other != t && !dist.contains_key(&other) {
                return None; // disconnected
            }
        }
        sp_dist.insert(t, dist);
        sp_pred.insert(t, pred);
    }

    // Step 2: Find MST of metric closure using Prim's algorithm
    // (Kruskal's would work too, but Prim is simpler to implement here)
    let mut in_mst: HashSet<NodeId> = HashSet::new();
    let mut mst_edges: Vec<(NodeId, NodeId, f64)> = Vec::new(); // (u, v, weight)

    // Start with first terminal
    let start = sorted_terminals[0];
    in_mst.insert(start);

    while in_mst.len() < sorted_terminals.len() {
        // Find minimum edge from in_mst to a non-mst terminal
        let mut best: Option<(NodeId, NodeId, f64)> = None;
        for &u in &in_mst {
            for &v in &sorted_terminals {
                if in_mst.contains(&v) { continue; }
                let d = sp_dist[&u].get(&v).copied().unwrap_or(f64::INFINITY);
                if d.is_finite() {
                    if best.is_none() || d < best.unwrap().2 {
                        best = Some((u, v, d));
                    }
                }
            }
        }
        let (u, v, w) = best?;
        in_mst.insert(v);
        mst_edges.push((u, v, w));
    }

    // Step 3: Expand each MST edge (u, v) to the actual shortest path
    let mut steiner_nodes: HashSet<NodeId> = terminal_set.clone();
    let mut steiner_edges: HashSet<(NodeId, NodeId)> = HashSet::new();

    for (u, v, _) in &mst_edges {
        let path = reconstruct_path(&sp_pred[u], *u, *v);
        for node in &path {
            steiner_nodes.insert(*node);
        }
        for window in path.windows(2) {
            let (a, b) = (window[0].min(window[1]), window[0].max(window[1]));
            steiner_edges.insert((a, b));
        }
    }

    // Step 4: Build the Steiner subgraph and extract MST to remove cycles
    // Map original NodeIds to local ids
    let mut sorted_nodes: Vec<NodeId> = steiner_nodes.into_iter().collect();
    sorted_nodes.sort_unstable();
    let node_local: HashMap<NodeId, NodeId> = sorted_nodes.iter()
        .enumerate()
        .map(|(i, &orig)| (orig, i as NodeId))
        .collect();

    let mut result = Graph::new(GraphConfig::simple());
    for _ in 0..sorted_nodes.len() { result.add_node(); }

    // Add edges using Kruskal-style union-find to build an MST
    let mut uf: Vec<usize> = (0..sorted_nodes.len()).collect();

    fn find(uf: &mut Vec<usize>, x: usize) -> usize {
        if uf[x] != x { uf[x] = find(uf, uf[x]); }
        uf[x]
    }

    // Collect edges with weights
    let mut edge_list: Vec<(f64, NodeId, NodeId)> = steiner_edges.iter()
        .map(|&(a, b)| {
            // Look up weight from original graph
            let w = graph.out_neighbors(a).iter()
                .find(|e| e.target == b)
                .map(|e| e.weight.unwrap_or(1.0))
                .unwrap_or(1.0);
            (w, a, b)
        })
        .collect();
    edge_list.sort_by(|x, y| x.0.partial_cmp(&y.0).unwrap());

    for (w, a, b) in edge_list {
        let la = node_local[&a] as usize;
        let lb = node_local[&b] as usize;
        let ra = find(&mut uf, la);
        let rb = find(&mut uf, lb);
        if ra != rb {
            uf[ra] = rb;
            result.add_edge(node_local[&a], node_local[&b], Some(w));
        }
    }

    Some(result)
}

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

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

    fn path_graph_4() -> Graph {
        // 0 -- 1 -- 2 -- 3 (unit weights)
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..4 { g.add_node(); }
        g.add_edge(0, 1, Some(1.0));
        g.add_edge(1, 2, Some(1.0));
        g.add_edge(2, 3, Some(1.0));
        g
    }

    fn triangle_graph() -> Graph {
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..3 { g.add_node(); }
        g.add_edge(0, 1, Some(1.0));
        g.add_edge(1, 2, Some(1.0));
        g.add_edge(0, 2, Some(2.0));
        g
    }

    // ── metric_closure ────────────────────────────────────────────────────────

    #[test]
    fn test_metric_closure_path() {
        let g = path_graph_4();
        let terminals = vec![0u64, 3];
        let mc = metric_closure(&g, &terminals).expect("connected");
        // Distance from 0 to 3 = 3
        assert!((mc.distances[&0][&3] - 3.0).abs() < 1e-9,
            "d(0,3) = {}", mc.distances[&0][&3]);
    }

    #[test]
    fn test_metric_closure_triangle() {
        let g = triangle_graph();
        let terminals = vec![0u64, 2];
        let mc = metric_closure(&g, &terminals).expect("connected");
        // Shortest path 0→2 = direct edge weight 2 vs 0→1→2 = 2 (same!)
        // Dijkstra should give min = 2
        assert!((mc.distances[&0][&2] - 2.0).abs() < 1e-9,
            "d(0,2) = {}", mc.distances[&0][&2]);
    }

    #[test]
    fn test_metric_closure_disconnected() {
        // Two isolated nodes
        let mut g = Graph::new(GraphConfig::simple());
        g.add_node(); g.add_node();
        let mc = metric_closure(&g, &[0, 1]);
        assert!(mc.is_none(), "disconnected should return None");
    }

    #[test]
    fn test_metric_closure_single_terminal() {
        let g = path_graph_4();
        let mc = metric_closure(&g, &[1]).expect("single terminal ok");
        // Only one terminal, distances map has just node 1
        assert!(mc.distances.contains_key(&1));
    }

    // ── steiner_tree ──────────────────────────────────────────────────────────

    #[test]
    fn test_steiner_tree_path_endpoints() {
        // Path 0-1-2-3; terminals = {0, 3}; optimal Steiner tree = entire path (3 edges)
        let g = path_graph_4();
        let tree = steiner_tree(&g, &[0, 3]).expect("connected");
        // Tree should have 4 nodes and 3 edges
        assert_eq!(tree.node_count(), 4, "4 nodes in Steiner tree");
        assert_eq!(tree.edge_count(), 3, "3 edges in Steiner tree");
    }

    #[test]
    fn test_steiner_tree_path_middle() {
        // Path 0-1-2-3; terminals = {0, 2}; optimal = 0-1-2 (2 edges)
        let g = path_graph_4();
        let tree = steiner_tree(&g, &[0, 2]).expect("connected");
        assert_eq!(tree.node_count(), 3, "3 nodes spanning 0 and 2");
        assert_eq!(tree.edge_count(), 2, "2 edges");
    }

    #[test]
    fn test_steiner_tree_single_terminal() {
        let g = path_graph_4();
        let tree = steiner_tree(&g, &[2]).expect("single terminal");
        assert_eq!(tree.node_count(), 1);
        assert_eq!(tree.edge_count(), 0);
    }

    #[test]
    fn test_steiner_tree_empty_terminals() {
        let g = path_graph_4();
        let tree = steiner_tree(&g, &[]).expect("empty terminals ok");
        assert_eq!(tree.node_count(), 0);
    }

    #[test]
    fn test_steiner_tree_all_terminals() {
        // All nodes are terminals: MST = min spanning tree
        let mut g = Graph::new(GraphConfig::simple());
        for _ in 0..4 { g.add_node(); }
        g.add_edge(0, 1, Some(1.0));
        g.add_edge(1, 2, Some(1.0));
        g.add_edge(2, 3, Some(1.0));
        g.add_edge(0, 3, Some(10.0));
        let tree = steiner_tree(&g, &[0, 1, 2, 3]).expect("all terminals");
        // MST has n-1 = 3 edges
        assert_eq!(tree.edge_count(), 3, "MST of 4 nodes has 3 edges");
    }

    #[test]
    fn test_steiner_tree_disconnected_returns_none() {
        let mut g = Graph::new(GraphConfig::simple());
        g.add_node(); g.add_node(); // no edges
        let tree = steiner_tree(&g, &[0, 1]);
        assert!(tree.is_none(), "disconnected terminals → None");
    }
}
