//! Graph community detection algorithms. Includes Louvain and WalkTrap
//! implementations.

use rand::prelude::*;
use rayon::prelude::*;
use rustc_hash::{FxHashMap, FxHashSet};
use std::collections::BinaryHeap;
use std::collections::VecDeque;
use std::time::Instant;

use crate::core::math::sparse::coo_to_csr;
use crate::prelude::*;

/////////////////////
// General helpers //
/////////////////////

/// Helper function to remap communities by size
///
/// The non-determinism due to threading etc. can be a huge annoyance in
/// single cell. This function remaps the clusters by size, ensuring more (not
/// complete) reproducibility.
///
/// ### Params
///
/// * `labels` - The membership labels
///
/// ### Returns
///
/// Remapped labels with `0` -> biggest community by membership, `1` -> second
/// biggest, etc.
fn remap_communities_by_size(labels: &[usize]) -> Vec<usize> {
    let n = labels.len();
    let mut counts: FxHashMap<usize, usize> = FxHashMap::default();
    for &l in labels {
        *counts.entry(l).or_insert(0) += 1;
    }

    let mut sorted: Vec<(usize, usize)> = counts.into_iter().collect();
    sorted.sort_unstable_by_key(|a| std::cmp::Reverse(a.1));

    let mut remap: FxHashMap<usize, usize> = FxHashMap::default();
    for (new_label, (old_label, _)) in sorted.into_iter().enumerate() {
        remap.insert(old_label, new_label);
    }

    (0..n).map(|i| remap[&labels[i]]).collect()
}

/////////////
// Louvain //
/////////////

/////////////
// Helpers //
/////////////

/// One Phase 1 pass with pruning.
///
/// Maintains a queue of dirty nodes. A node is dirty when first added or when
/// one of its neighbours has just moved community. Drains until empty.
///
/// ### Params
///
/// * `graph` - The [`SparseGraph`] - must be undirected!
/// * `resolution` - The resolution parameter to use
/// * `max_iter` - The maximum iterations to run the algorithm for. In this
///   case, it safeguards against degenerated graphs and should not trigger.
/// * `rng` - The random number generator
///
/// ### Returns
///
/// Returns `(contiguous labels in [0, n_distinct), count)`.
fn louvain_one_level<T>(
    graph: &SparseGraph<T>,
    resolution: T,
    max_iter: usize,
    rng: &mut StdRng,
) -> (Vec<usize>, usize)
where
    T: BixverseFloat + BixverseNumeric + Clone + std::iter::Sum,
{
    let n = graph.get_node_number();

    let m: T = (0..n)
        .map(|i| graph.get_neighbours(i).1.iter().copied().sum::<T>())
        .sum::<T>()
        / T::from_f64(2.0).unwrap();

    let two_m = T::from_f64(2.0).unwrap() * m;
    let res_over_two_m = resolution / two_m;

    let mut degrees = vec![T::zero(); n];
    for i in 0..n {
        degrees[i] = graph.get_neighbours(i).1.iter().copied().sum();
    }

    let mut communities: Vec<u32> = (0..n as u32).collect();
    let mut comm_degree_sums = degrees.clone();
    let mut neighbour_weights = vec![T::zero(); n];
    let mut comm_active = vec![false; n];
    let mut active_comms = Vec::with_capacity(256);

    let epsilon = T::from_f64(1e-10).unwrap();

    let mut initial_order: Vec<u32> = (0..n as u32).collect();
    initial_order.shuffle(rng);
    let mut queue: VecDeque<u32> = initial_order.into_iter().collect();
    let mut in_queue = vec![true; n];

    let max_evals = max_iter.saturating_mul(n);
    let mut evals = 0usize;

    while let Some(node) = queue.pop_front() {
        if evals >= max_evals {
            break;
        }
        evals += 1;

        let node_idx = node as usize;
        in_queue[node_idx] = false;

        let current_comm = communities[node_idx] as usize;
        let k_i = degrees[node_idx];
        let k_i_scaled = k_i * res_over_two_m;

        let (neighbours, weights) = graph.get_neighbours(node_idx);

        for (&neighbour, &weight) in neighbours.iter().zip(weights.iter()) {
            let comm = communities[neighbour] as usize;
            if !comm_active[comm] {
                comm_active[comm] = true;
                active_comms.push(comm);
            }
            neighbour_weights[comm] += weight;
        }

        let mut best_comm = current_comm;
        let mut best_delta = T::zero();

        for &comm in &active_comms {
            if comm != current_comm {
                let delta = neighbour_weights[comm] - k_i_scaled * comm_degree_sums[comm];
                if delta > best_delta {
                    best_delta = delta;
                    best_comm = comm;
                }
            }
        }

        for &comm in &active_comms {
            neighbour_weights[comm] = T::zero();
            comm_active[comm] = false;
        }
        active_comms.clear();

        if best_comm != current_comm && best_delta > epsilon {
            communities[node_idx] = best_comm as u32;
            comm_degree_sums[current_comm] -= k_i;
            comm_degree_sums[best_comm] += k_i;

            for nb in neighbours {
                if nb != node_idx && !in_queue[nb] {
                    queue.push_back(nb as u32);
                    in_queue[nb] = true;
                }
            }
        }
    }

    let mut comm_map = vec![u32::MAX; n];
    let mut label = 0u32;
    for c in &mut communities {
        let idx = *c as usize;
        if comm_map[idx] == u32::MAX {
            comm_map[idx] = label;
            label += 1;
        }
        *c = comm_map[idx];
    }

    let n_distinct = label as usize;
    (
        communities.iter().map(|&c| c as usize).collect(),
        n_distinct,
    )
}

