arrow-graph 0.2.0

use arrow::record_batch::RecordBatch;
use arrow::array::{StringArray, Float64Array, UInt32Array};
use arrow::datatypes::{DataType, Field, Schema};
use std::sync::Arc;
use std::collections::HashMap;
use crate::algorithms::{GraphAlgorithm, AlgorithmParams};
use crate::graph::ArrowGraph;
use crate::error::{GraphError, Result};

pub struct PageRank;

impl PageRank {
    /// PageRank algorithm with power iteration and early termination
    fn compute_pagerank(
        &self,
        graph: &ArrowGraph,
        damping_factor: f64,
        max_iterations: usize,
        tolerance: f64,
    ) -> Result<HashMap<String, f64>> {
        let node_count = graph.node_count();
        if node_count == 0 {
            return Ok(HashMap::new());
        }
        
        // Initialize PageRank scores
        let initial_score = 1.0 / node_count as f64;
        let mut current_scores: HashMap<String, f64> = HashMap::new();
        let mut next_scores: HashMap<String, f64> = HashMap::new();
        
        // Initialize all nodes with equal probability
        for node_id in graph.node_ids() {
            current_scores.insert(node_id.clone(), initial_score);
            next_scores.insert(node_id.clone(), 0.0);
        }
        
        // Calculate out-degrees for each node
        let mut out_degrees: HashMap<String, usize> = HashMap::new();
        for node_id in graph.node_ids() {
            let degree = graph.neighbors(node_id).map(|n| n.len()).unwrap_or(0);
            out_degrees.insert(node_id.clone(), degree);
        }
        
        // Power iteration
        for iteration in 0..max_iterations {
            // Reset next scores
            for score in next_scores.values_mut() {
                *score = (1.0 - damping_factor) / node_count as f64;
            }
            
            // Distribute PageRank scores
            for node_id in graph.node_ids() {
                let current_score = current_scores.get(node_id).unwrap_or(&0.0);
                let out_degree = out_degrees.get(node_id).unwrap_or(&0);
                
                if *out_degree > 0 {
                    let contribution = current_score * damping_factor / *out_degree as f64;
                    
                    if let Some(neighbors) = graph.neighbors(node_id) {
                        for neighbor in neighbors {
                            if let Some(neighbor_score) = next_scores.get_mut(neighbor) {
                                *neighbor_score += contribution;
                            }
                        }
                    }
                } else {
                    // Handle dangling nodes - distribute equally to all nodes
                    let dangling_contribution = current_score * damping_factor / node_count as f64;
                    for score in next_scores.values_mut() {
                        *score += dangling_contribution;
                    }
                }
            }
            
            // Check for convergence
            let mut diff = 0.0;
            for node_id in graph.node_ids() {
                let old_score = current_scores.get(node_id).unwrap_or(&0.0);
                let new_score = next_scores.get(node_id).unwrap_or(&0.0);
                diff += (new_score - old_score).abs();
            }
            
            // Early termination if converged
            if diff < tolerance {
                log::debug!("PageRank converged after {} iterations", iteration + 1);
                break;
            }
            
            // Swap scores for next iteration
            std::mem::swap(&mut current_scores, &mut next_scores);
        }
        
        Ok(current_scores)
    }
}

impl GraphAlgorithm for PageRank {
    fn execute(&self, graph: &ArrowGraph, params: &AlgorithmParams) -> Result<RecordBatch> {
        let damping_factor: f64 = params.get("damping_factor").unwrap_or(0.85);
        let max_iterations: usize = params.get("max_iterations").unwrap_or(100);
        let tolerance: f64 = params.get("tolerance").unwrap_or(1e-6);
        
        // Validate parameters
        if !(0.0..=1.0).contains(&damping_factor) {
            return Err(GraphError::invalid_parameter(
                "damping_factor must be between 0.0 and 1.0"
            ));
        }
        
        if max_iterations == 0 {
            return Err(GraphError::invalid_parameter(
                "max_iterations must be greater than 0"
            ));
        }
        
        if tolerance <= 0.0 {
            return Err(GraphError::invalid_parameter(
                "tolerance must be greater than 0.0"
            ));
        }
        
        let scores = self.compute_pagerank(graph, damping_factor, max_iterations, tolerance)?;
        
        // Convert to Arrow RecordBatch
        let schema = Arc::new(Schema::new(vec![
            Field::new("node_id", DataType::Utf8, false),
            Field::new("pagerank_score", DataType::Float64, false),
        ]));
        
        let mut node_ids = Vec::new();
        let mut pagerank_scores = Vec::new();
        
        // Sort by PageRank score (descending) for consistent output
        let mut sorted_scores: Vec<(&String, &f64)> = scores.iter().collect();
        sorted_scores.sort_by(|a, b| b.1.partial_cmp(a.1).unwrap_or(std::cmp::Ordering::Equal));
        
        for (node_id, score) in sorted_scores {
            node_ids.push(node_id.clone());
            pagerank_scores.push(*score);
        }
        
        RecordBatch::try_new(
            schema,
            vec![
                Arc::new(StringArray::from(node_ids)),
                Arc::new(Float64Array::from(pagerank_scores)),
            ],
        ).map_err(GraphError::from)
    }
    
