//! Spanning tree algorithms.
//!
//! Ported from classical algorithm sources, adapted to NetworKit's Graph structure.
//!
//! Algorithms:
//! - Kruskal's MST: Minimum spanning tree via union-find
//! - Prim's MST: Minimum spanning tree via greedy neighbor expansion
//! - Maximum spanning tree (reverse Kruskal's)

use super::super::graph::{Graph, NodeId};
use std::collections::{BinaryHeap, HashMap};
use std::cmp::Ordering;

// ─────────────────────────────────────────────────────────────────────────────
// Union-Find for Kruskal's

struct UnionFind {
    parent: Vec<usize>,
    rank: Vec<usize>,
}

impl UnionFind {
    fn new(n: usize) -> Self {
        UnionFind {
            parent: (0..n).collect(),
            rank: vec![0; n],
        }
    }

    fn find(&mut self, x: usize) -> usize {
        if self.parent[x] != x {
            self.parent[x] = self.find(self.parent[x]); // Path compression
        }
        self.parent[x]
    }

    fn union(&mut self, x: usize, y: usize) -> bool {
        let rx = self.find(x);
        let ry = self.find(y);
        if rx == ry {
            return false; // Already in the same component
        }
        // Union by rank
        match self.rank[rx].cmp(&self.rank[ry]) {
            Ordering::Less => self.parent[rx] = ry,
            Ordering::Greater => self.parent[ry] = rx,
            Ordering::Equal => {
                self.parent[ry] = rx;
                self.rank[rx] += 1;
            }
        }
        true
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Kruskal's MST

/// Minimum spanning tree using Kruskal's algorithm.
///
/// Finds the minimum spanning tree (or forest for disconnected graphs) by
/// sorting edges by weight and adding them if they connect different components.
///
/// Ported from classical Kruskal's algorithm, adapted to NetworKit.
/// Reference: Kruskal, J.B. (1956). "On the shortest spanning subtree of a graph"
///
/// Time complexity: O(E log E)
///
/// # Returns
/// List of `(source, target, weight)` tuples forming the minimum spanning tree.
///
/// # Example
/// ```ignore
/// let mst = minimum_spanning_tree_kruskal(&graph);
/// let total_weight: f64 = mst.iter().map(|(_, _, w)| w).sum();
/// ```
pub fn minimum_spanning_tree_kruskal(graph: &Graph) -> Vec<(NodeId, NodeId, f64)> {
    spanning_tree_kruskal(graph, false)
}

/// Maximum spanning tree using Kruskal's algorithm (reverse order).
///
/// Finds the maximum spanning tree by sorting edges in descending order.
///
/// # Returns
/// List of `(source, target, weight)` tuples forming the maximum spanning tree.
pub fn maximum_spanning_tree_kruskal(graph: &Graph) -> Vec<(NodeId, NodeId, f64)> {
    spanning_tree_kruskal(graph, true)
}

fn spanning_tree_kruskal(graph: &Graph, maximize: bool) -> Vec<(NodeId, NodeId, f64)> {
    // Collect all edges
    let mut edges: Vec<(f64, NodeId, NodeId)> = graph
        .edges()
        .map(|(u, v, w, _)| (w.unwrap_or(1.0), u, v))
        .collect();

    // Sort by weight (ascending for MST, descending for max ST)
    if maximize {
        edges.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(Ordering::Equal));
    } else {
        edges.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(Ordering::Equal));
    }

    // Map node IDs to indices for union-find
    let nodes: Vec<NodeId> = graph.nodes().collect();
    let node_to_idx: HashMap<NodeId, usize> = nodes.iter().enumerate().map(|(i, &id)| (id, i)).collect();
    let mut uf = UnionFind::new(nodes.len());

    let mut result = Vec::new();

    for (weight, u, v) in edges {
        let iu = match node_to_idx.get(&u) {
            Some(&i) => i,
            None => continue,
        };
        let iv = match node_to_idx.get(&v) {
            Some(&i) => i,
            None => continue,
        };

        // Add edge if it connects two different components
        if uf.union(iu, iv) {
            result.push((u, v, weight));
        }
    }

    result
}

// ─────────────────────────────────────────────────────────────────────────────
// Prim's MST

/// State for Prim's priority queue.
#[derive(Copy, Clone)]
struct PrimState {
    weight: f64,
    node: NodeId,
}

impl Ord for PrimState {
    fn cmp(&self, other: &Self) -> Ordering {
        other.weight.partial_cmp(&self.weight).unwrap_or(Ordering::Equal)
    }
}

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

impl PartialEq for PrimState {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}

impl Eq for PrimState {}