/// Aggregate fine nodes by community into a coarser graph.
///
/// Internal community edges become self-loops on the super-node (with weight
/// 2x the internal edge weight, as both CSR directions accumulate). External
/// edges sum across all member-to-member connections. Total weight m is
/// preserved across aggregation.
///
/// ### Params
///
/// * `graph` - The [`SparseGraph`] - must be undirected!
/// * `communities` - The community membership from the finer level.
/// * `n_comms` - Number of communities
///
/// ### Returns
///
/// Coarsed sparse graph
fn aggregate_graph<T>(
    graph: &SparseGraph<T>,
    communities: &[usize],
    n_comms: usize,
) -> SparseGraph<T>
where
    T: BixverseFloat + BixverseNumeric + Clone + std::iter::Sum,
{
    let mut nodes_per_comm: Vec<Vec<usize>> = vec![Vec::new(); n_comms];
    for (i, &c) in communities.iter().enumerate() {
        nodes_per_comm[c].push(i);
    }

    let mut acc = vec![T::zero(); n_comms];
    let mut touched = vec![false; n_comms];
    let mut active: Vec<usize> = Vec::with_capacity(64);

    let mut rows = Vec::new();
    let mut cols = Vec::new();
    let mut vals = Vec::new();

    for ci in 0..n_comms {
        for &fine_i in &nodes_per_comm[ci] {
            let (neighbours, weights) = graph.get_neighbours(fine_i);
            for (&j, &w) in neighbours.iter().zip(weights.iter()) {
                let cj = communities[j];
                if !touched[cj] {
                    touched[cj] = true;
                    active.push(cj);
                }
                acc[cj] += w;
            }
        }

        for &cj in &active {
            rows.push(ci);
            cols.push(cj);
            vals.push(acc[cj]);
            acc[cj] = T::zero();
            touched[cj] = false;
        }
        active.clear();
    }

    let csr = coo_to_csr(
        &rows.index_cast(),
        &cols.index_cast(),
        &vals,
        (n_comms, n_comms),
    );
    SparseGraph::new(n_comms, csr, false)
}

//////////
// Main //
//////////

