leiden-rs 0.7.0

//! Shared algorithm internals for single-layer and multiplex Leiden.

use std::collections::VecDeque;

use crate::graph::GraphData;
use crate::graph::GraphDataBuilder;
use crate::partition::Partition;
use crate::quality::{MoveComponents, QualityFunction};
use rand::rngs::StdRng;
use rand::seq::SliceRandom;
#[cfg(feature = "rayon")]
use rayon::prelude::*;
use rustc_hash::FxHashMap;

/// Find the best community to move a node to.
///
/// `delta_fn(target_comm) -> f64` returns the quality delta for moving to that community.
/// Returns `(best_community, best_delta)`. Stays at current if no improvement > epsilon.
pub(crate) fn find_best_community(
    candidates: impl Iterator<Item = usize>,
    current_community: usize,
    epsilon: f64,
    max_comm_size: usize,
    comm_size: &[f64],
    node_weight: f64,
    delta_fn: impl Fn(usize) -> f64,
) -> (usize, f64) {
    let mut best_community = current_community;
    let mut best_delta = epsilon;
    for target_comm in candidates {
        if max_comm_size > 0 && comm_size[target_comm] + node_weight > max_comm_size as f64 {
            continue;
        }
        let delta = delta_fn(target_comm);
        if delta > best_delta {
            best_delta = delta;
            best_community = target_comm;
        }
    }
    (best_community, best_delta)
}

/// Apply a community move and update statistics arrays.
///
/// If `refined_map` is `Some`, updates `refined_map[node] = new_comm` (refinement phase).
/// If `refined_map` is `None`, calls `partition.move_node(node, new_comm)` (local moving).
/// In-degree update only when `comm_in_degree` is non-empty.
pub(crate) fn apply_move(
    partition: Option<&mut Partition>,
    refined_map: Option<&mut [usize]>,
    node: usize,
    old_comm: usize,
    new_comm: usize,
    k_v_out: f64,
    k_v_in: f64,
    node_weight: f64,
    comm_total_degree: &mut [f64],
    comm_in_degree: &mut [f64],
    comm_size: &mut [f64],
) {
    match refined_map {
        Some(map) => map[node] = new_comm,
        None => partition
            .expect("partition required when refined_map is None")
            .move_node(node, new_comm),
    }
    comm_total_degree[old_comm] -= k_v_out;
    comm_total_degree[new_comm] += k_v_out;
    if !comm_in_degree.is_empty() {
        comm_in_degree[old_comm] -= k_v_in;
        comm_in_degree[new_comm] += k_v_in;
    }
    comm_size[old_comm] -= node_weight;
    comm_size[new_comm] += node_weight;
}

/// Initialize community statistics arrays from layers and partition.
///
/// Returns `(comm_total_degree, comm_in_degree, comm_size)`.
/// `community_of_fn` maps node -> community index.
pub(crate) fn init_community_stats(
    n: usize,
    layers: &[GraphData],
    community_of_fn: impl Fn(usize) -> usize,
) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
    let mut comm_total_degree: Vec<f64> = vec![0.0; n];
    let mut comm_in_degree: Vec<f64> = vec![0.0; n];
    let mut comm_size: Vec<f64> = vec![0.0; n];
    for layer in layers {
        for node in 0..n {
            let comm = community_of_fn(node);
            comm_total_degree[comm] += layer.out_degree_of(node);
            comm_in_degree[comm] += layer.in_degree_of(node);
            comm_size[comm] += layer.node_weight(node);
        }
    }
    (comm_total_degree, comm_in_degree, comm_size)
}

/// Accumulate one node's edges into the aggregated edge map.
///
/// For directed graphs: maps (u→v) through orig_to_agg.
/// For undirected graphs: canonicalizes keys to (min, max) to avoid double-counting;
/// self-loops use (u, u); skips edges where v <= u (each undirected edge stored twice in CSR).
pub(crate) fn aggregate_node_edges_into(
    source: &GraphData,
    u: usize,
    orig_to_agg: &[usize],
    directed: bool,
    map: &mut FxHashMap<(usize, usize), f64>,
) {
    let ru = orig_to_agg[u];
    if directed {
        for (v, w) in source.neighbors(u) {
            let rv = orig_to_agg[v];
            *map.entry((ru, rv)).or_default() += w;
        }
    } else {
        for (v, w) in source.neighbors(u) {
            if u == v {
                *map.entry((ru, ru)).or_default() += w;
            } else if v > u {
                let rv = orig_to_agg[v];
                let key = if ru <= rv { (ru, rv) } else { (rv, ru) };
                *map.entry(key).or_default() += w;
            }
        }
    }
}

