leiden-rs 0.6.0

//! Core Leiden algorithm implementation.

use std::collections::VecDeque;

use rand::rngs::StdRng;
use rand::seq::SliceRandom;
use rand::SeedableRng;
#[cfg(feature = "rayon")]
use rayon::prelude::*;
use rustc_hash::FxHashMap;

use crate::partition::Partition;
use crate::quality::{GraphData, Modularity, MoveComponents, QualityFunction};
use gryf::core::{
    id::IntegerIdType, marker::Undirected, GraphBase, GraphRef, Neighbors, VertexSet,
};

#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

/// Quality function selection for the Leiden algorithm.
#[derive(Debug, Clone, Copy, Default)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub enum QualityType {
    /// Newman-Girvan modularity (default).
    #[default]
    Modularity,
    /// Constant Potts Model.
    CPM,
    /// Reichardt-Bornholdt with configuration model null.
    RBConfiguration,
    /// Reichardt-Bornholdt with Erdős-Rényi null model.
    RBER,
}

/// Configuration for the Leiden algorithm.
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct LeidenConfig {
    /// Maximum number of Leiden iterations (local move → refine → aggregate).
    pub max_iterations: usize,
    /// Resolution parameter γ controlling community granularity.
    pub resolution: f64,
    /// Optional RNG seed for reproducible results.
    pub seed: Option<u64>,
    /// Quality function to optimize.
    pub quality: QualityType,
    /// Convergence threshold: stop when quality improvement is below this value.
    pub epsilon: f64,
    /// Maximum number of nodes per community (0 = unlimited).
    pub max_comm_size: usize,
}

impl Default for LeidenConfig {
    fn default() -> Self {
        Self {
            max_iterations: 100,
            resolution: 1.0,
            seed: None,
            quality: QualityType::default(),
            epsilon: 1e-10,
            max_comm_size: 0,
        }
    }
}

impl LeidenConfig {
    /// Create a builder for configuring the Leiden algorithm.
    pub fn builder() -> LeidenConfigBuilder {
        LeidenConfigBuilder::default()
    }
}

/// Builder for [`LeidenConfig`].
#[derive(Debug, Clone, Default)]
pub struct LeidenConfigBuilder {
    max_iterations: Option<usize>,
    resolution: Option<f64>,
    seed: Option<u64>,
    quality: Option<QualityType>,
    epsilon: Option<f64>,
    max_comm_size: Option<usize>,
}

impl LeidenConfigBuilder {
    /// Set the maximum number of Leiden iterations.
    pub fn max_iterations(mut self, v: usize) -> Self {
        self.max_iterations = Some(v);
        self
    }

    /// Set the resolution parameter γ.
    pub fn resolution(mut self, v: f64) -> Self {
        self.resolution = Some(v);
        self
    }

    /// Set the RNG seed for reproducible results.
    pub fn seed(mut self, v: u64) -> Self {
        self.seed = Some(v);
        self
    }

    /// Set the quality function to optimize.
    pub fn quality(mut self, v: QualityType) -> Self {
        self.quality = Some(v);
        self
    }

    /// Set the convergence threshold.
    pub fn epsilon(mut self, v: f64) -> Self {
        self.epsilon = Some(v);
        self
    }

    /// Set the maximum number of nodes per community (0 = unlimited).
    pub fn max_comm_size(mut self, v: usize) -> Self {
        self.max_comm_size = Some(v);
        self
    }

    /// Build the configuration, applying defaults for unset fields.
    pub fn build(self) -> LeidenConfig {
        LeidenConfig {
            max_iterations: self.max_iterations.unwrap_or(100),
            resolution: self.resolution.unwrap_or(1.0),
            seed: self.seed,
            quality: self.quality.unwrap_or_default(),
            epsilon: self.epsilon.unwrap_or(1e-10),
            max_comm_size: self.max_comm_size.unwrap_or(0),
        }
    }
}

/// Result of running the Leiden algorithm, containing the partition and its quality score.
#[derive(Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct LeidenOutput {
    /// The detected community partition.
    pub partition: Partition,
    /// The quality score of the partition (higher is better).
    pub quality: f64,
}

impl LeidenOutput {
    /// Create a new Leiden output with the given partition and quality score.
    pub fn new(partition: Partition, quality: f64) -> Self {
        Self { partition, quality }
    }
}

/// The Leiden community detection algorithm.
pub struct Leiden {
    config: LeidenConfig,
}

impl Leiden {
    /// Create a new Leiden instance with the given configuration.
    pub fn new(config: LeidenConfig) -> Self {
        Self { config }
    }