/// Louvain community detection on a [`SparseGraph`].
///
/// ### Params
///
/// * `graph` - The [`SparseGraph`]
/// * `resolution` - Resolution parameter for the Louvain clustering
/// * `max_iter` - The maximum iterations to run the algorithm for. In this
///   case, it safeguards against degenerated graphs and should not trigger.
/// * `multi_level` - If `true`, run the full multi-level Louvain (Phase 1 +
///   Phase 2 aggregation, repeated). If `false`, run only one Phase 1 pass
///   (matches Phenograph / the original doubletdetection behaviour).
/// * `seed` - Seed for reproducibility purposes
///
/// ### Returns
///
/// Vector of communities
pub fn louvain_sparse_graph<T>(
    graph: &SparseGraph<T>,
    resolution: T,
    max_iter: usize,
    multi_level: bool,
    seed: usize,
) -> Result<Vec<usize>, BixverseErrors>
where
    T: BixverseFloat + BixverseNumeric + Clone + std::iter::Sum,
{
    if graph.is_directed() {
        return Err(BixverseErrors::GraphDirectedError);
    }

    let n_orig = graph.get_node_number();
    if n_orig == 0 {
        return Ok(Vec::new());
    }

    let mut rng = StdRng::seed_from_u64(seed as u64);
    let mut node_to_super: Vec<usize> = (0..n_orig).collect();

    let (comms, mut n_super) = louvain_one_level(graph, resolution, max_iter, &mut rng);
    for s in node_to_super.iter_mut() {
        *s = comms[*s];
    }

    if !multi_level || n_super == n_orig {
        return Ok(remap_communities_by_size(&node_to_super));
    }

    let mut current_graph = aggregate_graph(graph, &comms, n_super);

    loop {
        let (comms, new_n_super) =
            louvain_one_level(&current_graph, resolution, max_iter, &mut rng);
        for s in node_to_super.iter_mut() {
            *s = comms[*s];
        }
        if new_n_super == n_super {
            break;
        }
        n_super = new_n_super;
        current_graph = aggregate_graph(&current_graph, &comms, n_super);
    }

    Ok(remap_communities_by_size(&node_to_super))
}

//////////////
// Walktrap //
//////////////

/// Squared Euclidean distance between two index-sorted sparse vectors.
///
/// Iterates over the union of indices in a single pass; missing indices are
/// treated as zero.
///
/// ### Params
///
/// * `a` - First index-sorted sparse vector as `(index, value)` pairs
/// * `b` - Second index-sorted sparse vector as `(index, value)` pairs
///
/// ### Returns
///
/// Squared Euclidean distance between the two vectors
fn sparse_sq_dist<T>(a: &[(usize, T)], b: &[(usize, T)]) -> T
where
    T: BixverseFloat,
{
    let mut acc = T::zero();
    let mut i = 0;
    let mut j = 0;
    while i < a.len() && j < b.len() {
        let (ai, av) = a[i];
        let (bi, bv) = b[j];
        if ai == bi {
            let d = av - bv;
            acc += d * d;
            i += 1;
            j += 1;
        } else if ai < bi {
            acc += av * av;
            i += 1;
        } else {
            acc += bv * bv;
            j += 1;
        }
    }
    while i < a.len() {
        let v = a[i].1;
        acc += v * v;
        i += 1;
    }
    while j < b.len() {
        let v = b[j].1;
        acc += v * v;
        j += 1;
    }
    acc
}

