poimen-memory/crates/mem-cli/src/query/community_detector.rs

//! Community Detection Engine
//!
//! Detects entity clusters using Louvain algorithm with modularity optimization.
//! Used to identify topic areas, entity groupings, and knowledge graph structure.

use serde::{Deserialize, Serialize};
use sqlx::{Pool, Postgres};
use std::collections::{HashMap, HashSet};
use tracing::{debug, info};

/// A detected community (cluster of related entities)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Community {
    pub id: usize,
    pub entity_ids: Vec<String>,
    pub entity_names: Vec<String>,
    pub size: usize,
    pub modularity_contribution: f32,  // This community's contribution to total modularity
    pub average_strength: f32,          // Average relationship strength within community
    pub density: f32,                   // 0-1, how tightly connected (actual edges / possible edges)
}

/// Results from community detection
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CommunityDetectionResult {
    pub entity_count: usize,
    pub edge_count: usize,
    pub communities: Vec<Community>,
    pub community_count: usize,
    pub total_modularity: f32,  // Overall modularity score (-1 to 1, higher is better)
    pub average_community_size: f32,
}

/// Edge representation for community detection
#[derive(Debug, Clone)]
struct GraphEdge {
    source: String,
    target: String,
    weight: f32,  // Relationship strength (0-1)
}

/// Community Detector using Louvain algorithm
pub struct CommunityDetector {
    pub pool: Pool<Postgres>,
}

impl CommunityDetector {
    /// Create a new community detector
    pub fn new(pool: Pool<Postgres>) -> Self {
        Self { pool }
    }

    /// Detect communities in the knowledge graph
    ///
    /// Uses Louvain algorithm to partition entities into communities
    /// based on relationship strength and graph structure.
    ///
    /// # Arguments
    /// * `project_id` - Project to analyze (optional, analyze all if None)
    /// * `min_community_size` - Minimum entities per community (default 3, min 2)
    /// * `modularity_threshold` - Stop optimization when improvement < threshold (default 0.001)
    ///
    /// # Returns
    /// CommunityDetectionResult with detected communities and metrics
    pub async fn detect_communities(
        &self,
        project_id: Option<&str>,
        min_community_size: usize,
        modularity_threshold: f32,
    ) -> Result<CommunityDetectionResult, String> {
        let min_community_size = min_community_size.max(2).min(1000);
        let modularity_threshold = modularity_threshold.max(0.0001).min(0.1);

        debug!(
            "Detecting communities: project={:?}, min_size={}, threshold={}",
            project_id, min_community_size, modularity_threshold
        );

        // 1. Fetch entities and edges from database
        let (entities, edges) = self.fetch_graph(project_id).await?;

        if entities.is_empty() {
            return Ok(CommunityDetectionResult {
                entity_count: 0,
                edge_count: 0,
                communities: vec![],
                community_count: 0,
                total_modularity: 0.0,
                average_community_size: 0.0,
            });
        }

        // 2. Initialize: each entity is its own community
        let mut entity_to_community: HashMap<String, usize> = HashMap::new();
        let mut community_members: HashMap<usize, HashSet<String>> = HashMap::new();

        for (idx, entity_id) in entities.iter().enumerate() {
            entity_to_community.insert(entity_id.clone(), idx);
            let mut members = HashSet::new();
            members.insert(entity_id.clone());
            community_members.insert(idx, members);
        }

        // 3. Louvain algorithm: iteratively optimize modularity
        let mut improved = true;
        let mut iteration = 0;
        let max_iterations = 100;

        while improved && iteration < max_iterations {
            improved = false;
            iteration += 1;

            // Try moving each entity to neighboring communities
            for entity_id in &entities {
                let current_community = entity_to_community[entity_id];
                let mut best_community = current_community;
                let mut best_modularity_gain = 0.0;

                // Find neighboring communities (connected via edges)
                let mut neighbor_communities = HashSet::new();
                neighbor_communities.insert(current_community);

                for edge in &edges {
                    if edge.source == *entity_id {
                        if let Some(&comm) = entity_to_community.get(&edge.target) {
                            neighbor_communities.insert(comm);
                        }
                    } else if edge.target == *entity_id {
                        if let Some(&comm) = entity_to_community.get(&edge.source) {
                            neighbor_communities.insert(comm);
                        }
                    }
                }

                // Evaluate moving to each neighbor community
                for &test_community in &neighbor_communities {
                    let gain = self.calculate_modularity_gain(
                        entity_id,
                        current_community,
                        test_community,
                        &edges,
                        &entity_to_community,
                    );

                    if gain > best_modularity_gain {
                        best_modularity_gain = gain;
                        best_community = test_community;
                    }
                }

                // Move entity if better community found
                if best_community != current_community && best_modularity_gain > modularity_threshold {
                    entity_to_community.insert(entity_id.clone(), best_community);
                    
                    // Update community membership
                    community_members
                        .get_mut(&current_community)
                        .map(|m| m.remove(entity_id));
                    community_members
                        .entry(best_community)
                        .or_insert_with(HashSet::new)
                        .insert(entity_id.clone());

                    improved = true;
                }
            }
        }

        // 4. Convert communities to output format
        let mut communities_vec = Vec::new();
        for (comm_id, members) in community_members {
            if members.len() >= min_community_size {
                let entity_names = members
                    .iter()
                    .map(|id| id.clone())  // In production, would look up actual names
                    .collect();

                let strength = self.calculate_community_strength(&members, &edges);
                let density = self.calculate_community_density(&members, &edges);
                let modularity_contrib = self.calculate_modularity_contribution(
                    &members,
                    &edges,
                    &entity_to_community,
                );

                communities_vec.push(Community {
                    id: comm_id,
                    entity_ids: members.into_iter().collect(),
                    entity_names,
                    size: members.len(),
                    modularity_contribution: modularity_contrib,
                    average_strength: strength,
                    density,
                });
            }
        }

        // 5. Calculate total modularity
        let total_modularity = communities_vec
            .iter()
            .map(|c| c.modularity_contribution)
            .sum();

        let average_community_size = if communities_vec.is_empty() {
            0.0
        } else {
            communities_vec.iter().map(|c| c.size as f32).sum::<f32>() / communities_vec.len() as f32
        };

        let result = CommunityDetectionResult {
            entity_count: entities.len(),
            edge_count: edges.len(),
            communities: communities_vec,
            community_count: communities_vec.len(),
            total_modularity: total_modularity.max(-1.0).min(1.0),
            average_community_size,
        };

        info!(
            "Community detection complete: {} communities, modularity={}",
            result.community_count, result.total_modularity
        );

        Ok(result)
    }