    /// Run the Leiden algorithm on the given graph.
    ///
    /// Returns a [`LeidenOutput`] containing the community partition and its quality score.
    pub fn run<V, G>(
        &self,
        graph: &gryf::Graph<V, f64, Undirected, G>,
    ) -> crate::error::Result<LeidenOutput>
    where
        G: GraphBase<VertexId: IntegerIdType, EdgeType = Undirected>
            + Neighbors
            + VertexSet
            + GraphRef<V, f64>,
    {
        let original_n = graph.vertex_count();
        if original_n == 0 {
            return Ok(LeidenOutput::new(Partition::new(0), 0.0));
        }

        let mut data = GraphData::from_gryf_graph(graph)?;
        let mut partition = Partition::new(data.node_count());
        let mut flat_mapping: Vec<usize> = (0..data.node_count()).collect();

        let mut rng = match self.config.seed {
            Some(seed) => StdRng::seed_from_u64(seed),
            None => StdRng::from_rng(&mut rand::rng()),
        };

        let modularity = Modularity::with_resolution(self.config.resolution);
        let cpm = crate::quality::CPM::new(self.config.resolution);
        let rbconfig = crate::quality::RBConfiguration::with_resolution(self.config.resolution);
        let rber = crate::quality::RBER::new(self.config.resolution);
        let quality: &(dyn QualityFunction + Sync) = match self.config.quality {
            QualityType::Modularity => &modularity,
            QualityType::CPM => &cpm,
            QualityType::RBConfiguration => &rbconfig,
            QualityType::RBER => &rber,
        };

        for _iter in 0..self.config.max_iterations {
            let q_before = quality.total_quality(&data, &partition);
            let changed = local_moving_dispatch(
                &data,
                &mut partition,
                quality,
                &mut rng,
                self.config.max_comm_size,
                self.config.epsilon,
            );
            if !changed {
                break;
            }
            partition.renumber();

            let q_after = quality.total_quality(&data, &partition);
            if (q_after - q_before).abs() < self.config.epsilon {
                break;
            }

            let refined = refinement(&data, &partition, quality, &mut rng, self.config.epsilon);

            let (agg_data, orig_to_agg, agg_initial_partition) =
                aggregate(&data, &refined, &partition);

            for orig_node in 0..original_n {
                flat_mapping[orig_node] = orig_to_agg[flat_mapping[orig_node]];
            }

            data = agg_data;
            partition = agg_initial_partition;

            if data.node_count() <= 1 {
                break;
            }
        }

        let mut result = Partition::from_membership(vec![0; original_n]);
        for (orig_node, &agg_node) in flat_mapping.iter().enumerate() {
            let comm = partition.community_of(agg_node);
            result.move_node(orig_node, comm);
        }
        result.renumber();
        let q = quality.total_quality(&data, &result);
        Ok(LeidenOutput::new(result, q))
    }

    /// Run the Leiden algorithm and capture intermediate community structure
    /// at each aggregation level.
    ///
    /// Returns a [`HierarchicalOutput`](crate::hierarchy::HierarchicalOutput)
    /// containing the final partition, quality, and per-level membership vectors.
    pub fn run_hierarchical<V, G>(
        &self,
        graph: &gryf::Graph<V, f64, Undirected, G>,
    ) -> crate::error::Result<crate::hierarchy::HierarchicalOutput>
    where
        G: GraphBase<VertexId: IntegerIdType, EdgeType = Undirected>
            + Neighbors
            + VertexSet
            + GraphRef<V, f64>,
    {
        use crate::hierarchy::{HierarchicalOutput, HierarchyLevel};

        let original_n = graph.vertex_count();
        if original_n == 0 {
            return Ok(HierarchicalOutput {
                levels: vec![HierarchyLevel {
                    node_count: 0,
                    num_communities: 0,
                    quality: 0.0,
                    membership: vec![],
                }],
                partition: Partition::new(0),
                quality: 0.0,
            });
        }

        let mut data = GraphData::from_gryf_graph(graph)?;
        let mut partition = Partition::new(data.node_count());
        let mut flat_mapping: Vec<usize> = (0..data.node_count()).collect();

        let mut rng = match self.config.seed {
            Some(seed) => StdRng::seed_from_u64(seed),
            None => StdRng::from_rng(&mut rand::rng()),
        };

        let modularity = Modularity::with_resolution(self.config.resolution);
        let cpm = crate::quality::CPM::new(self.config.resolution);
        let rbconfig = crate::quality::RBConfiguration::with_resolution(self.config.resolution);
        let rber = crate::quality::RBER::new(self.config.resolution);
        let quality: &(dyn QualityFunction + Sync) = match self.config.quality {
            QualityType::Modularity => &modularity,
            QualityType::CPM => &cpm,
            QualityType::RBConfiguration => &rbconfig,
            QualityType::RBER => &rber,
        };

        let mut levels: Vec<HierarchyLevel> = Vec::new();

        for _iter in 0..self.config.max_iterations {
            let q_before = quality.total_quality(&data, &partition);
            let changed = local_moving_dispatch(
                &data,
                &mut partition,
                quality,
                &mut rng,
                self.config.max_comm_size,
                self.config.epsilon,
            );
            if !changed {
                break;
            }
            partition.renumber();

            let q_after = quality.total_quality(&data, &partition);

            let level_membership = resolve_to_original_nodes(&flat_mapping, &partition, original_n);
            levels.push(HierarchyLevel {
                node_count: data.node_count(),
                num_communities: partition.num_communities(),
                quality: q_after,
                membership: level_membership,
            });

            if (q_after - q_before).abs() < self.config.epsilon {
                break;
            }

            let refined = refinement(&data, &partition, quality, &mut rng, self.config.epsilon);

            let (agg_data, orig_to_agg, agg_initial_partition) =
                aggregate(&data, &refined, &partition);

            for orig_node in 0..original_n {
                flat_mapping[orig_node] = orig_to_agg[flat_mapping[orig_node]];
            }

            data = agg_data;
            partition = agg_initial_partition;

            if data.node_count() <= 1 {
                break;
            }
        }

        let mut result = Partition::from_membership(vec![0; original_n]);
        for (orig_node, &agg_node) in flat_mapping.iter().enumerate() {
            let comm = partition.community_of(agg_node);
            result.move_node(orig_node, comm);
        }
        result.renumber();
        let q = quality.total_quality(&data, &result);

        Ok(HierarchicalOutput {
            levels,
            partition: result,
            quality: q,
        })
    }
}

