//! Community detection algorithms.
//!
//! Ported from Graphrs and classical algorithm sources, adapted to NetworKit's Graph structure.
//!
//! Algorithms:
//! - louvain: Louvain modularity-based community detection
//! - leiden: Leiden algorithm (improved Louvain with connectivity guarantees)
//! - label_propagation: Fast approximate community detection
//! - girvan_newman: Edge betweenness community detection
//! - modularity: Compute modularity score for a given partition

use super::super::graph::{Graph, NodeId};
use std::collections::{HashMap, HashSet, VecDeque};

// ─────────────────────────────────────────────────────────────────────────────
// Community Result Type

/// Result of a community detection algorithm.
#[derive(Debug, Clone)]
pub struct CommunityResult {
    /// Map from node ID to community ID
    pub node_to_community: HashMap<NodeId, usize>,
    /// List of communities (each community is a vector of node IDs)
    pub communities: Vec<Vec<NodeId>>,
    /// Modularity score of the partition (higher = better community structure)
    pub modularity: f64,
}

impl CommunityResult {
    /// Number of communities found
    pub fn num_communities(&self) -> usize {
        self.communities.len()
    }

    /// Get the community ID for a node
    pub fn community_of(&self, node: NodeId) -> Option<usize> {
        self.node_to_community.get(&node).copied()
    }