/// Weighted average of two index-sorted sparse vectors:
/// `(n_a · a + n_b · b) / (n_a + n_b)`, returned index-sorted.
///
/// Iterates over the union of indices in a single pass; missing indices are
/// treated as zero.
///
/// ### Params
///
/// * `a` - First index-sorted sparse vector as `(index, value)` pairs
/// * `b` - Second index-sorted sparse vector as `(index, value)` pairs
/// * `n_a` - Weight (typically community size) for vector `a`
/// * `n_b` - Weight (typically community size) for vector `b`
///
/// ### Returns
///
/// Merged index-sorted sparse vector
fn merge_sparse_vecs<T>(
    a: &[(usize, T)],
    b: &[(usize, T)],
    n_a: usize,
    n_b: usize,
) -> Vec<(usize, T)>
where
    T: BixverseFloat,
{
    let na = T::from_usize(n_a).unwrap();
    let nb = T::from_usize(n_b).unwrap();
    let nc = na + nb;
    let mut out = Vec::with_capacity(a.len() + b.len());
    let mut i = 0;
    let mut j = 0;
    while i < a.len() && j < b.len() {
        let (ai, av) = a[i];
        let (bi, bv) = b[j];
        if ai == bi {
            out.push((ai, (na * av + nb * bv) / nc));
            i += 1;
            j += 1;
        } else if ai < bi {
            out.push((ai, na * av / nc));
            i += 1;
        } else {
            out.push((bi, nb * bv / nc));
            j += 1;
        }
    }
    while i < a.len() {
        let (ai, av) = a[i];
        out.push((ai, na * av / nc));
        i += 1;
    }
    while j < b.len() {
        let (bi, bv) = b[j];
        out.push((bi, nb * bv / nc));
        j += 1;
    }
    out
}

/// Index-sort a sparse vector, optionally keeping only its `k` largest-value
/// entries first.
///
/// `None` leaves the support untruncated. `Some(k)` keeps the top-`k` entries by
/// mass (the dropped tail holds negligible walk mass), bounding stored support
/// at `k` per community vector.
fn finalise_sparse_vec<T>(mut v: Vec<(usize, T)>, max_support: Option<usize>) -> Vec<(usize, T)>
where
    T: BixverseFloat,
{
    if let Some(k) = max_support
        && v.len() > k
    {
        v.select_nth_unstable_by(k, |a, b| {
            b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
        });
        v.truncate(k);
        v.shrink_to_fit(); // important... I need to reduce the allocations here...
    }
    v.sort_unstable_by_key(|&(idx, _)| idx);
    v
}

/// Compute random-walk probability vectors for all nodes (sparse)
///
/// For each node i, simulates a `walk_length`-step random walk starting at i
/// and records the landing probabilities, scaled by the inverse square root of
/// each node's degree (Pons & Latapy 2005). Probabilities are propagated through
/// hashmaps so only reached nodes are stored.
///
/// Transition probabilities are weight-proportional: `p(i→j) = w(i,j) / deg(i)`.
/// Isolated nodes (zero degree) yield an effectively zero vector.
///
/// ### Params
///
/// * `graph` - Sparse weighted graph
/// * `walk_length` - Number of steps per random walk
/// * `max_support` - Optional usize. If provided, the only the top max support
///   entries are kept and the detail is removed. This can help with large
///   memory data sets.
///
/// ### Returns
///
/// One index-sorted sparse vector per node, where entry `(j, v)` holds
/// `v = P(walk from i lands at j) / sqrt(deg(j))`.
fn compute_walk_probabilities<T>(
    graph: &SparseGraph<T>,
    walk_length: usize,
    max_support: Option<usize>,
) -> Vec<Vec<(usize, T)>>
where
    T: BixverseFloat + std::iter::Sum + Send + Sync,
{
    let n = graph.get_node_number();
    let epsilon = T::from_f64(1e-10).unwrap();

    let inv_sqrt_deg: Vec<T> = (0..n)
        .map(|k| {
            let (_, w) = graph.get_neighbours(k);
            let d: T = w.iter().copied().sum();
            if d > T::zero() {
                T::one() / d.sqrt()
            } else {
                T::zero()
            }
        })
        .collect();

    let transition_probs: Vec<Vec<(usize, T)>> = (0..n)
        .into_par_iter()
        .map(|i| {
            let (neighbours, weights) = graph.get_neighbours(i);
            let degree: T = weights.iter().copied().sum();
            if degree > T::zero() {
                neighbours
                    .iter()
                    .zip(weights.iter())
                    .map(|(&j, &w)| (j, w / degree))
                    .collect()
            } else {
                Vec::new()
            }
        })
        .collect();

    (0..n)
        .into_par_iter()
        .map(|i| {
            let mut probs: FxHashMap<usize, T> = FxHashMap::default();
            probs.insert(i, T::one());
            for _ in 0..walk_length {
                let mut next: FxHashMap<usize, T> = FxHashMap::default();
                for (&node, &p) in &probs {
                    if p > epsilon {
                        for &(nb, tp) in &transition_probs[node] {
                            *next.entry(nb).or_insert_with(T::zero) += p * tp;
                        }
                    }
                }
                probs = next;
            }
            let row: Vec<(usize, T)> = probs
                .into_iter()
                .map(|(k, p)| (k, p * inv_sqrt_deg[k]))
                .collect();
            finalise_sparse_vec(row, max_support)
        })
        .collect()
}

