leiden-rs 0.6.0

//! Multiplex/multilayer network community detection.

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::leiden::QualityType;
use crate::partition::Partition;
use crate::quality::{GraphData, Modularity, MoveComponents, QualityFunction};

/// Configuration for multiplex Leiden optimization.
#[derive(Debug, Clone)]
pub struct MultiplexConfig {
    /// Maximum number of Leiden iterations per level.
    pub max_iterations: usize,
    /// Resolution parameter γ for all layers.
    pub resolution: f64,
    /// Optional RNG seed for reproducible results.
    pub seed: Option<u64>,
    /// Quality function to optimize (applied to all layers).
    pub quality: QualityType,
    /// Convergence threshold.
    pub epsilon: f64,
    /// Maximum community size (0 = unlimited).
    pub max_comm_size: usize,
    /// Weight for each layer. Length must match the number of layers.
    /// Negative weights invert the quality (push nodes apart).
    pub layer_weights: Vec<f64>,
}

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

/// Result of multiplex Leiden optimization.
#[derive(Debug, Clone)]
pub struct MultiplexOutput {
    /// The community partition (same for all layers).
    pub partition: Partition,
    /// Total weighted quality score: Σ w_l * Q_l.
    pub quality: f64,
    /// Per-layer quality scores.
    pub layer_qualities: Vec<f64>,
}

/// Run Leiden community detection on multiple graph layers simultaneously.
///
/// All layers must have the same number of nodes. The algorithm optimizes a
/// weighted sum of quality functions: Q = Σ_l w_l * Q_l.
///
/// # Arguments
/// * `layers` - Graph data for each layer (one GraphData per layer)
/// * `config` - Multiplex configuration including layer weights
///
/// # Returns
/// A `MultiplexOutput` with the shared partition and quality scores.
pub fn run_multiplex(
    layers: &[GraphData],
    config: &MultiplexConfig,
) -> crate::error::Result<MultiplexOutput> {
    if layers.is_empty() {
        return Err(crate::error::LeidenError::InvalidParameter {
            message: "at least one layer is required".to_string(),
        });
    }

    let n = layers[0].node_count();
    for (i, layer) in layers.iter().enumerate() {
        if layer.node_count() != n {
            return Err(crate::error::LeidenError::InvalidParameter {
                message: format!(
                    "layer {} has {} nodes, expected {} (all layers must share the same vertex set)",
                    i,
                    layer.node_count(),
                    n
                ),
            });
        }
    }

    let layer_weights: Vec<f64> = if config.layer_weights.is_empty() {
        vec![1.0; layers.len()]
    } else if config.layer_weights.len() != layers.len() {
        return Err(crate::error::LeidenError::InvalidParameter {
            message: format!(
                "layer_weights has {} entries but there are {} layers",
                config.layer_weights.len(),
                layers.len()
            ),
        });
    } else {
        config.layer_weights.clone()
    };

    if n == 0 {
        return Ok(MultiplexOutput {
            partition: Partition::new(0),
            quality: 0.0,
            layer_qualities: vec![0.0; layers.len()],
        });
    }

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

    let original_n = n;
    let mut partition = Partition::new(n);
    let mut flat_mapping: Vec<usize> = (0..n).collect();

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

    {
        let q_before: f64 = layer_weights
            .iter()
            .zip(layers.iter())
            .map(|(w, layer)| w * quality.total_quality(layer, &partition))
            .sum();

        let changed = multiplex_local_moving(
            layers,
            &layer_weights,
            &mut partition,
            quality,
            &mut rng,
            config.max_comm_size,
            config.epsilon,
        );
        if changed {
            partition.renumber();

            let q_after: f64 = layer_weights
                .iter()
                .zip(layers.iter())
                .map(|(w, layer)| w * quality.total_quality(layer, &partition))
                .sum();

            if (q_after - q_before).abs() >= config.epsilon {
                let refined = multiplex_refinement(
                    layers,
                    &layer_weights,
                    &partition,
                    quality,
                    &mut rng,
                    config.epsilon,
                );

                let (agg_data, orig_to_agg, agg_initial) =
                    multiplex_aggregate(layers, &refined, &partition);

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

                partition = agg_initial;

                if agg_data.node_count() <= 1 {
                    let _ = agg_data;
                }
            }
        }
    }

    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 layer_qualities: Vec<f64> = layers
        .iter()
        .map(|layer| quality.total_quality(layer, &result))
        .collect();
    let total_quality: f64 = layer_weights
        .iter()
        .zip(layer_qualities.iter())
        .map(|(w, q)| w * q)
        .sum();

    Ok(MultiplexOutput {
        partition: result,
        quality: total_quality,
        layer_qualities,
    })
}

