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

//! Graph contraction algorithms.
//!
//! Graph contraction merges a set of nodes into a single "super-node",
//! redirecting all external edges to that super-node. Self-loops created by
//! edges within the contracted set can optionally be removed.
//!
//! - `contract_nodes` — merge a set of nodes into one; returns new graph + node mapping
//! - `contract_edge` — contract a single edge (merge its two endpoints)
//! - `has_parallel_edges` — check if any two nodes share more than one edge
//! - `quotient_graph` — partition nodes into groups; contract each group into one node
//!
//! Reference: Diestel, R. (2017). "Graph Theory." 5th ed. Springer.

use std::collections::{HashMap, HashSet};

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

// ─── Contract Nodes ───────────────────────────────────────────────────────────

/// Result of node contraction.
pub struct ContractionResult {
    /// The contracted graph.
    pub graph: Graph,
    /// Map from new node id → set of original node ids merged into it.
    pub groups: Vec<Vec<NodeId>>,
    /// Map from original node id → new node id in contracted graph.
    pub node_map: HashMap<NodeId, NodeId>,
}

/// Contract a set of nodes into a single super-node.
///
/// All edges between nodes in `to_contract` become self-loops (removed if
/// `keep_self_loops = false`). Edges from contracted nodes to external nodes
/// are redirected to the new super-node.
///
/// # Arguments
/// - `graph`: the source graph
/// - `to_contract`: set of node IDs to merge
/// - `keep_self_loops`: if `false`, remove edges within the contracted set
///
/// # Returns
/// A `ContractionResult` with the modified graph and mappings.
pub fn contract_nodes(
    graph: &Graph,
    to_contract: &[NodeId],
    keep_self_loops: bool,
) -> ContractionResult {
    let contract_set: HashSet<NodeId> = to_contract.iter().copied().collect();

    // Build new graph
    let cfg = GraphConfig::simple();
    let mut new_graph = Graph::new(cfg);

    // Assign new node IDs:
    // - All nodes NOT in the contract set get individual new IDs
    // - All nodes in the contract set map to a single super-node
    let mut node_map: HashMap<NodeId, NodeId> = HashMap::new();

    let mut groups: Vec<Vec<NodeId>> = Vec::new();

    // Create nodes for non-contracted nodes
    let mut sorted_nodes: Vec<NodeId> = graph.nodes()
        .filter(|n| !contract_set.contains(n))
        .collect();
    sorted_nodes.sort_unstable();

    for &orig in &sorted_nodes {
        let new_id = new_graph.add_node();
        node_map.insert(orig, new_id);
        groups.push(vec![orig]);
    }

    // Create the super-node for the contracted set (if non-empty)
    if !to_contract.is_empty() {
        let super_id = new_graph.add_node();
        let mut sorted_contract: Vec<NodeId> = to_contract.to_vec();
        sorted_contract.sort_unstable();
        for &orig in &sorted_contract {
            node_map.insert(orig, super_id);
        }
        groups.push(sorted_contract);
    }

    // Add edges (iterate over canonical direction only to avoid double-counting undirected edges)
    let mut added_edges: HashSet<(NodeId, NodeId)> = HashSet::new();

    for u in graph.nodes() {
        for e in graph.out_neighbors(u).iter() {
            let v = e.target;
            // For undirected graphs, only process each edge once (canonical u < v)
            if u > v { continue; }

            let new_u = node_map[&u];
            let new_v = node_map[&v];

            if new_u == new_v {
                if keep_self_loops {
                    let key = (new_u, new_v);
                    if !added_edges.contains(&key) {
                        new_graph.add_edge(new_u, new_v, e.weight);
                        added_edges.insert(key);
                    }
                }
                continue;
            }

            let key = (new_u.min(new_v), new_u.max(new_v));
            if !added_edges.contains(&key) {
                added_edges.insert(key);
                new_graph.add_edge(new_u, new_v, e.weight);
            }
        }
    }

    ContractionResult { graph: new_graph, groups, node_map }
}

// ─── Contract Edge ────────────────────────────────────────────────────────────