/// Compute Ward's linkage criterion between two communities
///
/// Returns the increase in total within-cluster variance that would result from
/// merging communities A and B, given their squared centroid distance:
///
/// ```Ward(A, B) = (n_a · n_b) / (n_a + n_b) · d²(A, B)```
///
/// ### Params
///
/// * `d_sq` - Squared distance between the community centroids
/// * `n_a` - Size of community A
/// * `n_b` - Size of community B
///
/// ### Returns
///
/// Ward criterion value (lower means a cheaper merge)
#[inline(always)]
fn ward_criterion<T>(d_sq: T, n_a: usize, n_b: usize) -> T
where
    T: BixverseFloat,
{
    let na = T::from_usize(n_a).unwrap();
    let nb = T::from_usize(n_b).unwrap();
    (na * nb) / (na + nb) * d_sq
}

/// Resolve final community labels from a union-find parent array
///
/// Follows each node's parent chain to its root and assigns a contiguous
/// integer label `0..num_communities` in traversal order. The result is a flat
/// label per node with no gaps.
///
/// ### Params
///
/// * `n` - Number of original nodes
/// * `parent` - Union-find parent array; `parent[x] == x` denotes a root
///
/// ### Returns
///
/// Vector of community labels, one per node
fn finalise_labels(n: usize, parent: &[usize]) -> Vec<usize> {
    let mut labels = vec![0usize; n];
    let mut label_map: FxHashMap<usize, usize> = FxHashMap::default();
    let mut next = 0usize;
    for node in 0..n {
        let mut x = node;
        while parent[x] != x {
            x = parent[x];
        }
        let label = *label_map.entry(x).or_insert_with(|| {
            let l = next;
            next += 1;
            l
        });
        labels[node] = label;
    }
    labels
}