fn multiplex_local_moving(
    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 mut community_total_degree: Vec<f64> = vec![0.0; n];
    let mut community_size: Vec<f64> = vec![0.0; n];

    for layer in layers {
        for node in 0..n {
            let comm = partition.community_of(node);
            community_total_degree[comm] += layer.degree_of(node);
            community_size[comm] += layer.node_weight(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 total_node_weight: f64 = layers[0].total_node_weight();
    let two_m_values: Vec<f64> = layers.iter().map(|l| 2.0 * l.total_weight()).collect();

    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: f64 = layers.iter().map(|l| l.degree_of(node)).sum();
        let node_weight = layers[0].node_weight(node);

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

        for layer in layers {
            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 neighbor_comm_weights[comm] == 0.0 {
                    if comm == current_community {
                        current_touched = true;
                    } else {
                        touched.push(comm);
                    }
                }
                neighbor_comm_weights[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 in &touched {
            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 mut total_delta = 0.0f64;
            for (layer_idx, layer) in layers.iter().enumerate() {
                let mut layer_k_v_to_target = 0.0f64;
                let mut layer_k_v_to_current = 0.0f64;
                let (targets, weights) = layer.neighbor_slices(node);
                for i in 0..targets.len() {
                    let neighbor = targets[i];
                    let w = weights[i];
                    if neighbor == node {
                        continue;
                    }
                    let comm = partition.community_of(neighbor);
                    if comm == target_comm {
                        layer_k_v_to_target += w;
                    }
                    if comm == current_community {
                        layer_k_v_to_current += w;
                    }
                }

                let layer_k_v = layer.degree_of(node);
                let two_m = two_m_values[layer_idx];

                let delta = quality.delta_move_from_components(&MoveComponents {
                    k_v: layer_k_v,
                    k_v_to_target: layer_k_v_to_target,
                    k_v_to_current: layer_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,
                });
                total_delta += layer_weights[layer_idx] * delta;
            }

            if total_delta > best_delta {
                best_delta = total_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 {
            partition.move_node(node, best_community);
            community_total_degree[current_community] -= k_v;
            community_total_degree[best_community] += k_v;
            community_size[current_community] -= node_weight;
            community_size[best_community] += node_weight;
            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;
                    }
                }
            }
        }
    }

    changed
}

struct RefineContext<'a> {
    layers: &'a [GraphData],
    layer_weights: &'a [f64],
    partition: &'a Partition,
    quality: &'a (dyn QualityFunction + Sync),
    total_weight: f64,
    epsilon: f64,
}

fn multiplex_refinement(
    layers: &[GraphData],
    layer_weights: &[f64],
    partition: &Partition,
    quality: &(dyn QualityFunction + Sync),
    rng: &mut StdRng,
    epsilon: f64,
) -> Partition {
    let n = layers[0].node_count();
    let mut refined = Partition::new(n);

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

    let ctx = RefineContext {
        layers,
        layer_weights,
        partition,
        quality,
        total_weight: m,
        epsilon,
    };

    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(|(comm, nodes)| multiplex_refine_community(&ctx, comm, nodes))
                .collect()
        } else {
            community_nodes
                .iter()
                .enumerate()
                .map(|(comm, nodes)| multiplex_refine_community(&ctx, comm, nodes))
                .collect()
        }
        #[cfg(not(feature = "rayon"))]
        {
            community_nodes
                .iter()
                .enumerate()
                .map(|(comm, nodes)| multiplex_refine_community(&ctx, comm, nodes))
                .collect()
        }
    };

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

    refined
}

