//! Louvain community detection on weighted undirected graphs.
//!
//! Finds natural communities by maximizing modularity. Supports weighted edges
//! (cosine similarity between leaf centroids) and a resolution parameter for
//! controlling granularity.
//!
//! Ported from `boids-wasm/src/louvain.rs` with enhancements:
//! - Weighted edges (f32 weights instead of unweighted)
//! - Resolution parameter (controls community granularity)
//! - Deterministic node ordering for reproducibility

/// Single-level Louvain community detection on a weighted undirected graph.
///
/// # Arguments
/// * `n` - Number of nodes
/// * `adj` - Weighted adjacency list: `adj[i]` contains `(neighbor_index, edge_weight)` pairs
/// * `resolution` - Modularity resolution parameter (1.0 = standard, >1 = more/smaller communities)
///
/// # Returns
/// * `(community_ids, num_communities)` - IDs compacted and sorted by size (largest = 0)
pub fn louvain_communities(
    n: usize,
    adj: &[Vec<(u32, f32)>],
    resolution: f32,
) -> (Vec<u32>, u32) {
    if n == 0 {
        return (Vec::new(), 0);
    }

    let mut community: Vec<u32> = (0..n as u32).collect();

    // Weighted degree of each node (sum of edge weights)
    let deg: Vec<f64> = adj
        .iter()
        .map(|a| a.iter().map(|(_, w)| *w as f64).sum())
        .collect();
    let m: f64 = deg.iter().sum::<f64>() / 2.0;
    if m < 1e-12 {
        return (vec![0; n], 1);
    }
    let inv_m = 1.0 / m;
    let inv_2m2 = 1.0 / (2.0 * m * m);
    let gamma = resolution as f64;

    // Sum of weighted degrees for each community
    let mut sigma_tot: Vec<f64> = deg.clone();

    // Scratch buffer for per-node community edge counts
    let mut comm_edges: Vec<(u32, f64)> = Vec::with_capacity(16);

    // Deterministic ordering: process nodes in index order
    let mut pass = 0u32;
    loop {
        let mut moves = 0u32;
        pass += 1;

        for i in 0..n {
            let ci = community[i];
            let ki = deg[i];
            if ki == 0.0 {
                continue;
            }

            // Count weighted edges from i to each neighboring community
            comm_edges.clear();
            let mut ki_own: f64 = 0.0;

            for &(j, w) in &adj[i] {
                let cj = community[j as usize];
                let w64 = w as f64;
                if cj == ci {
                    ki_own += w64;
                } else if let Some(entry) = comm_edges.iter_mut().find(|(c, _)| *c == cj) {
                    entry.1 += w64;
                } else {
                    comm_edges.push((cj, w64));
                }
            }

            // Cost of removing i from its current community
            // ΔQ_remove = ki_in/m - γ·ki·(Σ_tot - ki)/(2m²)
            let remove_cost = ki_own * inv_m - gamma * ki * (sigma_tot[ci as usize] - ki) * inv_2m2;

            let mut best_comm = ci;
            let mut best_gain: f64 = 0.0;

            // Sort neighbor communities for deterministic tie-breaking
            comm_edges.sort_by_key(|&(c, _)| c);

            for &(cj, ki_cj) in &comm_edges {
                // ΔQ_insert = ki_cj/m - γ·ki·Σ_tot[cj]/(2m²)
                let gain = ki_cj * inv_m - gamma * ki * sigma_tot[cj as usize] * inv_2m2;
                let delta = gain - remove_cost;
                if delta > best_gain {
                    best_gain = delta;
                    best_comm = cj;
                }
            }

            if best_comm != ci {
                sigma_tot[ci as usize] -= ki;
                sigma_tot[best_comm as usize] += ki;
                community[i] = best_comm;
                moves += 1;
            }
        }

        if moves == 0 || pass >= 20 {
            break;
        }
    }

    compact_by_size(&mut community)
}

/// Compact community IDs and sort by size (largest community = 0).
fn compact_by_size(community: &mut [u32]) -> (Vec<u32>, u32) {
    let n = community.len();
    if n == 0 {
        return (Vec::new(), 0);
    }

    // Count members per raw community ID
    let max_raw = community.iter().copied().max().unwrap_or(0) as usize;
    let mut counts = vec![0u32; max_raw + 1];
    for &c in community.iter() {
        counts[c as usize] += 1;
    }

    // Collect non-empty communities with their raw IDs
    let mut communities: Vec<(u32, u32)> = counts
        .iter()
        .enumerate()
        .filter(|(_, &count)| count > 0)
        .map(|(raw_id, &count)| (raw_id as u32, count))
        .collect();

    // Sort by count descending, then by raw_id ascending for determinism
    communities.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));

    // Build mapping: raw_id -> new_id
    let mut mapping = vec![0u32; max_raw + 1];
    for (new_id, &(raw_id, _)) in communities.iter().enumerate() {
        mapping[raw_id as usize] = new_id as u32;
    }

    // Apply mapping
    let result: Vec<u32> = community.iter().map(|&c| mapping[c as usize]).collect();
    let num = communities.len() as u32;

    (result, num)
}