/// Walktrap community detection (Pons & Latapy 2005)
///
/// Identifies communities by agglomerative clustering of random-walk probability
/// vectors. Nodes that tend to co-occur on short random walks are grouped
/// together.
///
/// Algorithm outline:
///
/// 1. Compute scaled walk-probability vectors for all nodes (sparse)
/// 2. Initialise each node as its own community
/// 3. Greedily merge the adjacent pair with the smallest Ward criterion
///    until `num_clusters` communities remain
/// 4. Resolve labels via union-find
///
/// Merged community vectors are updated as the weighted average:
/// `q[A∪B] = (n_a · q[A] + n_b · q[B]) / (n_a + n_b)`
///
/// ### Params
///
/// * `graph` - Sparse weighted graph to cluster
/// * `walk_length` - Number of steps for the random walks
/// * `num_clusters` - Target number of communities to return
/// * `max_support` - Optional usize. If provided, the only the top max support
///   entries are kept and the detail is removed. This can help with large
///   memory data sets and the tail tends to hold neglible mass. This becomes
///   an approximation in this case.
/// * `verbose` - Print timing information for each stage
///
/// ### Returns
///
/// Vector of community labels (0-indexed, contiguous), one per node
pub fn walktrap_sparse_graph<T>(
    graph: &SparseGraph<T>,
    walk_length: usize,
    num_clusters: usize,
    max_support: Option<usize>,
    verbose: bool,
) -> Vec<usize>
where
    T: BixverseFloat + std::iter::Sum + Send + Sync,
{
    let n = graph.get_node_number();
    if n == 0 {
        return Vec::new();
    }
    if n <= num_clusters {
        return (0..n).collect();
    }

    let walktrap_start = Instant::now();

    if verbose {
        match max_support {
            Some(ms) => println!(
                "Calculating walk probabilities with max support of {:?}.",
                ms
            ),
            None => println!("Calculating walk probabilities."),
        }
    }

    let t0 = Instant::now();
    let q = compute_walk_probabilities(graph, walk_length, max_support);
    if verbose {
        println!(" ... done in {:.2?}", t0.elapsed());
    }

    let mut comm_q: Vec<Option<Vec<(usize, T)>>> = q.into_iter().map(Some).collect();

    let mut comm_size: Vec<usize> = vec![1; n];
    let mut active: Vec<bool> = vec![true; n];
    let mut comm_parent: Vec<usize> = (0..n).collect();
    let mut adj: Vec<FxHashSet<usize>> = vec![FxHashSet::default(); n];

    let mut edges: Vec<(usize, usize)> = Vec::new();
    for i in 0..n {
        let (neighbours, _) = graph.get_neighbours(i);
        for j in neighbours {
            if i != j {
                adj[i].insert(j);
                if i < j {
                    edges.push((i, j));
                }
            }
        }
    }

    let t1 = Instant::now();
    let initial: Vec<(RevOrderedFloat<T>, usize, usize)> = edges
        .par_iter()
        .map(|&(i, j)| {
            let qi = comm_q[i].as_ref().unwrap();
            let qj = comm_q[j].as_ref().unwrap();
            let d_sq = sparse_sq_dist(qi, qj);
            let crit = ward_criterion(d_sq, 1, 1);
            (RevOrderedFloat(crit), i, j)
        })
        .collect();

    if verbose {
        println!("Computing initial distances.")
    }

    let mut heap: BinaryHeap<(RevOrderedFloat<T>, usize, usize)> =
        BinaryHeap::with_capacity(initial.len() * 2);
    for entry in initial {
        heap.push(entry);
    }
    if verbose {
        println!(" ... done in {:.2?}", t1.elapsed());
    }

    let t2 = Instant::now();
    let mut num_active = n;

    while num_active > num_clusters {
        let (a, b) = loop {
            match heap.pop() {
                Some((_, i, j)) => {
                    if active[i] && active[j] {
                        break (i, j);
                    }
                }
                None => {
                    if verbose {
                        println!(
                            "Heap exhausted with {} communities active (graph likely disconnected)",
                            num_active
                        );
                        println!("Finished merging communities: {:.2?}", t2.elapsed());
                        println!(
                            "Finished WalkTrap community detection: {:.2?}",
                            walktrap_start.elapsed()
                        );
                    }
                    return finalise_labels(n, &comm_parent);
                }
            }
        };

        let n_a = comm_size[a];
        let n_b = comm_size[b];
        let n_c = n_a + n_b;

        let new_q = {
            let qa = comm_q[a].as_ref().unwrap();
            let qb = comm_q[b].as_ref().unwrap();
            merge_sparse_vecs(qa, qb, n_a, n_b)
        };

        let mut new_adj: FxHashSet<usize> = FxHashSet::default();
        for &k in &adj[a] {
            if k != b {
                new_adj.insert(k);
            }
        }
        for &k in &adj[b] {
            if k != a {
                new_adj.insert(k);
            }
        }

        active[a] = false;
        active[b] = false;
        comm_q[a] = None;
        comm_q[b] = None;
        adj[a] = FxHashSet::default();
        adj[b] = FxHashSet::default();

        let c = comm_q.len();
        comm_q.push(Some(new_q));
        comm_size.push(n_c);
        active.push(true);
        comm_parent.push(c);
        comm_parent[a] = c;
        comm_parent[b] = c;
        adj.push(new_adj.clone());

        for &k in &new_adj {
            adj[k].remove(&a);
            adj[k].remove(&b);
            adj[k].insert(c);
        }

        let new_adj_vec: Vec<usize> = new_adj.into_iter().collect();
        let new_criteria: Vec<(RevOrderedFloat<T>, usize, usize)> = new_adj_vec
            .par_iter()
            .map(|&k| {
                let qc = comm_q[c].as_ref().unwrap();
                let qk = comm_q[k].as_ref().unwrap();
                let d_sq = sparse_sq_dist(qc, qk);
                let crit = ward_criterion(d_sq, n_c, comm_size[k]);
                (RevOrderedFloat(crit), c, k)
            })
            .collect();

        for entry in new_criteria {
            heap.push(entry);
        }

        num_active -= 1;
    }

    if verbose {
        println!("Finished merging communities: {:.2?}", t2.elapsed());
        println!(
            "Finished WalkTrap community detection: {:.2?}",
            walktrap_start.elapsed()
        );
    }

    let labels = finalise_labels(n, &comm_parent);

    remap_communities_by_size(&labels)
}