fn multiplex_refine_community(
    ctx: &RefineContext,
    community: usize,
    nodes: &[usize],
) -> Vec<(usize, usize)> {
    if nodes.len() <= 1 {
        return Vec::new();
    }

    let two_m = 2.0 * ctx.total_weight;
    let total_node_weight: f64 = ctx.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();

    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] += ctx.layers.iter().map(|l| l.degree_of(node)).sum::<f64>();
        comm_size[node] += ctx.layers[0].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: f64 = ctx.layers.iter().map(|l| l.degree_of(node)).sum();

        for layer in ctx.layers {
            let (targets, weights) = layer.neighbor_slices(node);
            for i in 0..targets.len() {
                let neighbor = targets[i];
                let weight = weights[i];
                if ctx.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 sigma_tot_current = comm_total_degree[current_refined];

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

        for &target_rc in &touched {
            let sigma_tot_target = comm_total_degree[target_rc];

            let total_delta: f64 = ctx
                .layers
                .iter()
                .enumerate()
                .map(|(idx, layer)| {
                    let mut layer_k_v_to_target = 0.0f64;
                    let mut layer_k_v_to_current = 0.0f64;
                    let (targets, weights) = layer.neighbor_slices(node);
                    for i in 0..targets.len() {
                        let nb = targets[i];
                        let w = weights[i];
                        if ctx.partition.community_of(nb) != community || nb == node {
                            continue;
                        }
                        let rc = refined_map[nb];
                        if rc == target_rc {
                            layer_k_v_to_target += w;
                        }
                        if rc == current_refined {
                            layer_k_v_to_current += w;
                        }
                    }
                    let delta = ctx.quality.delta_move_from_components(&MoveComponents {
                        k_v: layer.degree_of(node),
                        k_v_to_target: layer_k_v_to_target,
                        k_v_to_current: layer_k_v_to_current,
                        sigma_tot_target,
                        sigma_tot_current,
                        two_m,
                        n_target: comm_size[target_rc],
                        n_current: comm_size[current_refined],
                        node_weight: ctx.layers[0].node_weight(node),
                        total_node_weight,
                    });
                    ctx.layer_weights[idx] * delta
                })
                .sum();

            if total_delta > best_delta {
                best_delta = total_delta;
                best_refined = target_rc;
            }
        }

        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 = ctx.layers[0].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()
        .map(|&node| (node, refined_map[node]))
        .collect()
}

fn multiplex_aggregate(
    layers: &[GraphData],
    refined_partition: &Partition,
    coarse_partition: &Partition,
) -> (GraphData, Vec<usize>, Partition) {
    let n = layers[0].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 layer in layers {
        for u in 0..n {
            let ru = orig_to_agg[u];
            for (v, w) in layer.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;
            }
        }
    }

    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;
        }
    }

    for u in 0..agg_n {
        let start = agg_offsets[u];
        let end = agg_offsets[u + 1];
        if end - start <= 1 {
            continue;
        }
        let slice_t = &mut agg_targets[start..end];
        let slice_w = &mut agg_weights[start..end];
        let len = slice_t.len();
        let mut indices: Vec<usize> = (0..len).collect();
        indices.sort_by_key(|&i| slice_t[i]);
        let sorted_t: Vec<usize> = indices.iter().map(|&i| slice_t[i]).collect();
        let sorted_w: Vec<f64> = indices.iter().map(|&i| slice_w[i]).collect();
        slice_t.copy_from_slice(&sorted_t);
        slice_w.copy_from_slice(&sorted_w);
    }

    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() {
        let total_degree: f64 = layers.iter().map(|l| l.degree_of(orig)).sum();
        agg_degree[agg_node] += total_degree;
        agg_node_weight[agg_node] += layers[0].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 crate::quality::GraphData;

    #[test]
    fn test_multiplex_two_layers() {
        let edges1 = [
            (0, 1, 1.0),
            (1, 2, 1.0),
            (0, 2, 1.0),
            (3, 4, 1.0),
            (4, 5, 1.0),
            (3, 5, 1.0),
            (2, 3, 1.0),
        ];
        let layer1 = GraphData::from_edgelist(&edges1, 6).unwrap();

        let edges2 = [
            (0, 1, 1.0),
            (1, 2, 1.0),
            (0, 2, 1.0),
            (3, 4, 1.0),
            (4, 5, 1.0),
            (3, 5, 1.0),
            (1, 4, 1.0),
        ];
        let layer2 = GraphData::from_edgelist(&edges2, 6).unwrap();

        let config = MultiplexConfig {
            seed: Some(42),
            layer_weights: vec![1.0, 1.0],
            ..Default::default()
        };

        let result = run_multiplex(&[layer1, layer2], &config).unwrap();
        assert!(result.partition.num_communities() >= 1);
        assert_eq!(result.layer_qualities.len(), 2);
    }

    #[test]
    fn test_multiplex_single_layer_matches_standard() {
        let edges = [
            (0, 1, 1.0),
            (1, 2, 1.0),
            (0, 2, 1.0),
            (3, 4, 1.0),
            (4, 5, 1.0),
            (3, 5, 1.0),
            (2, 3, 1.0),
        ];
        let layer = GraphData::from_edgelist(&edges, 6).unwrap();

        let config = MultiplexConfig {
            seed: Some(42),
            layer_weights: vec![1.0],
            ..Default::default()
        };

        let result = run_multiplex(&[layer], &config).unwrap();
        assert!(result.partition.num_communities() >= 1);
    }

    #[test]
    fn test_multiplex_mismatched_nodes() {
        let layer1 = GraphData::from_edgelist(&[(0, 1, 1.0)], 2).unwrap();
        let layer2 = GraphData::from_edgelist(&[(0, 1, 1.0), (1, 2, 1.0)], 3).unwrap();

        let config = MultiplexConfig {
            layer_weights: vec![1.0, 1.0],
            ..Default::default()
        };

        assert!(run_multiplex(&[layer1, layer2], &config).is_err());
    }

    #[test]
    fn test_multiplex_empty_layers() {
        let config = MultiplexConfig::default();
        assert!(run_multiplex(&[], &config).is_err());
    }

    #[test]
    fn test_multiplex_weighted_layers() {
        let edges1 = [
            (0, 1, 1.0),
            (1, 2, 1.0),
            (0, 2, 1.0),
            (3, 4, 1.0),
            (4, 5, 1.0),
            (3, 5, 1.0),
            (2, 3, 0.01),
        ];
        let edges2 = [(0, 3, 1.0), (1, 4, 1.0), (2, 5, 1.0)];
        let layer1 = GraphData::from_edgelist(&edges1, 6).unwrap();
        let layer2 = GraphData::from_edgelist(&edges2, 6).unwrap();

        let config = MultiplexConfig {
            seed: Some(42),
            layer_weights: vec![10.0, 0.1],
            ..Default::default()
        };

        let result = run_multiplex(&[layer1, layer2], &config).unwrap();
        assert!(
            result.partition.num_communities() >= 2,
            "expected >= 2 communities, got {}",
            result.partition.num_communities()
        );
    }
}