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,
                    size: members.len(),
                    entity_ids: members.into_iter().collect(),
                    entity_names,
                    modularity_contribution: modularity_contrib,
                    average_strength: strength,
                    density,
                });
            }
        }

        // 5. Calculate total modularity
        let total_modularity: f64 = communities_vec
            .iter()
            .map(|c| c.modularity_contribution as f64)
            .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(),
            community_count: communities_vec.len(),
            communities: communities_vec,
            total_modularity: total_modularity.max(-1.0).min(1.0) as f32,
            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)
    }
}