fn resolve_to_original_nodes(
    flat_mapping: &[usize],
    partition: &Partition,
    original_n: usize,
) -> Vec<usize> {
    let mut result = vec![0; original_n];
    for orig_node in 0..original_n {
        let agg_node = flat_mapping[orig_node];
        result[orig_node] = partition.community_of(agg_node);
    }
    result
}

fn local_moving_sequential(
    data: &GraphData,
    partition: &mut Partition,
    quality: &dyn QualityFunction,
    rng: &mut StdRng,
    max_comm_size: usize,
    epsilon: f64,
) -> bool {
    let n = data.node_count();
    if n == 0 {
        return false;
    }

    let mut community_total_degree: Vec<f64> = vec![0.0; n];
    let mut community_size: Vec<f64> = vec![0.0; n];
    for node in 0..n {
        let comm = partition.community_of(node);
        community_total_degree[comm] += data.degree_of(node);
        community_size[comm] += data.node_weight(node);
    }

    let mut order: Vec<usize> = (0..n).filter(|&node| data.degree_of(node) > 0.0).collect();
    order.shuffle(rng);
    let mut queue: VecDeque<usize> = order.into_iter().collect();
    let mut in_queue: Vec<bool> = vec![false; n];
    for &node in &queue {
        in_queue[node] = true;
    }
    let m = data.total_weight();

    if m <= 0.0 {
        return false;
    }

    let two_m = 2.0 * m;
    let total_node_weight: f64 = data.total_node_weight();

    let mut changed = false;
    let mut neighbor_comm_weights: Vec<f64> = vec![0.0; n];
    let mut touched: Vec<usize> = Vec::with_capacity(64);

    while let Some(node) = queue.pop_front() {
        in_queue[node] = false;
        touched.clear();
        let current_community = partition.community_of(node);
        let k_v = data.degree_of(node);

        neighbor_comm_weights[current_community] = 0.0;
        let mut current_touched = false;

        let (targets, weights) = data.neighbor_slices(node);
        for i in 0..targets.len() {
            let neighbor = targets[i];
            let weight = weights[i];
            if neighbor == node {
                continue;
            }
            let comm = partition.community_of(neighbor);
            if neighbor_comm_weights[comm] == 0.0 {
                if comm == current_community {
                    current_touched = true;
                } else {
                    touched.push(comm);
                }
            }
            neighbor_comm_weights[comm] += weight;
        }

        let k_v_to_current = neighbor_comm_weights[current_community];
        let sigma_tot_current = community_total_degree[current_community];
        let node_weight = data.node_weight(node);

        let mut best_community = current_community;
        let mut best_delta = epsilon;

        for &target_comm in &touched {
            let k_v_to_target = neighbor_comm_weights[target_comm];
            if max_comm_size > 0 && community_size[target_comm] + node_weight > max_comm_size as f64
            {
                continue;
            }
            let sigma_tot_target = community_total_degree[target_comm];
            let delta = quality.delta_move_from_components(&MoveComponents {
                k_v,
                k_v_to_target,
                k_v_to_current,
                sigma_tot_target,
                sigma_tot_current,
                two_m,
                n_target: community_size[target_comm],
                n_current: community_size[current_community],
                node_weight,
                total_node_weight,
            });
            if delta > best_delta {
                best_delta = delta;
                best_community = target_comm;
            }
        }

        if current_touched {
            neighbor_comm_weights[current_community] = 0.0;
        }
        for &comm in &touched {
            neighbor_comm_weights[comm] = 0.0;
        }

        if best_community != current_community {
            let w = data.node_weight(node);
            partition.move_node(node, best_community);
            community_total_degree[current_community] -= k_v;
            community_total_degree[best_community] += k_v;
            community_size[current_community] -= w;
            community_size[best_community] += w;
            changed = true;

            let (targets, _) = data.neighbor_slices(node);
            for &neighbor in targets {
                if !in_queue[neighbor] {
                    queue.push_back(neighbor);
                    in_queue[neighbor] = true;
                }
            }
        }
    }

    changed
}

/// Greedy distance-1 graph coloring.
/// Returns (colors: Vec<usize>, num_colors: usize) where colors[node] is the color assignment.
/// Uses at most max_degree + 1 colors. O(V + E) time.
#[cfg(feature = "rayon")]
fn greedy_coloring(data: &GraphData, order: &[usize]) -> (Vec<usize>, usize) {
    let n = data.node_count();
    let mut colors = vec![0usize; n];
    let mut used_colors: Vec<bool> = Vec::new();
    let mut num_colors = 1usize;

    for &node in order {
        let (targets, _) = data.neighbor_slices(node);
        let max_neighbor_color = targets.iter().map(|&t| colors[t]).max().unwrap_or(0);
        if used_colors.len() <= max_neighbor_color {
            used_colors.resize(max_neighbor_color + 1, false);
        }
        for &neighbor in targets {
            used_colors[colors[neighbor]] = true;
        }

        let mut color = 0;
        while color < used_colors.len() && used_colors[color] {
            color += 1;
        }
        colors[node] = color;
        if color + 1 > num_colors {
            num_colors = color + 1;
        }

        for &neighbor in targets {
            used_colors[colors[neighbor]] = false;
        }
    }

    (colors, num_colors)
}

/// Threshold below which sequential local moving is used (avoids parallelization overhead).
#[cfg(feature = "rayon")]
const PARALLEL_THRESHOLD: usize = 500;