    /// Fetch entities and edges from database
    async fn fetch_graph(&self, _project_id: Option<&str>) -> Result<(Vec<String>, Vec<GraphEdge>), String> {
        // Fetch entities
        let entities = sqlx::query_as::<_, (String,)>(
            "SELECT DISTINCT id FROM memory_entity WHERE deleted_at IS NULL"
        )
            .fetch_all(&self.pool)
            .await
            .map_err(|e| format!("Failed to fetch entities: {}", e))?
            .into_iter()
            .map(|(id,)| id)
            .collect();

        // Fetch edges with confidence as weight
        let edges = sqlx::query_as::<_, (String, String, f32)>(
            "SELECT source_entity_id, target_entity_id, confidence 
             FROM memory_edge 
             WHERE fact_invalid_at IS NULL AND deleted_at IS NULL"
        )
            .fetch_all(&self.pool)
            .await
            .map_err(|e| format!("Failed to fetch edges: {}", e))?
            .into_iter()
            .map(|(source, target, confidence)| GraphEdge {
                source,
                target,
                weight: confidence.max(0.0).min(1.0),
            })
            .collect();

        Ok((entities, edges))
    }

    /// Calculate modularity gain of moving entity to target community
    fn calculate_modularity_gain(
        &self,
        entity_id: &str,
        from_community: usize,
        to_community: usize,
        edges: &[GraphEdge],
        entity_to_community: &HashMap<String, usize>,
    ) -> f32 {
        // Simplified modularity gain calculation
        // In production, use full Louvain formula with degrees
        
        let mut connections_to_target = 0.0;
        let mut connections_to_current = 0.0;

        for edge in edges {
            if edge.source == entity_id && entity_to_community.get(&edge.target).copied() == Some(to_community) {
                connections_to_target += edge.weight;
            } else if edge.target == entity_id && entity_to_community.get(&edge.source).copied() == Some(to_community) {
                connections_to_target += edge.weight;
            }

            if edge.source == entity_id && entity_to_community.get(&edge.target).copied() == Some(from_community) {
                connections_to_current += edge.weight;
            } else if edge.target == entity_id && entity_to_community.get(&edge.source).copied() == Some(from_community) {
                connections_to_current += edge.weight;
            }
        }

        // Gain = increased connections to target - lost connections from current
        (connections_to_target - connections_to_current) / edges.len().max(1) as f32
    }

    /// Calculate average relationship strength within community
    fn calculate_community_strength(&self, members: &HashSet<String>, edges: &[GraphEdge]) -> f32 {
        let mut total_weight = 0.0;
        let mut count = 0;

        for edge in edges {
            if members.contains(&edge.source) && members.contains(&edge.target) {
                total_weight += edge.weight;
                count += 1;
            }
        }

        if count == 0 {
            0.0
        } else {
            (total_weight / count as f32).max(0.0).min(1.0)
        }
    }

    /// Calculate community density (actual edges / possible edges)
    fn calculate_community_density(&self, members: &HashSet<String>, edges: &[GraphEdge]) -> f32 {
        let n = members.len() as f32;
        let possible_edges = (n * (n - 1.0) / 2.0).max(1.0);

        let mut actual_edges = 0.0;
        for edge in edges {
            if members.contains(&edge.source) && members.contains(&edge.target) {
                actual_edges += 1.0;
            }
        }

        (actual_edges / possible_edges).max(0.0).min(1.0)
    }