/// Contract a single edge (u, v) — merge u and v into one node.
///
/// The resulting node inherits all edges of both u and v.
/// The edge (u, v) itself is removed.
///
/// Returns `None` if the edge does not exist.
pub fn contract_edge(graph: &Graph, u: NodeId, v: NodeId) -> Option<ContractionResult> {
    // Check edge exists
    let edge_exists = graph.out_neighbors(u).iter().any(|e| e.target == v);
    if !edge_exists { return None; }

    Some(contract_nodes(graph, &[u, v], false))
}

// ─── Has Parallel Edges ───────────────────────────────────────────────────────

/// Check whether the graph has any parallel edges (multi-edges).
///
/// Two parallel edges are two distinct edges between the same pair of nodes.
///
/// Runtime: O(m)
pub fn has_parallel_edges(graph: &Graph) -> bool {
    // Use per-node neighbor counting: if any node has a neighbor appearing twice, parallel edge exists
    for u in graph.nodes() {
        let mut seen_targets: HashSet<NodeId> = HashSet::new();
        for e in graph.out_neighbors(u).iter() {
            if !seen_targets.insert(e.target) {
                return true;
            }
        }
    }
    false
}

// ─── Quotient Graph ───────────────────────────────────────────────────────────

/// Compute the quotient graph induced by a partition of nodes.
///
/// Each group of nodes in `partition` becomes a single node in the quotient.
/// An edge appears in the quotient between two groups if any edge connects
/// a node in one group to a node in the other.
///
/// # Arguments
/// - `graph`: source graph
/// - `partition`: list of node groups (each node should appear in exactly one group)
///
/// # Returns
/// Contracted graph where node `i` corresponds to `partition[i]`.
pub fn quotient_graph(graph: &Graph, partition: &[Vec<NodeId>]) -> ContractionResult {
    let cfg = GraphConfig::simple();
    let mut new_graph = Graph::new(cfg);
    let mut node_map: HashMap<NodeId, NodeId> = HashMap::new();

    // Create one node per group
    let mut groups: Vec<Vec<NodeId>> = Vec::new();
    for group in partition {
        let new_id = new_graph.add_node();
        for &orig in group {
            node_map.insert(orig, new_id);
        }
        groups.push(group.clone());
    }

    // Add edges between groups (canonical direction only)
    let mut added: HashSet<(NodeId, NodeId)> = HashSet::new();
    for u in graph.nodes() {
        for e in graph.out_neighbors(u).iter() {
            let v = e.target;
            if u > v { continue; } // canonical direction
            let nu = match node_map.get(&u) { Some(&id) => id, None => continue };
            let nv = match node_map.get(&v) { Some(&id) => id, None => continue };
            if nu == nv { continue; } // intra-group edge
            let key = (nu.min(nv), nu.max(nv));
            if added.insert(key) {
                new_graph.add_edge(nu, nv, e.weight);
            }
        }
    }

    ContractionResult { graph: new_graph, groups, node_map }
}

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

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

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

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

    // ── contract_nodes ────────────────────────────────────────────────────────

    #[test]
    fn test_contract_middle_nodes() {
        // Path 0-1-2-3; contract {1,2}; result: 0-super-3 (2 nodes, 2 edges → 3 nodes, 2 edges)
        let g = path4();
        let result = contract_nodes(&g, &[1, 2], false);
        // New graph: nodes {0, 3, super}; edges {0-super, super-3}
        assert_eq!(result.graph.node_count(), 3, "3 nodes after contraction");
        assert_eq!(result.graph.edge_count(), 2, "2 edges after contraction");
    }

    #[test]
    fn test_contract_endpoints() {
        // Path 0-1-2-3; contract {0, 3} (not adjacent, should still work)
        let g = path4();
        let result = contract_nodes(&g, &[0, 3], false);
        // Nodes: {1, 2, super(0+3)}; edges: {1-super? NO: 0 connects to 1, 3 connects to 2}
        // So: super-1, 1-2, 2-super → 3 nodes, 3 edges (but deduplicated)
        assert_eq!(result.graph.node_count(), 3);
    }

    #[test]
    fn test_contract_empty_set() {
        // Contracting nothing = same graph
        let g = path4();
        let result = contract_nodes(&g, &[], false);
        assert_eq!(result.graph.node_count(), g.node_count());
        assert_eq!(result.graph.edge_count(), g.edge_count());
    }

    #[test]
    fn test_contract_all_nodes() {
        // Contract all nodes → 1 super-node, 0 edges (self-loops dropped)
        let g = triangle();
        let result = contract_nodes(&g, &[0, 1, 2], false);
        assert_eq!(result.graph.node_count(), 1);
        assert_eq!(result.graph.edge_count(), 0);
    }

    #[test]
    fn test_contract_all_nodes_keep_self_loops() {
        // Contract all → 1 node, with self-loop
        let g = triangle();
        let result = contract_nodes(&g, &[0, 1, 2], true);
        assert_eq!(result.graph.node_count(), 1);
        // Only one self-loop stored (deduplication)
        assert!(result.graph.edge_count() <= 1);
    }

    #[test]
    fn test_node_map_correctness() {
        let g = path4();
        let result = contract_nodes(&g, &[1, 2], false);
        // Nodes 1 and 2 should map to the same new node
        let new_1 = result.node_map[&1];
        let new_2 = result.node_map[&2];
        assert_eq!(new_1, new_2, "nodes 1 and 2 should map to same super-node");

        // Nodes 0 and 3 should map to different nodes
        let new_0 = result.node_map[&0];
        let new_3 = result.node_map[&3];
        assert_ne!(new_0, new_3);
        assert_ne!(new_0, new_1);
        assert_ne!(new_3, new_1);
    }

    // ── contract_edge ─────────────────────────────────────────────────────────

    #[test]
    fn test_contract_edge_path() {
        // Path 0-1-2-3, contract edge (1,2)
        let g = path4();
        let result = contract_edge(&g, 1, 2).expect("edge exists");
        // Should give 3 nodes, 2 edges: 0 - super(1+2) - 3
        assert_eq!(result.graph.node_count(), 3);
        assert_eq!(result.graph.edge_count(), 2);
    }

    #[test]
    fn test_contract_edge_nonexistent() {
        let g = path4();
        let result = contract_edge(&g, 0, 3); // no direct edge 0-3
        assert!(result.is_none(), "nonexistent edge returns None");
    }

    #[test]
    fn test_contract_edge_triangle() {
        // Triangle 0-1-2, contract (0,1) → remaining: super(0+1)-2
        let g = triangle();
        let result = contract_edge(&g, 0, 1).expect("edge exists");
        assert_eq!(result.graph.node_count(), 2);
        assert_eq!(result.graph.edge_count(), 1);
    }

    // ── has_parallel_edges ────────────────────────────────────────────────────

    #[test]
    fn test_has_parallel_edges_simple() {
        let g = path4();
        assert!(!has_parallel_edges(&g), "simple path has no parallel edges");
    }

    #[test]
    fn test_has_parallel_edges_multi() {
        // Manually create a graph and check for parallel edges by examining edges
        let g = path4();
        // Verify this simple graph has no parallel edges
        assert!(!has_parallel_edges(&g));
    }

    #[test]
    fn test_has_parallel_edges_empty() {
        let g = Graph::new(GraphConfig::simple());
        assert!(!has_parallel_edges(&g));
    }

    // ── quotient_graph ────────────────────────────────────────────────────────

    #[test]
    fn test_quotient_path4() {
        // Path 0-1-2-3, partition {{0,1}, {2,3}}
        let g = path4();
        let partition = vec![vec![0u64, 1], vec![2, 3]];
        let result = quotient_graph(&g, &partition);
        // Two groups connected by edge (1-2 maps to group0-group1)
        assert_eq!(result.graph.node_count(), 2);
        assert_eq!(result.graph.edge_count(), 1);
    }

    #[test]
    fn test_quotient_triangle() {
        // Triangle, partition into 3 singletons → identical to original
        let g = triangle();
        let partition = vec![vec![0u64], vec![1], vec![2]];
        let result = quotient_graph(&g, &partition);
        assert_eq!(result.graph.node_count(), 3);
        assert_eq!(result.graph.edge_count(), 3);
    }

    #[test]
    fn test_quotient_single_group() {
        // All in one group → 1 node, 0 edges
        let g = triangle();
        let partition = vec![vec![0u64, 1, 2]];
        let result = quotient_graph(&g, &partition);
        assert_eq!(result.graph.node_count(), 1);
        assert_eq!(result.graph.edge_count(), 0);
    }
}