#[cfg(feature = "rayon")]
fn local_moving_parallel(
    data: &GraphData,
    partition: &mut Partition,
    quality: &(dyn QualityFunction + Sync),
    rng: &mut StdRng,
    max_comm_size: usize,
    epsilon: f64,
) -> bool {
    let n = data.node_count();
    if n == 0 {
        return false;
    }

    let m = data.total_weight();
    if m <= 0.0 {
        return false;
    }
    let two_m = 2.0 * m;
    let total_node_weight: f64 = data.total_node_weight();

    let mut order: Vec<usize> = (0..n).filter(|&node| data.degree_of(node) > 0.0).collect();
    order.shuffle(rng);
    let (colors, num_colors) = greedy_coloring(data, &order);

    let mut color_groups: Vec<Vec<usize>> = vec![Vec::new(); num_colors];
    for &node in &order {
        color_groups[colors[node]].push(node);
    }

    let mut community_total_degree: Vec<f64> = vec![0.0; n];
    let mut community_size: Vec<f64> = vec![0.0; n];
    for node in 0..n {
        let comm = partition.community_of(node);
        community_total_degree[comm] += data.degree_of(node);
        community_size[comm] += data.node_weight(node);
    }

    let mut changed = false;
    let mut any_move = true;

    while any_move {
        any_move = false;

        for group in &color_groups {
            if group.is_empty() {
                continue;
            }

            // Phase 1: Parallel compute best moves (read-only access to community totals)
            let moves: Vec<(usize, usize, usize, f64, f64)> = group
                .par_iter()
                .filter_map(|&node| {
                    let current_community = partition.community_of(node);
                    let k_v = data.degree_of(node);
                    let node_weight = data.node_weight(node);

                    let (targets, weights) = data.neighbor_slices(node);
                    let mut neighbor_comm_weights: Vec<(usize, f64)> = Vec::new();
                    let mut k_v_to_current = 0.0f64;

                    for i in 0..targets.len() {
                        let neighbor = targets[i];
                        let weight = weights[i];
                        if neighbor == node {
                            continue;
                        }
                        let comm = partition.community_of(neighbor);
                        if comm == current_community {
                            k_v_to_current += weight;
                        } else if let Some(entry) =
                            neighbor_comm_weights.iter_mut().find(|(c, _)| *c == comm)
                        {
                            entry.1 += weight;
                        } else {
                            neighbor_comm_weights.push((comm, weight));
                        }
                    }

                    let sigma_tot_current = community_total_degree[current_community];
                    let mut best_community = current_community;
                    let mut best_delta = epsilon;

                    for (target_comm, k_v_to_target) in &neighbor_comm_weights {
                        if max_comm_size > 0
                            && community_size[*target_comm] + node_weight > max_comm_size as f64
                        {
                            continue;
                        }
                        let sigma_tot_target = community_total_degree[*target_comm];
                        let delta = quality.delta_move_from_components(&MoveComponents {
                            k_v,
                            k_v_to_target: *k_v_to_target,
                            k_v_to_current,
                            sigma_tot_target,
                            sigma_tot_current,
                            two_m,
                            n_target: community_size[*target_comm],
                            n_current: community_size[current_community],
                            node_weight,
                            total_node_weight,
                        });
                        if delta > best_delta {
                            best_delta = delta;
                            best_community = *target_comm;
                        }
                    }

                    if best_community != current_community {
                        Some((node, current_community, best_community, k_v, node_weight))
                    } else {
                        None
                    }
                })
                .collect();

            // Phase 2: Sequential apply moves
            for (node, old_comm, new_comm, k_v, node_weight) in moves {
                partition.move_node(node, new_comm);
                community_total_degree[old_comm] -= k_v;
                community_total_degree[new_comm] += k_v;
                community_size[old_comm] -= node_weight;
                community_size[new_comm] += node_weight;
                any_move = true;
                changed = true;
            }
        }
    }

    changed
}

fn local_moving_dispatch(
    data: &GraphData,
    partition: &mut Partition,
    quality: &(dyn QualityFunction + Sync),
    rng: &mut StdRng,
    max_comm_size: usize,
    epsilon: f64,
) -> bool {
    #[cfg(feature = "rayon")]
    {
        if data.node_count() >= PARALLEL_THRESHOLD {
            return local_moving_parallel(data, partition, quality, rng, max_comm_size, epsilon);
        }
    }
    let _ = data;
    let _ = partition;
    let _ = quality;
    let _ = rng;
    let _ = max_comm_size;
    let _ = epsilon;
    local_moving_sequential(data, partition, quality, rng, max_comm_size, epsilon)
}