/// Build the orig_to_agg mapping from a refined partition.
///
/// Returns `(orig_to_agg, agg_n)` where `orig_to_agg[original_node] = aggregate_node_id`
/// and `agg_n` is the number of aggregate nodes.
pub(crate) fn build_orig_to_agg_mapping(
    n: usize,
    refined_partition: &Partition,
) -> (Vec<usize>, usize) {
    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;
    }
    (orig_to_agg, next_id)
}

/// Build the aggregated graph from collected edges and node weights.
///
/// Takes ownership of `orig_to_agg` and returns `(GraphData, Vec<usize>, Partition)`.
/// The `node_weight_fn` closure provides the weight for each original node.
pub(crate) fn build_aggregated_graph(
    orig_to_agg: Vec<usize>,
    agg_n: usize,
    directed: bool,
    agg_edges_map: FxHashMap<(usize, usize), f64>,
    coarse_partition: &Partition,
    node_weight_fn: impl Fn(usize) -> f64,
) -> (crate::graph::GraphData, Vec<usize>, Partition) {
    let mut agg_edges: Vec<((usize, usize), f64)> = agg_edges_map.into_iter().collect();
    agg_edges.sort_by_key(|&((u, v), _)| (u, v));

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

    let mut builder = GraphDataBuilder::new(agg_n);
    if directed {
        builder = builder.directed();
    }
    for &((u, v), w) in &agg_edges {
        builder.add_edge(u, v, w).expect("internal builder error");
    }
    for (node, &nw) in agg_node_weight.iter().enumerate() {
        if nw != 1.0 {
            builder
                .set_node_weight(node, nw)
                .expect("internal builder error");
        }
    }
    let agg_data = builder.build().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)
}

/// Generic refinement wrapper: collect community nodes, shuffle, dispatch refinement
/// per community (sequential or parallel), and apply moves.
///
/// The `refine_fn` closure receives `(community_index, community_node_list)` and
/// returns a list of `(node, new_community)` moves.
pub(crate) fn refinement_generic(
    n: usize,
    partition: &Partition,
    rng: &mut StdRng,
    refine_fn: impl Fn(usize, &[usize]) -> Vec<(usize, usize)> + Send + Sync,
) -> Partition {
    let mut refined = Partition::new(n);

    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")]
        {
            let par_threshold = std::cmp::max(4, rayon::current_num_threads() * 2);
            if num_comms > par_threshold {
                community_nodes
                    .par_iter()
                    .enumerate()
                    .map(|(community, nodes)| refine_fn(community, nodes))
                    .collect()
            } else {
                community_nodes
                    .iter()
                    .enumerate()
                    .map(|(community, nodes)| refine_fn(community, nodes))
                    .collect()
            }
        }
        #[cfg(not(feature = "rayon"))]
        {
            community_nodes
                .iter()
                .enumerate()
                .map(|(community, nodes)| refine_fn(community, nodes))
                .collect()
        }
    };

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

    refined.renumber();
    refined
}