/// Minimum spanning tree using Prim's algorithm.
///
/// Grows the MST greedily from a starting node, always choosing the minimum
/// weight edge that connects a new node to the current tree.
///
/// Better than Kruskal's for dense graphs.
///
/// Reference: Prim, R.C. (1957). "Shortest connection networks and some generalizations"
///
/// Time complexity: O((V + E) log V)
///
/// # Returns
/// List of `(source, target, weight)` tuples forming the minimum spanning tree.
///
/// # Example
/// ```ignore
/// let mst = minimum_spanning_tree_prim(&graph);
/// ```
pub fn minimum_spanning_tree_prim(graph: &Graph) -> Vec<(NodeId, NodeId, f64)> {
    let mut in_mst: HashMap<NodeId, bool> = HashMap::new();
    let mut min_weight: HashMap<NodeId, f64> = HashMap::new();
    let mut parent: HashMap<NodeId, NodeId> = HashMap::new();
    let mut heap = BinaryHeap::new();

    // Start from the first node
    let start = match graph.nodes().next() {
        Some(n) => n,
        None => return vec![],
    };

    min_weight.insert(start, 0.0);
    heap.push(PrimState { weight: 0.0, node: start });

    let mut result = Vec::new();

    while let Some(PrimState { weight, node }) = heap.pop() {
        if in_mst.contains_key(&node) {
            continue;
        }
        in_mst.insert(node, true);

        // Add edge to MST (skip the root)
        if let Some(&par) = parent.get(&node) {
            result.push((par, node, weight));
        }

        // Explore neighbors
        for neighbor in graph.out_neighbors(node) {
            let edge_weight = neighbor.weight.unwrap_or(1.0);
            let curr_min = min_weight.get(&neighbor.target).copied().unwrap_or(f64::INFINITY);

            if edge_weight < curr_min {
                min_weight.insert(neighbor.target, edge_weight);
                parent.insert(neighbor.target, node);
                heap.push(PrimState { weight: edge_weight, node: neighbor.target });
            }
        }
    }

    result
}

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

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

    fn make_undirected_weighted() -> (Graph, Vec<NodeId>) {
        use super::super::super::graph::GraphConfig;
        let mut g = Graph::new(GraphConfig::simple());
        let nodes: Vec<NodeId> = (0..5).map(|_| g.add_node()).collect();
        // Classic MST test case
        //   0 -2- 1 -3- 2
        //   |  \      /
        //   6   5   4
        //   |     \  |
        //   4 --3-- 3
        g.add_edge(nodes[0], nodes[1], Some(2.0));
        g.add_edge(nodes[0], nodes[3], Some(6.0));
        g.add_edge(nodes[1], nodes[2], Some(3.0));
        g.add_edge(nodes[1], nodes[3], Some(8.0));
        g.add_edge(nodes[1], nodes[4], Some(5.0));
        g.add_edge(nodes[2], nodes[4], Some(7.0));
        g.add_edge(nodes[3], nodes[4], Some(9.0));
        (g, nodes)
    }

    #[test]
    fn test_kruskal_mst() {
        let (g, _nodes) = make_undirected_weighted();
        let mst = minimum_spanning_tree_kruskal(&g);

        // MST should have V-1 = 4 edges
        assert_eq!(mst.len(), 4);

        // Total weight should be minimum possible
        let total: f64 = mst.iter().map(|(_, _, w)| w).sum();
        // Optimal MST: edges 2, 3, 5, 6 = total 16
        assert_eq!(total, 16.0);
    }

    #[test]
    fn test_prim_mst() {
        let (g, _nodes) = make_undirected_weighted();
        let mst = minimum_spanning_tree_prim(&g);

        assert_eq!(mst.len(), 4);

        let total: f64 = mst.iter().map(|(_, _, w)| w).sum();
        assert_eq!(total, 16.0);
    }

    #[test]
    fn test_maximum_spanning_tree() {
        let (g, _nodes) = make_undirected_weighted();
        let max_st = maximum_spanning_tree_kruskal(&g);

        assert_eq!(max_st.len(), 4);

        let total: f64 = max_st.iter().map(|(_, _, w)| w).sum();
        // Max ST picks highest weight edges: 9, 8, 7, 6 = 30
        assert_eq!(total, 30.0);
    }

    #[test]
    fn test_empty_graph() {
        use super::super::super::graph::GraphConfig;
        let g = Graph::new(GraphConfig::simple());
        let mst = minimum_spanning_tree_kruskal(&g);
        assert!(mst.is_empty());
    }

    #[test]
    fn test_single_node() {
        use super::super::super::graph::GraphConfig;
        let mut g = Graph::new(GraphConfig::simple());
        g.add_node();
        let mst = minimum_spanning_tree_kruskal(&g);
        assert!(mst.is_empty()); // No edges to form MST
    }
}