fn refine_community(
    data: &GraphData,
    partition: &Partition,
    quality: &(dyn QualityFunction + Sync),
    community: usize,
    nodes: &[usize],
    m: f64,
    epsilon: f64,
) -> Vec<(usize, usize)> {
    if nodes.len() <= 1 {
        return Vec::new();
    }

    let two_m = 2.0 * m;
    let total_node_weight: f64 = data.total_node_weight();

    let max_node_id = nodes.iter().copied().max().unwrap_or(0);
    let mut refined_map: Vec<usize> = (0..=max_node_id).collect();

    let mut comm_total_degree: Vec<f64> = vec![0.0; max_node_id + 1];
    let mut comm_size: Vec<f64> = vec![0.0; max_node_id + 1];
    for &node in nodes {
        comm_total_degree[node] += data.degree_of(node);
        comm_size[node] += data.node_weight(node);
    }

    let mut neighbor_comm_weights: Vec<f64> = vec![0.0; max_node_id + 1];
    let mut touched: Vec<usize> = Vec::new();

    for &node in nodes {
        let current_refined = refined_map[node];
        let k_v = data.degree_of(node);

        let (targets, weights) = data.neighbor_slices(node);
        for i in 0..targets.len() {
            let neighbor = targets[i];
            let weight = weights[i];
            if partition.community_of(neighbor) != community {
                continue;
            }
            if neighbor == node {
                continue;
            }
            let rc = refined_map[neighbor];
            if neighbor_comm_weights[rc] == 0.0 && rc != current_refined {
                touched.push(rc);
            }
            neighbor_comm_weights[rc] += weight;
        }

        let k_v_to_current = neighbor_comm_weights[current_refined];
        let sigma_tot_current = comm_total_degree[current_refined];

        let mut best_refined = current_refined;
        let mut best_delta = epsilon;

        for &target_rc in &touched {
            let k_v_to_target = neighbor_comm_weights[target_rc];
            let sigma_tot_target = comm_total_degree[target_rc];
            let n_current = comm_size[current_refined];
            let n_target = comm_size[target_rc];
            let delta = quality.delta_move_from_components(&MoveComponents {
                k_v,
                k_v_to_target,
                k_v_to_current,
                sigma_tot_target,
                sigma_tot_current,
                two_m,
                n_target,
                n_current,
                node_weight: data.node_weight(node),
                total_node_weight,
            });
            if delta > best_delta {
                best_delta = delta;
                best_refined = target_rc;
            }
        }

        // Cleanup: reset touched entries for next iteration
        for &rc in &touched {
            neighbor_comm_weights[rc] = 0.0;
        }
        neighbor_comm_weights[current_refined] = 0.0;
        touched.clear();

        if best_refined != current_refined {
            let w = data.node_weight(node);
            refined_map[node] = best_refined;
            comm_total_degree[current_refined] -= k_v;
            comm_total_degree[best_refined] += k_v;
            comm_size[current_refined] -= w;
            comm_size[best_refined] += w;
        }
    }

    nodes
        .iter()
        .filter_map(|&node| {
            let rc = refined_map[node];
            if rc != node {
                Some((node, rc))
            } else {
                None
            }
        })
        .collect()
}

fn refinement(
    data: &GraphData,
    partition: &Partition,
    quality: &(dyn QualityFunction + Sync),
    rng: &mut StdRng,
    epsilon: f64,
) -> Partition {
    let n = data.node_count();
    let mut refined = Partition::new(n);

    let m = data.total_weight();
    if m <= 0.0 {
        return refined;
    }

    // Collect community node lists and shuffle sequentially (preserves seed determinism)
    let num_comms = partition.num_communities();
    let mut community_nodes: Vec<Vec<usize>> = vec![Vec::new(); num_comms];
    for node in 0..n {
        community_nodes[partition.community_of(node)].push(node);
    }
    for nodes in &mut community_nodes {
        nodes.shuffle(rng);
    }

    let results: Vec<Vec<(usize, usize)>> = {
        #[cfg(feature = "rayon")]
        if num_comms > 32 {
            community_nodes
                .par_iter()
                .enumerate()
                .map(|(community, nodes)| {
                    refine_community(data, partition, quality, community, nodes, m, epsilon)
                })
                .collect()
        } else {
            community_nodes
                .iter()
                .enumerate()
                .map(|(community, nodes)| {
                    refine_community(data, partition, quality, community, nodes, m, epsilon)
                })
                .collect()
        }
        #[cfg(not(feature = "rayon"))]
        {
            community_nodes
                .iter()
                .enumerate()
                .map(|(community, nodes)| {
                    refine_community(data, partition, quality, community, nodes, m, epsilon)
                })
                .collect()
        }
    };

    for moves in &results {
        for &(node, new_comm) in moves {
            refined.move_node(node, new_comm);
        }
    }

    refined.renumber();
    refined
}