/// Build a weighted adjacency list from leaf centroids using KNN cosine similarity.
///
/// # Arguments
/// * `centroids` - Flat array of L2-normalized centroids, shape (n_leaves, dim)
/// * `n_leaves` - Number of leaves
/// * `dim` - Embedding dimension
/// * `k` - Number of nearest neighbors per leaf
///
/// # Returns
/// Symmetric weighted adjacency list where weights are cosine similarities (clamped to [0, 1]).
pub fn build_knn_graph(
    centroids: &[f32],
    n_leaves: usize,
    dim: usize,
    k: usize,
) -> Vec<Vec<(u32, f32)>> {
    assert_eq!(centroids.len(), n_leaves * dim);
    let k = k.min(n_leaves - 1);

    // Compute all pairwise cosine similarities (centroids assumed L2-normalized)
    let mut adj: Vec<Vec<(u32, f32)>> = vec![Vec::new(); n_leaves];

    for i in 0..n_leaves {
        let ci = &centroids[i * dim..(i + 1) * dim];

        // Compute similarities to all other leaves
        let mut sims: Vec<(u32, f32)> = (0..n_leaves)
            .filter(|&j| j != i)
            .map(|j| {
                let cj = &centroids[j * dim..(j + 1) * dim];
                let sim: f32 = ci.iter().zip(cj.iter()).map(|(a, b)| a * b).sum();
                (j as u32, sim.max(0.0)) // clamp negative similarities
            })
            .collect();

        // Keep top-k by similarity
        sims.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        sims.truncate(k);

        for &(j, w) in &sims {
            adj[i].push((j, w));
        }
    }

    // Make symmetric: if i→j exists but j→i doesn't, add j→i
    for i in 0..n_leaves {
        let edges: Vec<(u32, f32)> = adj[i].clone();
        for (j, w) in edges {
            let j_idx = j as usize;
            if !adj[j_idx].iter().any(|&(n, _)| n == i as u32) {
                adj[j_idx].push((i as u32, w));
            }
        }
    }

    adj
}

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

    #[test]
    fn test_empty_graph() {
        let (labels, k) = louvain_communities(0, &[], 1.0);
        assert_eq!(labels.len(), 0);
        assert_eq!(k, 0);
    }

    #[test]
    fn test_single_node() {
        let adj = vec![vec![]];
        let (labels, k) = louvain_communities(1, &adj, 1.0);
        assert_eq!(labels, vec![0]);
        assert_eq!(k, 1);
    }

    #[test]
    fn test_disconnected_nodes() {
        // 3 disconnected nodes with 0 degree → m=0 → all placed in community 0
        let adj: Vec<Vec<(u32, f32)>> = vec![vec![], vec![], vec![]];
        let (labels, k) = louvain_communities(3, &adj, 1.0);
        assert_eq!(k, 1, "Zero-degree nodes all land in one community");
        assert!(labels.iter().all(|&l| l == 0));
    }

    #[test]
    fn test_two_cliques() {
        // Two triangles connected by a weak edge:
        // Clique A: 0-1-2 (weight 1.0 each)
        // Clique B: 3-4-5 (weight 1.0 each)
        // Bridge: 2-3 (weight 0.1)
        let adj = vec![
            // Node 0: connected to 1, 2
            vec![(1, 1.0), (2, 1.0)],
            // Node 1: connected to 0, 2
            vec![(0, 1.0), (2, 1.0)],
            // Node 2: connected to 0, 1, 3 (bridge)
            vec![(0, 1.0), (1, 1.0), (3, 0.1)],
            // Node 3: connected to 4, 5, 2 (bridge)
            vec![(4, 1.0), (5, 1.0), (2, 0.1)],
            // Node 4: connected to 3, 5
            vec![(3, 1.0), (5, 1.0)],
            // Node 5: connected to 3, 4
            vec![(3, 1.0), (4, 1.0)],
        ];

        let (labels, k) = louvain_communities(6, &adj, 1.0);
        assert_eq!(k, 2, "Should find 2 communities");

        // Nodes in same clique should have same label
        assert_eq!(labels[0], labels[1]);
        assert_eq!(labels[0], labels[2]);
        assert_eq!(labels[3], labels[4]);
        assert_eq!(labels[3], labels[5]);

        // Cliques should have different labels
        assert_ne!(labels[0], labels[3]);
    }

    #[test]
    fn test_complete_graph() {
        // K4 with uniform weights → should stay as one community
        let adj = vec![
            vec![(1, 1.0), (2, 1.0), (3, 1.0)],
            vec![(0, 1.0), (2, 1.0), (3, 1.0)],
            vec![(0, 1.0), (1, 1.0), (3, 1.0)],
            vec![(0, 1.0), (1, 1.0), (2, 1.0)],
        ];

        let (labels, k) = louvain_communities(4, &adj, 1.0);
        assert_eq!(k, 1, "Complete graph should be one community");
        assert!(labels.iter().all(|&l| l == 0));
    }

    #[test]
    fn test_resolution_parameter() {
        // Two cliques connected by moderate edge — higher resolution should split more
        let adj = vec![
            vec![(1, 1.0), (2, 1.0)],
            vec![(0, 1.0), (2, 1.0)],
            vec![(0, 1.0), (1, 1.0), (3, 0.5)],
            vec![(4, 1.0), (5, 1.0), (2, 0.5)],
            vec![(3, 1.0), (5, 1.0)],
            vec![(3, 1.0), (4, 1.0)],
        ];

        let (_, k_low) = louvain_communities(6, &adj, 0.5);
        let (_, k_high) = louvain_communities(6, &adj, 2.0);

        // Higher resolution should produce >= as many communities
        assert!(k_high >= k_low, "Higher resolution should produce more communities: k_high={k_high}, k_low={k_low}");
    }

    #[test]
    fn test_largest_community_is_zero() {
        // Asymmetric: clique of 4 + pair of 2
        let adj = vec![
            vec![(1, 1.0), (2, 1.0), (3, 1.0)],
            vec![(0, 1.0), (2, 1.0), (3, 1.0)],
            vec![(0, 1.0), (1, 1.0), (3, 1.0)],
            vec![(0, 1.0), (1, 1.0), (2, 1.0)],
            vec![(5, 1.0)],
            vec![(4, 1.0)],
        ];

        let (labels, k) = louvain_communities(6, &adj, 1.0);
        assert_eq!(k, 2);

        // The 4-node clique should be community 0 (largest)
        assert_eq!(labels[0], 0);
        assert_eq!(labels[1], 0);
        assert_eq!(labels[2], 0);
        assert_eq!(labels[3], 0);
        assert_eq!(labels[4], 1);
        assert_eq!(labels[5], 1);
    }

    #[test]
    fn test_build_knn_graph_simple() {
        // 3 unit vectors in 2D: (1,0), (0,1), (0.7,0.7)≈normalized
        let s = 1.0 / (2.0_f32).sqrt();
        let centroids = vec![
            1.0, 0.0, // node 0
            0.0, 1.0, // node 1
            s, s,     // node 2
        ];

        let adj = build_knn_graph(&centroids, 3, 2, 2);
        assert_eq!(adj.len(), 3);

        // Each node should have 2 neighbors (k=2, n-1=2)
        for a in &adj {
            assert_eq!(a.len(), 2);
        }

        // Node 2 should have highest similarity to both 0 and 1 (~0.707)
        let sim_2_0 = adj[2].iter().find(|&&(n, _)| n == 0).unwrap().1;
        let sim_2_1 = adj[2].iter().find(|&&(n, _)| n == 1).unwrap().1;
        assert!((sim_2_0 - s).abs() < 0.01);
        assert!((sim_2_1 - s).abs() < 0.01);
    }

    #[test]
    fn test_knn_graph_symmetry() {
        // Verify graph is symmetric
        let centroids = vec![
            1.0, 0.0, 0.0,
            0.0, 1.0, 0.0,
            0.0, 0.0, 1.0,
            0.577, 0.577, 0.577,
        ];

        let adj = build_knn_graph(&centroids, 4, 3, 2);

        for i in 0..4 {
            for &(j, _w) in &adj[i] {
                let reverse = adj[j as usize]
                    .iter()
                    .find(|&&(n, _)| n == i as u32);
                assert!(
                    reverse.is_some(),
                    "Edge {i}→{j} exists but {j}→{i} does not"
                );
            }
        }
    }
}