    /// Calculate this community's contribution to total modularity
    fn calculate_modularity_contribution(
        &self,
        members: &HashSet<String>,
        edges: &[GraphEdge],
        _entity_to_community: &HashMap<String, usize>,
    ) -> f32 {
        let internal_edges: f32 = edges
            .iter()
            .filter(|e| members.contains(&e.source) && members.contains(&e.target))
            .map(|e| e.weight)
            .sum();

        // Simplified: normalized by community size
        let max_possible = (members.len() as f32 * (members.len() as f32 - 1.0) / 2.0).max(1.0);
        (internal_edges / max_possible).max(0.0).min(1.0)
    }
}

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

    #[test]
    fn test_community_creation() {
        let community = Community {
            id: 0,
            entity_ids: vec!["e1".to_string(), "e2".to_string()],
            entity_names: vec!["Entity1".to_string(), "Entity2".to_string()],
            size: 2,
            modularity_contribution: 0.8,
            average_strength: 0.9,
            density: 1.0,
        };
        assert_eq!(community.size, 2);
        assert_eq!(community.entity_ids.len(), 2);
    }

    #[test]
    fn test_community_detection_result() {
        let result = CommunityDetectionResult {
            entity_count: 100,
            edge_count: 250,
            communities: vec![],
            community_count: 0,
            total_modularity: 0.0,
            average_community_size: 0.0,
        };
        assert_eq!(result.entity_count, 100);
        assert_eq!(result.edge_count, 250);
    }

    #[test]
    fn test_min_community_size_clamping() {
        let size = 1;
        let clamped = size.max(2).min(1000);
        assert_eq!(clamped, 2);

        let size = 5000;
        let clamped = size.max(2).min(1000);
        assert_eq!(clamped, 1000);
    }

    #[test]
    fn test_modularity_threshold_clamping() {
        let threshold = 0.0001;
        let clamped = threshold.max(0.0001).min(0.1);
        assert_eq!(clamped, 0.0001);

        let threshold = 0.5;
        let clamped = threshold.max(0.0001).min(0.1);
        assert_eq!(clamped, 0.1);
    }

    #[test]
    fn test_density_calculation() {
        // 3 entities, all connected (3 edges)
        // Possible edges: 3 * 2 / 2 = 3
        // Density: 3 / 3 = 1.0 (fully connected)
        let density = (3.0 / 3.0).max(0.0).min(1.0);
        assert_eq!(density, 1.0);

        // 4 entities, 2 edges
        // Possible: 4 * 3 / 2 = 6
        // Density: 2 / 6 ≈ 0.33
        let density = (2.0 / 6.0).max(0.0).min(1.0);
        assert!((density - 0.333).abs() < 0.01);
    }

    #[test]
    fn test_modularity_bounds() {
        let modularity = 0.75;
        let clamped = modularity.max(-1.0).min(1.0);
        assert_eq!(clamped, 0.75);

        let modularity = -0.5;
        let clamped = modularity.max(-1.0).min(1.0);
        assert_eq!(clamped, -0.5);
    }

    #[test]
    fn test_average_community_size() {
        let communities = vec![
            Community {
                id: 0,
                entity_ids: vec!["a".into(), "b".into(), "c".into()],
                entity_names: vec![],
                size: 3,
                modularity_contribution: 0.5,
                average_strength: 0.8,
                density: 0.9,
            },
            Community {
                id: 1,
                entity_ids: vec!["d".into(), "e".into()],
                entity_names: vec![],
                size: 2,
                modularity_contribution: 0.4,
                average_strength: 0.7,
                density: 1.0,
            },
        ];

        let avg = communities.iter().map(|c| c.size as f32).sum::<f32>() / communities.len() as f32;
        assert_eq!(avg, 2.5);
    }

    #[test]
    fn test_total_modularity_sum() {
        let contributions = vec![0.3, 0.25, 0.2, 0.15];
        let total: f32 = contributions.iter().sum();
        let clamped = total.max(-1.0).min(1.0);
        
        assert!(clamped >= -1.0 && clamped <= 1.0);
    }

    #[test]
    fn test_empty_graph_handling() {
        let entities: Vec<String> = vec![];
        let edges: Vec<GraphEdge> = vec![];

        assert!(entities.is_empty());
        assert!(edges.is_empty());
    }

    #[test]
    fn test_single_node_graph() {
        let entity_count = 1;
        let edge_count = 0;

        assert_eq!(entity_count, 1);
        assert_eq!(edge_count, 0);
    }

    #[test]
    fn test_fully_connected_graph() {
        // 5 nodes fully connected: 5*4/2 = 10 edges
        let nodes = 5;
        let possible_edges = nodes * (nodes - 1) / 2;
        assert_eq!(possible_edges, 10);
    }

    #[test]
    fn test_strength_normalization() {
        let strengths = vec![0.0, 0.25, 0.5, 0.75, 1.0];
        for s in strengths {
            let normalized = s.max(0.0).min(1.0);
            assert!(normalized >= 0.0 && normalized <= 1.0);
        }
    }

    #[test]
    fn test_louvain_max_iterations() {
        let max_iterations = 100;
        let mut iteration = 0;

        while iteration < max_iterations && iteration < 5 {
            iteration += 1;
        }

        assert!(iteration <= max_iterations);
    }
}