fn aggregate(
    data: &GraphData,
    refined_partition: &Partition,
    coarse_partition: &Partition,
) -> (GraphData, Vec<usize>, Partition) {
    let n = data.node_count();

    let mut orig_to_agg: Vec<usize> = vec![0; n];
    let mut comm_to_agg: FxHashMap<usize, usize> = FxHashMap::default();
    let mut next_id = 0usize;
    for (node, entry) in orig_to_agg.iter_mut().enumerate() {
        let c = refined_partition.community_of(node);
        let agg_id = *comm_to_agg.entry(c).or_insert_with(|| {
            let id = next_id;
            next_id += 1;
            id
        });
        *entry = agg_id;
    }

    let agg_n = next_id;

    let mut agg_edges: FxHashMap<(usize, usize), f64> = FxHashMap::default();
    for u in 0..n {
        let ru = orig_to_agg[u];
        for (v, w) in data.neighbors(u) {
            if u == v {
                *agg_edges.entry((ru, ru)).or_default() += w;
                continue;
            }
            if v <= u {
                continue;
            }
            let rv = orig_to_agg[v];
            let key = if ru <= rv { (ru, rv) } else { (rv, ru) };
            *agg_edges.entry(key).or_default() += w;
        }
    }

    // Build flat CSR directly from FxHashMap
    // First, collect degree counts per node
    let mut agg_degree_count: Vec<usize> = vec![0; agg_n];
    for &(u, v) in agg_edges.keys() {
        agg_degree_count[u] += 1;
        if u != v {
            agg_degree_count[v] += 1;
        }
    }

    let mut agg_offsets = vec![0usize; agg_n + 1];
    for i in 0..agg_n {
        agg_offsets[i + 1] = agg_offsets[i] + agg_degree_count[i];
    }

    let total_slots = agg_offsets[agg_n];
    let mut agg_targets = vec![0usize; total_slots];
    let mut agg_weights = vec![0.0f64; total_slots];
    let mut insert_pos = agg_offsets.clone();

    for ((u, v), w) in &agg_edges {
        let pos = insert_pos[*u];
        agg_targets[pos] = *v;
        agg_weights[pos] = *w;
        insert_pos[*u] += 1;
        if *u != *v {
            let pos = insert_pos[*v];
            agg_targets[pos] = *u;
            agg_weights[pos] = *w;
            insert_pos[*v] += 1;
        }
    }

    // Sort each node's neighbor list by target ID
    for u in 0..agg_n {
        let start = agg_offsets[u];
        let end = agg_offsets[u + 1];
        if end - start <= 1 {
            continue;
        }
        let slice_targets = &mut agg_targets[start..end];
        let slice_weights = &mut agg_weights[start..end];
        let len = slice_targets.len();
        let mut indices: Vec<usize> = (0..len).collect();
        indices.sort_by_key(|&i| slice_targets[i]);
        let sorted_targets: Vec<usize> = indices.iter().map(|&i| slice_targets[i]).collect();
        let sorted_weights: Vec<f64> = indices.iter().map(|&i| slice_weights[i]).collect();
        slice_targets.copy_from_slice(&sorted_targets);
        slice_weights.copy_from_slice(&sorted_weights);
    }

    let mut agg_degree: Vec<f64> = vec![0.0; agg_n];
    let mut agg_node_weight: Vec<f64> = vec![0.0; agg_n];
    for (orig, &agg_node) in orig_to_agg.iter().enumerate() {
        agg_degree[agg_node] += data.degree_of(orig);
        agg_node_weight[agg_node] += data.node_weight(orig);
    }

    let agg_data = GraphData::from_parts(
        agg_n,
        agg_offsets,
        agg_targets,
        agg_weights,
        agg_degree,
        agg_node_weight,
    )
    .expect("internal CSR construction failed");

    let mut agg_initial = Partition::new(agg_n);
    for (orig, &agg_node) in orig_to_agg.iter().enumerate() {
        let coarse_comm = coarse_partition.community_of(orig);
        agg_initial.move_node(agg_node, coarse_comm);
    }
    agg_initial.renumber();

    (agg_data, orig_to_agg, agg_initial)
}

#[cfg(test)]
mod tests {
    use super::*;
    use gryf::Graph;
    use rand::prelude::IndexedRandom;

    fn make_two_cliques() -> Graph<i32, f64, Undirected> {
        let mut graph = Graph::new_undirected();
        let nodes: Vec<_> = (0..10).map(|i| graph.add_vertex(i)).collect();

        for i in 0..5 {
            for j in (i + 1)..5 {
                graph.add_edge(&nodes[i], &nodes[j], 1.0);
            }
        }
        for i in 5..10 {
            for j in (i + 1)..10 {
                graph.add_edge(&nodes[i], &nodes[j], 1.0);
            }
        }
        graph.add_edge(&nodes[0], &nodes[5], 1.0);
        graph
    }

    #[test]
    fn test_two_cliques() {
        let graph = make_two_cliques();
        let leiden = Leiden::new(LeidenConfig::default());
        let partition = leiden.run(&graph).unwrap().partition;

        assert_eq!(partition.num_communities(), 2);

        let comm0 = partition.community_of(0);
        for i in 1..5 {
            assert_eq!(partition.community_of(i), comm0, "node {i}");
        }
        let comm5 = partition.community_of(5);
        for i in 6..10 {
            assert_eq!(partition.community_of(i), comm5, "node {i}");
        }
        assert_ne!(comm0, comm5);
    }

    #[test]
    fn test_single_node() {
        let mut graph: Graph<i32, f64, _> = Graph::new_undirected();
        graph.add_vertex(1);

        let leiden = Leiden::new(LeidenConfig::default());
        let partition = leiden.run(&graph).unwrap().partition;
        assert_eq!(partition.num_communities(), 1);
    }

    #[test]
    fn test_ring_of_cliques() {
        let mut graph = Graph::new_undirected();
        let nodes: Vec<_> = (0..12).map(|i| graph.add_vertex(i)).collect();

        for i in 0..4 {
            for j in (i + 1)..4 {
                graph.add_edge(&nodes[i], &nodes[j], 1.0);
            }
        }
        for i in 4..8 {
            for j in (i + 1)..8 {
                graph.add_edge(&nodes[i], &nodes[j], 1.0);
            }
        }
        for i in 8..12 {
            for j in (i + 1)..12 {
                graph.add_edge(&nodes[i], &nodes[j], 1.0);
            }
        }
        graph.add_edge(&nodes[0], &nodes[4], 0.1);
        graph.add_edge(&nodes[4], &nodes[8], 0.1);
        graph.add_edge(&nodes[8], &nodes[0], 0.1);

        let leiden = Leiden::new(LeidenConfig::default());
        let partition = leiden.run(&graph).unwrap().partition;

        assert_eq!(partition.num_communities(), 3);
    }