/// Unified local moving for single-layer and multiplex Leiden.
///
/// Single-layer: `layers = &[data], layer_weights = &[1.0]`
/// Multi-layer: pass all layers and their weights.
///
/// Uses per-layer pre-computed neighbor weight buffers instead of re-scanning,
/// accumulating out-edge and in-edge weights separately for each layer in one pass.
pub(crate) fn local_moving_generic(
    layers: &[GraphData],
    layer_weights: &[f64],
    partition: &mut Partition,
    quality: &(dyn QualityFunction + Sync),
    rng: &mut StdRng,
    max_comm_size: usize,
    epsilon: f64,
) -> bool {
    let n = layers[0].node_count();
    if n == 0 {
        return false;
    }

    let num_layers = layers.len();

    let total_weight: f64 = layers.iter().map(|l| l.total_weight()).sum();
    if total_weight <= 0.0 {
        return false;
    }

    let two_m_values: Vec<f64> = layers.iter().map(|l| 2.0 * l.total_weight()).collect();
    let total_node_weight: f64 = layers[0].total_node_weight();

    let (mut comm_total_degree, mut comm_in_degree, mut comm_size) =
        init_community_stats(n, layers, |node| partition.community_of(node));

    let mut order: Vec<usize> = (0..n)
        .filter(|&node| layers.iter().any(|l| l.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 mut layer_out_weights: Vec<Vec<f64>> = vec![vec![0.0; n]; num_layers];
    let mut layer_in_weights: Vec<Vec<f64>> = vec![vec![0.0; n]; num_layers];

    let mut changed = false;
    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);

        for l in 0..num_layers {
            layer_out_weights[l][current_community] = 0.0;
            layer_in_weights[l][current_community] = 0.0;
        }

        let mut current_touched = false;

        for (l, layer) in layers.iter().enumerate() {
            let (targets, weights) = layer.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 layer_out_weights[l][comm] == 0.0 && layer_in_weights[l][comm] == 0.0 {
                    if comm == current_community {
                        current_touched = true;
                    } else if !touched.contains(&comm) {
                        touched.push(comm);
                    }
                }
                layer_out_weights[l][comm] += weight;
            }
            let (in_targets, in_weights) = layer.in_neighbor_slices(node);
            for i in 0..in_targets.len() {
                let neighbor = in_targets[i];
                let weight = in_weights[i];
                if neighbor == node {
                    continue;
                }
                let comm = partition.community_of(neighbor);
                if layer_out_weights[l][comm] == 0.0 && layer_in_weights[l][comm] == 0.0 {
                    if comm == current_community {
                        current_touched = true;
                    } else if !touched.contains(&comm) {
                        touched.push(comm);
                    }
                }
                layer_in_weights[l][comm] += weight;
            }
        }

        let k_v_out: f64 = layers.iter().map(|l| l.out_degree_of(node)).sum();
        let k_v_in: f64 = layers.iter().map(|l| l.in_degree_of(node)).sum();
        let node_weight = layers[0].node_weight(node);

        let sigma_tot_current_out = comm_total_degree[current_community];
        let sigma_tot_current_in = comm_in_degree[current_community];

        let (best_community, _) = find_best_community(
            touched.iter().copied(),
            current_community,
            epsilon,
            max_comm_size,
            &comm_size,
            node_weight,
            |target_comm| {
                let mut total_delta = 0.0f64;
                for (l, layer) in layers.iter().enumerate() {
                    let delta = quality.delta_move_from_components(&MoveComponents {
                        two_m: two_m_values[l],
                        node_weight,
                        total_node_weight,
                        k_v_out: layer.out_degree_of(node),
                        k_v_to_target_out: layer_out_weights[l][target_comm],
                        k_v_to_current_out: layer_out_weights[l][current_community],
                        sigma_tot_target_out: comm_total_degree[target_comm],
                        sigma_tot_current_out,
                        k_v_in: layer.in_degree_of(node),
                        k_v_to_target_in: layer_in_weights[l][target_comm],
                        k_v_to_current_in: layer_in_weights[l][current_community],
                        sigma_tot_target_in: comm_in_degree[target_comm],
                        sigma_tot_current_in,
                        n_target: comm_size[target_comm],
                        n_current: comm_size[current_community],
                        directed: layer.is_directed(),
                    });
                    total_delta += layer_weights[l] * delta;
                }
                total_delta
            },
        );

        for l in 0..num_layers {
            if current_touched {
                layer_out_weights[l][current_community] = 0.0;
                layer_in_weights[l][current_community] = 0.0;
            }
            for &comm in &touched {
                layer_out_weights[l][comm] = 0.0;
                layer_in_weights[l][comm] = 0.0;
            }
        }

        if best_community != current_community {
            apply_move(
                Some(&mut *partition),
                None,
                node,
                current_community,
                best_community,
                k_v_out,
                k_v_in,
                node_weight,
                &mut comm_total_degree,
                &mut comm_in_degree,
                &mut comm_size,
            );
            changed = true;

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

    changed
}

/// Unified refinement for single-layer and multiplex Leiden.
///
/// Refines communities by splitting them based on quality improvement.
/// Only considers neighbors within the same coarse community.
/// Uses per-layer pre-computed neighbor weights for delta computation.
///
/// Single-layer: `layers = &[data], layer_weights = &[1.0]`
/// Multi-layer: pass all layers and their weights.
pub(crate) fn refine_community_generic(
    layers: &[GraphData],
    layer_weights: &[f64],
    partition: &Partition,
    quality: &(dyn QualityFunction + Sync),
    community: usize,
    nodes: &[usize],
    epsilon: f64,
) -> Vec<(usize, usize)> {
    if nodes.len() <= 1 {
        return Vec::new();
    }

    let num_layers = layers.len();
    let total_node_weight: f64 = layers[0].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();

    // Init stats: sum degrees across layers, indexed by node (= refined community initially)
    let mut comm_total_degree: Vec<f64> = vec![0.0; max_node_id + 1];
    let mut comm_in_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] += layers.iter().map(|l| l.out_degree_of(node)).sum::<f64>();
        comm_in_degree[node] += layers.iter().map(|l| l.in_degree_of(node)).sum::<f64>();
        comm_size[node] += layers[0].node_weight(node);
    }

    // Per-layer neighbor weight buffers
    let mut layer_out_weights: Vec<Vec<f64>> = vec![vec![0.0; max_node_id + 1]; num_layers];
    let mut layer_in_weights: Vec<Vec<f64>> = vec![vec![0.0; max_node_id + 1]; num_layers];
    let mut touched: Vec<usize> = Vec::new();

    for &node in nodes {
        let current_refined = refined_map[node];

        // Reset per-layer current refined buffers
        for l in 0..num_layers {
            layer_out_weights[l][current_refined] = 0.0;
            layer_in_weights[l][current_refined] = 0.0;
        }

        // Accumulate per-layer neighbor weights (with community filter)
        for (l, layer) in layers.iter().enumerate() {
            let (targets, weights) = layer.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 layer_out_weights[l][rc] == 0.0
                    && layer_in_weights[l][rc] == 0.0
                    && rc != current_refined
                    && !touched.contains(&rc)
                {
                    touched.push(rc);
                }
                layer_out_weights[l][rc] += weight;
            }
            let (in_targets, in_weights) = layer.in_neighbor_slices(node);
            for i in 0..in_targets.len() {
                let neighbor = in_targets[i];
                let weight = in_weights[i];
                if partition.community_of(neighbor) != community {
                    continue;
                }
                if neighbor == node {
                    continue;
                }
                let rc = refined_map[neighbor];
                if layer_out_weights[l][rc] == 0.0
                    && layer_in_weights[l][rc] == 0.0
                    && rc != current_refined
                    && !touched.contains(&rc)
                {
                    touched.push(rc);
                }
                layer_in_weights[l][rc] += weight;
            }
        }

        let k_v_out: f64 = layers.iter().map(|l| l.out_degree_of(node)).sum();
        let k_v_in: f64 = layers.iter().map(|l| l.in_degree_of(node)).sum();
        let node_weight = layers[0].node_weight(node);
        let sigma_tot_current_out = comm_total_degree[current_refined];
        let sigma_tot_current_in = comm_in_degree[current_refined];

        // Find best refined community
        let (best_refined, _) = find_best_community(
            touched.iter().copied(),
            current_refined,
            epsilon,
            0, // max_comm_size = 0 (unlimited)
            &comm_size,
            node_weight,
            |target_comm| {
                let mut total_delta = 0.0f64;
                for (l, layer) in layers.iter().enumerate() {
                    let delta = quality.delta_move_from_components(&MoveComponents {
                        two_m: 2.0 * layer.total_weight(),
                        node_weight,
                        total_node_weight,
                        k_v_out: layer.out_degree_of(node),
                        k_v_to_target_out: layer_out_weights[l][target_comm],
                        k_v_to_current_out: layer_out_weights[l][current_refined],
                        sigma_tot_target_out: comm_total_degree[target_comm],
                        sigma_tot_current_out,
                        k_v_in: layer.in_degree_of(node),
                        k_v_to_target_in: layer_in_weights[l][target_comm],
                        k_v_to_current_in: layer_in_weights[l][current_refined],
                        sigma_tot_target_in: comm_in_degree[target_comm],
                        sigma_tot_current_in,
                        n_target: comm_size[target_comm],
                        n_current: comm_size[current_refined],
                        directed: layer.is_directed(),
                    });
                    total_delta += layer_weights[l] * delta;
                }
                total_delta
            },
        );

        // Clear per-layer touched buffers
        for l in 0..num_layers {
            for &rc in &touched {
                layer_out_weights[l][rc] = 0.0;
                layer_in_weights[l][rc] = 0.0;
            }
            layer_out_weights[l][current_refined] = 0.0;
            layer_in_weights[l][current_refined] = 0.0;
        }
        touched.clear();

        // Apply move if better
        if best_refined != current_refined {
            apply_move(
                None,
                Some(&mut refined_map),
                node,
                current_refined,
                best_refined,
                k_v_out,
                k_v_in,
                node_weight,
                &mut comm_total_degree,
                &mut comm_in_degree,
                &mut comm_size,
            );
        }
    }

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