///////////
// Tests //
///////////

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::math::sparse::CompressedSparseData2;

    fn build_barbell_graph() -> SparseGraph<f64> {
        let indptr = vec![0, 2, 4, 7, 10, 12, 14];
        let indices = vec![1, 2, 0, 2, 0, 1, 3, 2, 4, 5, 3, 5, 3, 4];
        let data = vec![1.0; 14];
        let csr =
            CompressedSparseData2::<f64, f64>::new_csr(&data, &indices, &indptr, None, (6, 6));
        SparseGraph::new(6, csr, false)
    }

    fn build_two_cliques_graph() -> SparseGraph<f64> {
        // K4 {0..3} + K4 {4..7} bridged by edge 3-4
        let mut nbrs: Vec<Vec<usize>> = vec![Vec::new(); 8];
        for i in 0..4 {
            for j in 0..4 {
                if i != j {
                    nbrs[i].push(j);
                }
            }
        }
        for i in 4..8 {
            for j in 4..8 {
                if i != j {
                    nbrs[i].push(j);
                }
            }
        }
        nbrs[3].push(4);
        nbrs[4].push(3);

        let mut indptr = vec![0usize];
        let mut indices = Vec::new();
        let mut data = Vec::new();
        for list in &mut nbrs {
            list.sort();
            for &j in list.iter() {
                indices.push(j);
                data.push(1.0);
            }
            indptr.push(indices.len());
        }

        let csr = CompressedSparseData2::<f64, f64>::new_csr(
            &data,
            &indices.index_cast(),
            &indptr.index_cast(),
            None,
            (8, 8),
        );
        SparseGraph::new(8, csr, false)
    }

    #[test]
    fn test_louvain_barbell() {
        let graph = build_barbell_graph();
        let comms = louvain_sparse_graph(&graph, 1.0, 10, true, 42).unwrap();
        assert_eq!(comms[0], comms[1]);
        assert_eq!(comms[1], comms[2]);
        assert_eq!(comms[3], comms[4]);
        assert_eq!(comms[4], comms[5]);
        assert_ne!(comms[2], comms[3]);
    }

    #[test]
    fn test_walktrap_barbell() {
        let graph = build_barbell_graph();
        let comms = walktrap_sparse_graph(&graph, 3, 2, None, false);
        assert_eq!(comms[0], comms[1]);
        assert_eq!(comms[1], comms[2]);
        assert_eq!(comms[3], comms[4]);
        assert_eq!(comms[4], comms[5]);
        assert_ne!(comms[2], comms[3]);
    }

    #[test]
    fn test_walktrap_two_cliques() {
        let graph = build_two_cliques_graph();
        let comms = walktrap_sparse_graph(&graph, 4, 2, None, false);
        for i in 1..4 {
            assert_eq!(comms[i], comms[0]);
        }
        for i in 5..8 {
            assert_eq!(comms[i], comms[4]);
        }
        assert_ne!(comms[0], comms[4]);
    }
}