    #[test]
    fn test_empty_graph() {
        let graph: Graph<i32, f64, _> = Graph::new_undirected();
        let leiden = Leiden::new(LeidenConfig::default());
        let partition = leiden.run(&graph).unwrap().partition;
        assert_eq!(partition.num_communities(), 0);
    }

    #[test]
    fn test_disconnected_graph() {
        let mut graph = Graph::new_undirected();
        let nodes: Vec<_> = (0..6).map(|i| graph.add_vertex(i)).collect();

        // Triangle 0-1-2
        for i in 0..3 {
            for j in (i + 1)..3 {
                graph.add_edge(&nodes[i], &nodes[j], 1.0);
            }
        }
        // Triangle 3-4-5
        for i in 3..6 {
            for j in (i + 1)..6 {
                graph.add_edge(&nodes[i], &nodes[j], 1.0);
            }
        }

        let leiden = Leiden::new(LeidenConfig::default());
        let partition = leiden.run(&graph).unwrap().partition;

        assert_eq!(partition.num_communities(), 2);

        let comm0 = partition.community_of(0);
        assert_eq!(partition.community_of(1), comm0);
        assert_eq!(partition.community_of(2), comm0);

        let comm3 = partition.community_of(3);
        assert_eq!(partition.community_of(4), comm3);
        assert_eq!(partition.community_of(5), comm3);

        assert_ne!(comm0, comm3);
    }

    #[test]
    fn test_single_edge() {
        let mut graph = Graph::new_undirected();
        let nodes: Vec<_> = (0..2).map(|i| graph.add_vertex(i)).collect();
        graph.add_edge(&nodes[0], &nodes[1], 1.0);

        let leiden = Leiden::new(LeidenConfig::default());
        let partition = leiden.run(&graph).unwrap().partition;

        assert_eq!(partition.num_communities(), 1);
        assert_eq!(partition.community_of(0), partition.community_of(1));
    }

    #[test]
    fn test_weighted_graph() {
        let mut graph = Graph::new_undirected();
        let nodes: Vec<_> = (0..6).map(|i| graph.add_vertex(i)).collect();

        // Within-group edges: weight 5.0
        for i in 0..3 {
            for j in (i + 1)..3 {
                graph.add_edge(&nodes[i], &nodes[j], 5.0);
            }
        }
        for i in 3..6 {
            for j in (i + 1)..6 {
                graph.add_edge(&nodes[i], &nodes[j], 5.0);
            }
        }
        // Single bridge edge: weight 0.1
        graph.add_edge(&nodes[0], &nodes[3], 0.1);

        let leiden = Leiden::new(LeidenConfig::default());
        let partition = leiden.run(&graph).unwrap().partition;

        assert_eq!(partition.num_communities(), 2);

        let comm0 = partition.community_of(0);
        for i in 1..3 {
            assert_eq!(partition.community_of(i), comm0, "node {i}");
        }
        let comm3 = partition.community_of(3);
        for i in 4..6 {
            assert_eq!(partition.community_of(i), comm3, "node {i}");
        }
        assert_ne!(comm0, comm3);
    }

    #[test]
    fn test_large_random_graph_determinism() {
        let mut rng = StdRng::seed_from_u64(12345);

        let mut graph = Graph::new_undirected();
        let nodes: Vec<_> = (0..100).map(|i| graph.add_vertex(i)).collect();

        let clusters: Vec<Vec<usize>> = (0..4).map(|c| (c * 25..(c + 1) * 25).collect()).collect();

        // Within-cluster edges: ~5 random connections per node, weight 1.0
        for cluster in &clusters {
            for &node in cluster {
                let others: Vec<usize> = cluster.iter().filter(|&&x| x != node).copied().collect();
                let count = std::cmp::min(5, others.len());
                let chosen: Vec<usize> = others.choose_multiple(&mut rng, count).copied().collect();
                for &neighbor in &chosen {
                    if neighbor > node {
                        graph.add_edge(&nodes[node], &nodes[neighbor], 1.0);
                    }
                }
            }
        }

        // Between-cluster edges: ~3 per cluster pair, weight 0.1
        for i in 0..4 {
            for j in (i + 1)..4 {
                let mut pairs: Vec<(usize, usize)> = Vec::new();
                for &a in &clusters[i] {
                    for &b in &clusters[j] {
                        pairs.push((a, b));
                    }
                }
                let count = std::cmp::min(3, pairs.len());
                for &(a, b) in pairs.choose_multiple(&mut rng, count) {
                    graph.add_edge(&nodes[a], &nodes[b], 0.1);
                }
            }
        }

        let leiden = Leiden::new(LeidenConfig {
            seed: Some(42),
            ..Default::default()
        });
        let partition1 = leiden.run(&graph).unwrap().partition;

        let leiden2 = Leiden::new(LeidenConfig {
            seed: Some(42),
            ..Default::default()
        });
        let partition2 = leiden2.run(&graph).unwrap().partition;

        assert_eq!(
            partition1.as_slice(),
            partition2.as_slice(),
            "same seed should produce identical partitions"
        );
        assert!(
            partition1.num_communities() >= 2,
            "expected at least 2 communities, got {}",
            partition1.num_communities()
        );
    }

    #[test]
    fn test_star_graph() {
        let mut graph = Graph::new_undirected();
        let nodes: Vec<_> = (0..11).map(|i| graph.add_vertex(i)).collect();

        for i in 1..11 {
            graph.add_edge(&nodes[0], &nodes[i], 1.0);
        }

        let leiden = Leiden::new(LeidenConfig::default());
        let partition = leiden.run(&graph).unwrap().partition;

        assert!(partition.num_communities() >= 1);
    }