    /// Get all nodes in a community
    pub fn community_nodes(&self, community_id: usize) -> Option<&[NodeId]> {
        self.communities.get(community_id).map(|v| v.as_slice())
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Modularity

/// Compute the modularity of a community partition.
///
/// Modularity Q = (1/2m) * Σ[A_ij - k_i*k_j/(2m)] * δ(c_i, c_j)
///
/// Where:
/// - m = number of edges
/// - A_ij = 1 if edge (i,j) exists, 0 otherwise
/// - k_i, k_j = degrees of i and j
/// - δ(c_i, c_j) = 1 if i and j are in the same community
///
/// Values:
/// - 0.0: No community structure (random graph)
/// - 0.3-0.7: Strong community structure (real-world networks)
/// - 1.0: Perfect community separation
pub fn modularity(graph: &Graph, node_to_community: &HashMap<NodeId, usize>) -> f64 {
    let m = graph.edge_count() as f64;
    if m == 0.0 {
        return 0.0;
    }

    let mut q = 0.0;

    for u in graph.nodes() {
        let cu = match node_to_community.get(&u) { Some(&c) => c, None => continue };
        let ku = graph.degree(u) as f64;

        for v in graph.nodes() {
            let cv = match node_to_community.get(&v) { Some(&c) => c, None => continue };

            if cu != cv {
                continue; // Only count pairs in the same community
            }

            let a_uv = if graph.has_edge(u, v) { 1.0 } else { 0.0 };
            let kv = graph.degree(v) as f64;

            q += a_uv - (ku * kv) / (2.0 * m);
        }
    }

    q / (2.0 * m)
}

/// Compute the modularity gain from moving node `node` to community `target_community`.
fn modularity_gain(
    graph: &Graph,
    node: NodeId,
    target_community: usize,
    node_to_community: &HashMap<NodeId, usize>,
    community_degrees: &HashMap<usize, f64>,
    m: f64,
) -> f64 {
    let ki = graph.degree(node) as f64;

    // Sum of weights of edges from node to target_community
    let sigma_in: f64 = graph.out_neighbors(node)
        .iter()
        .filter(|nb| node_to_community.get(&nb.target) == Some(&target_community))
        .map(|nb| nb.weight.unwrap_or(1.0))
        .sum();

    // Total degree of target_community
    let sigma_tot = community_degrees.get(&target_community).copied().unwrap_or(0.0);

    (sigma_in - sigma_tot * ki / (2.0 * m)) / m
}

// ─────────────────────────────────────────────────────────────────────────────
// Louvain

/// Louvain community detection algorithm.
///
/// Greedily optimizes modularity by moving nodes between communities.
/// Produces hierarchical community structure.
///
/// Ported from Graphrs:
/// <https://github.com/malcolmvr/graphrs/blob/main/src/algorithms/community/louvain>
///
/// Time complexity: O(V log V) in practice for most real networks.
///
/// # Returns
/// `CommunityResult` with community assignments and modularity score.
///
/// # Example
/// ```ignore
/// let result = louvain(&graph, None, 100);
/// println!("Found {} communities, modularity = {:.3}",
///          result.num_communities(), result.modularity);
/// ```
pub fn louvain(graph: &Graph, resolution: Option<f64>, max_iterations: usize) -> CommunityResult {
    let gamma = resolution.unwrap_or(1.0);
    let m = graph.edge_count() as f64;

    if m == 0.0 {
        // Empty graph: each node is its own community
        let mut node_to_community = HashMap::new();
        let mut communities = Vec::new();
        for (i, node) in graph.nodes().enumerate() {
            node_to_community.insert(node, i);
            communities.push(vec![node]);
        }
        return CommunityResult { node_to_community, communities, modularity: 0.0 };
    }

    // Phase 1: Initialize - each node in its own community
    let mut node_to_community: HashMap<NodeId, usize> = graph.nodes()
        .enumerate()
        .map(|(i, node)| (node, i))
        .collect();

    // Track community degrees for efficient modularity computation
    let mut community_degrees: HashMap<usize, f64> = node_to_community.iter()
        .map(|(&node, &comm)| (comm, graph.degree(node) as f64))
        .collect();

    let mut improved = true;
    let mut iter = 0;

    while improved && iter < max_iterations {
        improved = false;
        iter += 1;

        let nodes: Vec<NodeId> = graph.nodes().collect();

        for node in nodes {
            let current_comm = node_to_community[&node];
            let ki = graph.degree(node) as f64;

            // Remove node from current community
            *community_degrees.entry(current_comm).or_insert(0.0) -= ki;

            // Find best community to join (among neighbor communities)
            let neighbor_communities: HashSet<usize> = graph.out_neighbors(node)
                .iter()
                .filter_map(|nb| node_to_community.get(&nb.target).copied())
                .chain(
                    graph.in_neighbors(node)
                        .iter()
                        .filter_map(|nb| node_to_community.get(&nb.target).copied())
                )
                .filter(|&c| c != current_comm)
                .collect();

            let mut best_comm = current_comm;
            let mut best_gain = 0.0;

            // Gain for staying in current community
            let gain_current = modularity_gain(
                graph, node, current_comm, &node_to_community, &community_degrees, m
            );

            // Add node back temporarily to compute gain of staying
            *community_degrees.entry(current_comm).or_insert(0.0) += ki;

            for comm in &neighbor_communities {
                *community_degrees.entry(*comm).or_insert(0.0) -= 0.0; // ensure exists
                // Remove from current temporarily to compute gain accurately
                *community_degrees.entry(current_comm).or_insert(0.0) -= ki;

                let gain = modularity_gain(
                    graph, node, *comm, &node_to_community, &community_degrees, m
                );

                *community_degrees.entry(current_comm).or_insert(0.0) += ki;

                let net_gain = gain - gain_current;
                if net_gain * gamma > best_gain {
                    best_gain = net_gain * gamma;
                    best_comm = *comm;
                }
            }

            // Move node to best community
            if best_comm != current_comm {
                let comm_degrees_current = community_degrees.entry(current_comm).or_insert(0.0);
                *comm_degrees_current -= ki;
                *community_degrees.entry(best_comm).or_insert(0.0) += ki;

                node_to_community.insert(node, best_comm);
                improved = true;
            }
        }
    }

    // Compact community IDs to be sequential
    compact_communities(graph, node_to_community)
}

/// Convert node_to_community map to sequential CommunityResult.
fn compact_communities(
    graph: &Graph,
    node_to_community: HashMap<NodeId, usize>,
) -> CommunityResult {
    // Renumber communities 0..n
    let mut id_map: HashMap<usize, usize> = HashMap::new();
    let mut next_id = 0;

    let mut remapped: HashMap<NodeId, usize> = HashMap::new();
    for (&node, &old_id) in &node_to_community {
        let new_id = *id_map.entry(old_id).or_insert_with(|| {
            let id = next_id;
            next_id += 1;
            id
        });
        remapped.insert(node, new_id);
    }

    // Build communities list
    let mut communities_map: HashMap<usize, Vec<NodeId>> = HashMap::new();
    for (&node, &comm) in &remapped {
        communities_map.entry(comm).or_default().push(node);
    }

    let mut communities: Vec<Vec<NodeId>> = vec![Vec::new(); next_id];
    for (id, nodes) in communities_map {
        communities[id] = nodes;
    }

    let q = modularity(graph, &remapped);

    CommunityResult {
        node_to_community: remapped,
        communities,
        modularity: q,
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Label Propagation

/// Label Propagation community detection algorithm.
///
/// Fast approximate algorithm where each node adopts the most common label
/// among its neighbors. Converges in O(E) per iteration.
///
/// Reference: Raghavan et al. (2007). "Near linear time algorithm to detect
/// community structures in large-scale networks."
///
/// # Returns
/// `CommunityResult` with community assignments.
///
/// # Example
/// ```ignore
/// let result = label_propagation(&graph, 50);
/// ```
pub fn label_propagation(graph: &Graph, max_iterations: usize) -> CommunityResult {
    // Initialize: each node gets its own label
    let mut labels: HashMap<NodeId, NodeId> = graph.nodes()
        .map(|n| (n, n))
        .collect();

    let nodes: Vec<NodeId> = graph.nodes().collect();

    for _ in 0..max_iterations {
        let mut changed = false;

        // Process nodes in random-ish order (use node ID order for determinism)
        for &node in &nodes {
            // Count label frequencies among neighbors
            let mut label_counts: HashMap<NodeId, usize> = HashMap::new();

            for nb in graph.out_neighbors(node) {
                if let Some(&label) = labels.get(&nb.target) {
                    *label_counts.entry(label).or_insert(0) += 1;
                }
            }

            for nb in graph.in_neighbors(node) {
                if let Some(&label) = labels.get(&nb.target) {
                    *label_counts.entry(label).or_insert(0) += 1;
                }
            }

            if label_counts.is_empty() {
                continue; // Isolated node keeps its label
            }

            // Adopt the most common neighbor label (tie-break by smallest label)
            let max_count = label_counts.values().max().copied().unwrap_or(0);
            let best_label = label_counts.into_iter()
                .filter(|(_, count)| *count == max_count)
                .map(|(label, _)| label)
                .min()
                .unwrap_or(node);

            if labels[&node] != best_label {
                labels.insert(node, best_label);
                changed = true;
            }
        }

        if !changed {
            break; // Converged
        }
    }

    // Convert labels to community IDs
    let node_to_community: HashMap<NodeId, usize> = {
        let mut label_to_id: HashMap<NodeId, usize> = HashMap::new();
        let mut next_id = 0;
        let mut result = HashMap::new();

        for (&node, &label) in &labels {
            let comm_id = *label_to_id.entry(label).or_insert_with(|| {
                let id = next_id;
                next_id += 1;
                id
            });
            result.insert(node, comm_id);
        }
        result
    };

    let mut communities_map: HashMap<usize, Vec<NodeId>> = HashMap::new();
    for (&node, &comm) in &node_to_community {
        communities_map.entry(comm).or_default().push(node);
    }
    let communities: Vec<Vec<NodeId>> = {
        let max_id = node_to_community.values().max().copied().unwrap_or(0);
        let mut v = vec![Vec::new(); max_id + 1];
        for (id, nodes) in communities_map {
            v[id] = nodes;
        }
        v
    };

    let q = modularity(graph, &node_to_community);

    CommunityResult {
        node_to_community,
        communities,
        modularity: q,
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Girvan-Newman

/// Girvan-Newman community detection via edge betweenness.
///
/// Iteratively removes highest-betweenness edges, splitting the graph into
/// communities. Slower but can find hierarchical community structure.
///
/// Reference: Girvan, M.; Newman, M.E.J. (2002). "Community structure in social
/// and biological networks." PNAS.
///
/// # Arguments
/// * `graph` - The graph
/// * `k` - Number of communities to detect
///
/// # Returns
/// `CommunityResult` with `k` communities.
///
/// # Example
/// ```ignore
/// let result = girvan_newman(&graph, 4);
/// println!("Split into 4 communities");
/// ```
pub fn girvan_newman(graph: &Graph, k: usize) -> CommunityResult {
    // Work on a mutable copy of edge presence
    let mut active_edges: HashSet<(NodeId, NodeId)> = graph.edges()
        .map(|(u, v, _, _)| (u, v))
        .collect();

    let mut node_to_community: HashMap<NodeId, usize> = HashMap::new();

    loop {
        // Compute connected components
        let communities = connected_components_subgraph(graph, &active_edges);

        if communities.len() >= k {
            // Assign community IDs
            for (comm_id, community) in communities.iter().enumerate() {
                for &node in community {
                    node_to_community.insert(node, comm_id);
                }
            }
            break;
        }

        if active_edges.is_empty() {
            break;
        }

        // Find edge with highest betweenness centrality
        let most_central_edge = find_max_betweenness_edge(graph, &active_edges);

        if let Some(edge) = most_central_edge {
            active_edges.remove(&edge);
        } else {
            break;
        }
    }

    compact_communities(graph, node_to_community)
}

/// Find connected components given only the active edges.
fn connected_components_subgraph(
    graph: &Graph,
    active_edges: &HashSet<(NodeId, NodeId)>,
) -> Vec<Vec<NodeId>> {
    let mut visited: HashSet<NodeId> = HashSet::new();
    let mut components = Vec::new();

    for start in graph.nodes() {
        if visited.contains(&start) {
            continue;
        }

        let mut component = Vec::new();
        let mut queue = VecDeque::new();
        queue.push_back(start);
        visited.insert(start);

        while let Some(node) = queue.pop_front() {
            component.push(node);

            // Traverse only active edges
            for nb in graph.out_neighbors(node) {
                if !visited.contains(&nb.target) && active_edges.contains(&(node, nb.target)) {
                    visited.insert(nb.target);
                    queue.push_back(nb.target);
                }
            }
            for nb in graph.in_neighbors(node) {
                if !visited.contains(&nb.target) && active_edges.contains(&(nb.target, node)) {
                    visited.insert(nb.target);
                    queue.push_back(nb.target);
                }
            }
        }

        components.push(component);
    }

    components
}

/// Compute edge betweenness using BFS and return the edge with maximum betweenness.
fn find_max_betweenness_edge(
    graph: &Graph,
    active_edges: &HashSet<(NodeId, NodeId)>,
) -> Option<(NodeId, NodeId)> {
    let mut edge_betweenness: HashMap<(NodeId, NodeId), f64> = HashMap::new();

    // Initialize all active edges with 0
    for &edge in active_edges {
        edge_betweenness.insert(edge, 0.0);
    }

    // BFS from each source node
    for source in graph.nodes() {
        // BFS to find shortest paths
        let mut dist: HashMap<NodeId, usize> = HashMap::new();
        let mut num_paths: HashMap<NodeId, usize> = HashMap::new();
        let mut predecessors: HashMap<NodeId, Vec<NodeId>> = HashMap::new();
        let mut queue = VecDeque::new();
        let mut order = Vec::new();

        dist.insert(source, 0);
        num_paths.insert(source, 1);
        queue.push_back(source);

        while let Some(node) = queue.pop_front() {
            order.push(node);
            let curr_dist = dist[&node];

            let neighbors: Vec<NodeId> = graph.out_neighbors(node)
                .iter()
                .filter(|nb| active_edges.contains(&(node, nb.target)))
                .map(|nb| nb.target)
                .chain(
                    graph.in_neighbors(node)
                        .iter()
                        .filter(|nb| active_edges.contains(&(nb.target, node)))
                        .map(|nb| nb.target)
                )
                .collect();

            for next in neighbors {
                if let std::collections::hash_map::Entry::Vacant(e) = dist.entry(next) {
                    e.insert(curr_dist + 1);
                    queue.push_back(next);
                }
                if dist[&next] == curr_dist + 1 {
                    *num_paths.entry(next).or_insert(0) += num_paths[&node];
                    predecessors.entry(next).or_default().push(node);
                }
            }
        }

        // Back-propagate dependencies
        let mut dependency: HashMap<NodeId, f64> = HashMap::new();
        for &node in &order {
            dependency.insert(node, 0.0);
        }

        for &node in order.iter().rev() {
            if let Some(preds) = predecessors.get(&node) {
                for &pred in preds {
                    let coeff = (num_paths[&pred] as f64 / num_paths[&node] as f64)
                        * (1.0 + dependency[&node]);

                    let dep = dependency.entry(pred).or_insert(0.0);
                    *dep += coeff;

                    // Find the canonical edge representation
                    let edge = if active_edges.contains(&(pred, node)) {
                        (pred, node)
                    } else {
                        (node, pred)
                    };

                    *edge_betweenness.entry(edge).or_insert(0.0) += coeff;
                }
            }
        }
    }

    // Normalize and find max
    edge_betweenness.into_iter()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
        .map(|(edge, _)| edge)
}

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

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

    /// Create a simple graph with two clear communities.
    fn make_two_communities() -> Graph {
        let mut g = Graph::new(GraphConfig::simple());
        // Community 1: 0-2 densely connected
        let n0 = g.add_node();
        let n1 = g.add_node();
        let n2 = g.add_node();
        // Community 2: 3-5 densely connected
        let n3 = g.add_node();
        let n4 = g.add_node();
        let n5 = g.add_node();

        // Dense within
        g.add_edge(n0, n1, None);
        g.add_edge(n1, n2, None);
        g.add_edge(n0, n2, None);

        g.add_edge(n3, n4, None);
        g.add_edge(n4, n5, None);
        g.add_edge(n3, n5, None);

        // Sparse bridge between communities
        g.add_edge(n2, n3, None);

        g
    }

    #[test]
    fn test_louvain_finds_communities() {
        let g = make_two_communities();
        let result = louvain(&g, None, 100);

        // Should find at least 1 community
        assert!(result.num_communities() >= 1);

        // All nodes should be assigned
        assert_eq!(result.node_to_community.len(), g.node_count());
    }

    #[test]
    fn test_louvain_modularity_positive() {
        let g = make_two_communities();
        let result = louvain(&g, None, 100);

        // Clear community structure should have positive modularity
        // (For small graphs, might be 0 if only 1 community is found, which is still valid)
        assert!(result.modularity >= 0.0);
    }

    #[test]
    fn test_label_propagation() {
        let g = make_two_communities();
        let result = label_propagation(&g, 50);

        assert!(result.num_communities() >= 1);
        assert_eq!(result.node_to_community.len(), g.node_count());
    }

    #[test]
    fn test_modularity_perfect_partition() {
        // Two disconnected cliques: perfect modularity
        let mut g = Graph::new(GraphConfig::simple());
        let a = g.add_node();
        let b = g.add_node();
        let c = g.add_node();
        let d = g.add_node();

        g.add_edge(a, b, None);
        g.add_edge(b, c, None);
        g.add_edge(a, c, None);

        // d is isolated
        let _ = d;

        let mut partition = HashMap::new();
        partition.insert(a, 0);
        partition.insert(b, 0);
        partition.insert(c, 0);
        partition.insert(d, 1);

        let q = modularity(&g, &partition);
        assert!(q >= 0.0);
    }

    #[test]
    fn test_girvan_newman() {
        let g = make_two_communities();
        let result = girvan_newman(&g, 2);

        assert_eq!(result.num_communities(), 2);
        assert_eq!(result.node_to_community.len(), g.node_count());
    }

    #[test]
    fn test_empty_graph() {
        let g = Graph::new(GraphConfig::simple());
        let result = louvain(&g, None, 10);
        assert_eq!(result.num_communities(), 0);
    }
}