    fn name(&self) -> &'static str {
        "pagerank"
    }
    
    fn description(&self) -> &'static str {
        "Calculate PageRank scores using power iteration with early termination"
    }
}

pub struct BetweennessCentrality;

impl BetweennessCentrality {
    /// Calculate betweenness centrality using Brandes' algorithm
    fn compute_betweenness_centrality(&self, graph: &ArrowGraph) -> Result<HashMap<String, f64>> {
        let mut centrality: HashMap<String, f64> = HashMap::new();
        
        // Initialize centrality scores
        for node_id in graph.node_ids() {
            centrality.insert(node_id.clone(), 0.0);
        }
        
        // For each node as source
        for source in graph.node_ids() {
            let mut stack = Vec::new();
            let mut paths: HashMap<String, Vec<String>> = HashMap::new();
            let mut num_paths: HashMap<String, f64> = HashMap::new();
            let mut distances: HashMap<String, i32> = HashMap::new();
            let mut delta: HashMap<String, f64> = HashMap::new();
            
            // Initialize
            for node_id in graph.node_ids() {
                paths.insert(node_id.clone(), Vec::new());
                num_paths.insert(node_id.clone(), 0.0);
                distances.insert(node_id.clone(), -1);
                delta.insert(node_id.clone(), 0.0);
            }
            
            num_paths.insert(source.clone(), 1.0);
            distances.insert(source.clone(), 0);
            
            // BFS
            let mut queue = std::collections::VecDeque::new();
            queue.push_back(source.clone());
            
            while let Some(current) = queue.pop_front() {
                stack.push(current.clone());
                
                if let Some(neighbors) = graph.neighbors(&current) {
                    for neighbor in neighbors {
                        let current_dist = *distances.get(&current).unwrap_or(&-1);
                        let neighbor_dist = *distances.get(neighbor).unwrap_or(&-1);
                        
                        // First time we reach this neighbor
                        if neighbor_dist < 0 {
                            queue.push_back(neighbor.clone());
                            distances.insert(neighbor.clone(), current_dist + 1);
                        }
                        
                        // Shortest path to neighbor via current
                        if neighbor_dist == current_dist + 1 {
                            let current_paths = *num_paths.get(&current).unwrap_or(&0.0);
                            let neighbor_paths = num_paths.get_mut(neighbor).unwrap();
                            *neighbor_paths += current_paths;
                            
                            paths.get_mut(neighbor).unwrap().push(current.clone());
                        }
                    }
                }
            }
            
            // Accumulation
            while let Some(w) = stack.pop() {
                if let Some(predecessors) = paths.get(&w) {
                    for predecessor in predecessors {
                        let w_delta = *delta.get(&w).unwrap_or(&0.0);
                        let w_paths = *num_paths.get(&w).unwrap_or(&0.0);
                        let pred_paths = *num_paths.get(predecessor).unwrap_or(&0.0);
                        
                        if pred_paths > 0.0 {
                            let contribution = (pred_paths / w_paths) * (1.0 + w_delta);
                            *delta.get_mut(predecessor).unwrap() += contribution;
                        }
                    }
                }
                
                if w != *source {
                    let w_delta = *delta.get(&w).unwrap_or(&0.0);
                    *centrality.get_mut(&w).unwrap() += w_delta;
                }
            }
        }
        
        // Normalize for undirected graphs
        let node_count = graph.node_count() as f64;
        if node_count > 2.0 {
            let normalization = 2.0 / ((node_count - 1.0) * (node_count - 2.0));
            for score in centrality.values_mut() {
                *score *= normalization;
            }
        }
        
        Ok(centrality)
    }
}

impl GraphAlgorithm for BetweennessCentrality {
    fn execute(&self, _graph: &ArrowGraph, _params: &AlgorithmParams) -> Result<RecordBatch> {
        todo!("Implement betweenness centrality - complex algorithm, implementing in future version")
    }
    
    fn name(&self) -> &'static str {
        "betweenness_centrality"
    }
    
    fn description(&self) -> &'static str {
        "Calculate betweenness centrality using Brandes' algorithm"
    }
}