    #[test]
    fn test_isolated_nodes_with_edges() {
        let mut graph = Graph::new_undirected();
        let nodes: Vec<_> = (0..9).map(|i| graph.add_vertex(i)).collect();

        // Clique 0-3
        for i in 0..4 {
            for j in (i + 1)..4 {
                graph.add_edge(&nodes[i], &nodes[j], 1.0);
            }
        }
        // Clique 4-7
        for i in 4..8 {
            for j in (i + 1)..8 {
                graph.add_edge(&nodes[i], &nodes[j], 1.0);
            }
        }
        // Node 8 is isolated (no edges)

        let leiden = Leiden::new(LeidenConfig::default());
        let partition = leiden.run(&graph).unwrap().partition;

        assert_eq!(partition.num_communities(), 3);

        let comm0 = partition.community_of(0);
        for i in 1..4 {
            assert_eq!(partition.community_of(i), comm0, "node {i}");
        }
        let comm4 = partition.community_of(4);
        for i in 5..8 {
            assert_eq!(partition.community_of(i), comm4, "node {i}");
        }
        // Node 8 should be in its own community
        let comm8 = partition.community_of(8);
        assert_ne!(comm8, comm0);
        assert_ne!(comm8, comm4);
    }

    #[test]
    fn test_seed_determinism_different_seeds() {
        let graph = make_two_cliques();

        let config1 = LeidenConfig {
            seed: Some(1),
            ..Default::default()
        };
        let leiden1 = Leiden::new(config1);
        let partition1 = leiden1.run(&graph).unwrap().partition;

        let config2 = LeidenConfig {
            seed: Some(2),
            ..Default::default()
        };
        let leiden2 = Leiden::new(config2);
        let partition2 = leiden2.run(&graph).unwrap().partition;

        assert_eq!(partition1.num_communities(), 2);
        assert_eq!(partition2.num_communities(), 2);
    }

    #[test]
    fn test_self_loop() {
        let mut graph = Graph::new_undirected();
        let nodes: Vec<_> = (0..4).map(|i| graph.add_vertex(i)).collect();

        graph.add_edge(&nodes[0], &nodes[1], 1.0);
        graph.add_edge(&nodes[2], &nodes[3], 1.0);
        graph.add_edge(&nodes[0], &nodes[0], 2.0);

        let leiden = Leiden::new(LeidenConfig::default());
        let partition = leiden.run(&graph).unwrap().partition;

        assert_eq!(partition.num_communities(), 2);
        assert_eq!(partition.community_of(0), partition.community_of(1));
        assert_eq!(partition.community_of(2), partition.community_of(3));
        assert_ne!(partition.community_of(0), partition.community_of(2));
    }

    #[test]
    fn test_hierarchical_matches_run() {
        let graph = make_two_cliques();
        let leiden = Leiden::new(LeidenConfig {
            seed: Some(42),
            ..Default::default()
        });
        let normal = leiden.run(&graph).unwrap();
        let hierarchical = leiden.run_hierarchical(&graph).unwrap();

        assert_eq!(
            normal.partition.as_slice(),
            hierarchical.partition.as_slice()
        );
        assert!((normal.quality - hierarchical.quality).abs() < 1e-10);
        assert!(!hierarchical.levels.is_empty());
    }

    #[test]
    fn test_hierarchical_levels_decrease_nodes() {
        let graph = make_two_cliques();
        let leiden = Leiden::new(LeidenConfig {
            seed: Some(42),
            ..Default::default()
        });
        let h = leiden.run_hierarchical(&graph).unwrap();

        for window in h.levels.windows(2) {
            assert!(window[1].node_count <= window[0].node_count);
        }

        if !h.levels.is_empty() {
            let last_level = &h.levels[h.levels.len() - 1];
            for node in 0..graph.vertex_count() {
                assert_eq!(last_level.membership[node], h.partition.community_of(node));
            }
        }
    }

    #[test]
    fn test_hierarchical_community_of_at_level() {
        let graph = make_two_cliques();
        let leiden = Leiden::new(LeidenConfig {
            seed: Some(42),
            ..Default::default()
        });
        let h = leiden.run_hierarchical(&graph).unwrap();

        if h.levels.len() >= 2 {
            assert!(h.levels[0].num_communities >= h.partition.num_communities());
        }
    }

    #[test]
    #[cfg(feature = "rayon")]
    fn test_coloring_basic() {
        use crate::quality::GraphData;
        let graph = make_two_cliques();
        let data = GraphData::from_gryf_graph(&graph).unwrap();
        let order: Vec<usize> = (0..data.node_count()).collect();
        let (colors, num_colors) = greedy_coloring(&data, &order);

        for node in 0..data.node_count() {
            let (targets, _) = data.neighbor_slices(node);
            for &neighbor in targets {
                if neighbor != node {
                    assert_ne!(
                        colors[node], colors[neighbor],
                        "Adjacent nodes {} and {} have same color {}",
                        node, neighbor, colors[node]
                    );
                }
            }
        }
        assert!(num_colors > 0);
    }

    #[test]
    fn test_parallel_matches_sequential() {
        let graph = make_two_cliques();
        let leiden = Leiden::new(LeidenConfig {
            seed: Some(42),
            ..Default::default()
        });
        let result_normal = leiden.run(&graph).unwrap();
        assert!(result_normal.partition.num_communities() >= 1);
    }
}