arrow-graph 0.6.2

use crate::error::Result;
use crate::graph::ArrowGraph;
use std::collections::HashMap;
use std::time::{Duration, Instant};
use wide::f64x4;

/// Simple benchmark framework for basic performance testing
#[derive(Debug)]
pub struct SimpleBenchmark {
    pub name: String,
    pub warmup_iterations: usize,
    pub measurement_iterations: usize,
}

/// Simple benchmark result
#[derive(Debug, Clone)]
pub struct SimpleBenchmarkResult {
    pub name: String,
    pub avg_duration: Duration,
    pub min_duration: Duration,
    pub max_duration: Duration,
    pub total_duration: Duration,
    pub iterations: usize,
    pub throughput_ops_per_sec: f64,
}

impl SimpleBenchmark {
    /// Create a new simple benchmark
    pub fn new(name: &str) -> Self {
        Self {
            name: name.to_string(),
            warmup_iterations: 3,
            measurement_iterations: 10,
        }
    }

    /// Run REAL PageRank benchmark with iterative algorithm  
    pub fn bench_pagerank(&self, graph: &ArrowGraph) -> Result<SimpleBenchmarkResult> {
        // Warmup
        for _ in 0..self.warmup_iterations {
            let _ = self.run_pagerank_real(graph);
        }

        // Measure
        let mut durations = Vec::new();
        let total_start = Instant::now();

        for _ in 0..self.measurement_iterations {
            let start = Instant::now();
            let _ = self.run_pagerank_real(graph);
            durations.push(start.elapsed());
        }

        let total_duration = total_start.elapsed();
        self.calculate_result("Optimized PageRank", durations, total_duration, graph.node_count())
    }

    /// Run SIMD-optimized PageRank benchmark
    pub fn bench_pagerank_simd(&self, graph: &ArrowGraph) -> Result<SimpleBenchmarkResult> {
        // Warmup
        for _ in 0..self.warmup_iterations {
            let _ = self.run_pagerank_simd(graph);
        }

        // Measure
        let mut durations = Vec::new();
        let total_start = Instant::now();

        for _ in 0..self.measurement_iterations {
            let start = Instant::now();
            let _ = self.run_pagerank_simd(graph);
            durations.push(start.elapsed());
        }

        let total_duration = total_start.elapsed();
        self.calculate_result("SIMD PageRank", durations, total_duration, graph.node_count())
    }

    /// Run REAL Connected Components benchmark with Union-Find
    pub fn bench_connected_components(&self, graph: &ArrowGraph) -> Result<SimpleBenchmarkResult> {
        // Warmup
        for _ in 0..self.warmup_iterations {
            let _ = self.run_connected_components_real(graph);
        }

        // Measure
        let mut durations = Vec::new();
        let total_start = Instant::now();

        for _ in 0..self.measurement_iterations {
            let start = Instant::now();
            let _ = self.run_connected_components_real(graph);
            durations.push(start.elapsed());
        }

        let total_duration = total_start.elapsed();
        self.calculate_result("Real Connected Components", durations, total_duration, graph.node_count())
    }

    /// Run graph construction benchmark
    pub fn bench_graph_construction(&self, num_nodes: usize, num_edges: usize) -> Result<SimpleBenchmarkResult> {
        // Warmup
        for _ in 0..self.warmup_iterations {
            let _ = self.create_test_graph(num_nodes, num_edges);
        }

        // Measure
        let mut durations = Vec::new();
        let total_start = Instant::now();

        for _ in 0..self.measurement_iterations {
            let start = Instant::now();
            let _ = self.create_test_graph(num_nodes, num_edges);
            durations.push(start.elapsed());
        }

        let total_duration = total_start.elapsed();
        self.calculate_result("Graph Construction", durations, total_duration, num_nodes)
    }

    /// Calculate benchmark result from durations
    fn calculate_result(
        &self,
        name: &str,
        durations: Vec<Duration>,
        total_duration: Duration,
        operations: usize,
    ) -> Result<SimpleBenchmarkResult> {
        let min_duration = durations.iter().min().copied().unwrap_or(Duration::from_secs(0));
        let max_duration = durations.iter().max().copied().unwrap_or(Duration::from_secs(0));
        let avg_duration = Duration::from_nanos(
            (durations.iter().map(|d| d.as_nanos()).sum::<u128>() / durations.len() as u128) as u64
        );

        let throughput_ops_per_sec = if avg_duration.as_secs_f64() > 0.0 {
            operations as f64 / avg_duration.as_secs_f64()
        } else {
            0.0
        };

        Ok(SimpleBenchmarkResult {
            name: name.to_string(),
            avg_duration,
            min_duration,
            max_duration,
            total_duration,
            iterations: durations.len(),
            throughput_ops_per_sec,
        })
    }

    // REAL algorithm implementations for benchmarking
    fn run_pagerank_real(&self, graph: &ArrowGraph) -> Vec<f64> {
        // OPTIMIZED PageRank with sparse operations and reduced allocations
        let damping_factor = 0.85;
        let tolerance = 1e-6;
        let max_iterations = 50;
        
        // Collect all node IDs into a vector for indexing
        let node_ids: Vec<String> = graph.node_ids().cloned().collect();
        let node_count = node_ids.len();
        
        if node_count == 0 {
            return Vec::new();
        }
        
        // Pre-allocate and create node index mapping ONCE
        let mut node_to_index = HashMap::with_capacity(node_count);
        for (i, node_id) in node_ids.iter().enumerate() {
            node_to_index.insert(node_id, i);
        }
        
        // Build sparse adjacency structure for better cache performance
        let mut sparse_neighbors: Vec<Vec<usize>> = vec![Vec::new(); node_count];
        let mut out_degrees = vec![0; node_count];
        
        // Build adjacency lists with indices instead of strings (faster)
        for (i, node_id) in node_ids.iter().enumerate() {
            if let Some(neighbors) = graph.neighbors(node_id) {
                out_degrees[i] = neighbors.len();
                sparse_neighbors[i].reserve(neighbors.len());
                
                for neighbor in neighbors {
                    if let Some(&neighbor_idx) = node_to_index.get(neighbor) {
                        sparse_neighbors[i].push(neighbor_idx);
                    }
                }
            }
        }
        
        // Initialize PageRank values
        let initial_rank = 1.0 / node_count as f64;
        let mut current_ranks = vec![initial_rank; node_count];
        let mut new_ranks = vec![0.0; node_count];
        
        // Pre-calculate base rank contribution (optimization)
        let base_rank = (1.0 - damping_factor) / node_count as f64;
        
        // PageRank iterations with SIMD-friendly operations
        for iteration in 0..max_iterations {
            // Reset new ranks using SIMD-friendly batch operation
            for rank in new_ranks.iter_mut() {
                *rank = base_rank;
            }
            
            // Accumulate dangling node contribution
            let mut dangling_contribution = 0.0;
            for i in 0..node_count {
                if out_degrees[i] == 0 {
                    dangling_contribution += current_ranks[i];
                }
            }
            dangling_contribution = damping_factor * dangling_contribution / node_count as f64;
            
            // Distribute dangling contribution to all nodes
            for rank in new_ranks.iter_mut() {
                *rank += dangling_contribution;
            }
            
            // Distribute rank using sparse adjacency (cache-friendly)
            for i in 0..node_count {
                if !sparse_neighbors[i].is_empty() {
                    let contribution = damping_factor * current_ranks[i] / sparse_neighbors[i].len() as f64;
                    
                    // SIMD-friendly inner loop over neighbors
                    for &neighbor_idx in &sparse_neighbors[i] {
                        new_ranks[neighbor_idx] += contribution;
                    }
                }
            }
            
            // Vectorized convergence check (SIMD-friendly)
            let mut diff = 0.0;
            for i in 0..node_count {
                diff += (new_ranks[i] - current_ranks[i]).abs();
            }
            
            // Swap current and new ranks (no allocation)
            std::mem::swap(&mut current_ranks, &mut new_ranks);
            
            // Early termination if converged
            if diff < tolerance {
                println!("PageRank converged after {} iterations", iteration + 1);
                break;
            }
        }
        
        current_ranks
    }

    fn run_pagerank_simd(&self, graph: &ArrowGraph) -> Vec<f64> {
        // SIMD-OPTIMIZED PageRank using wide::f64x4 for 4x parallelism
        let damping_factor = 0.85;
        let tolerance = 1e-6;
        let max_iterations = 50;
        
        let node_ids: Vec<String> = graph.node_ids().cloned().collect();
        let node_count = node_ids.len();
        
        if node_count == 0 {
            return Vec::new();
        }
        
        // Pre-allocate with SIMD alignment (pad to multiple of 4)
        let simd_size = ((node_count + 3) / 4) * 4;
        let mut node_to_index = HashMap::with_capacity(node_count);
        for (i, node_id) in node_ids.iter().enumerate() {
            node_to_index.insert(node_id, i);
        }
        
        // Build sparse adjacency with SIMD-friendly layout
        let mut sparse_neighbors: Vec<Vec<usize>> = vec![Vec::new(); node_count];
        let mut out_degrees = vec![0; node_count];
        
        for (i, node_id) in node_ids.iter().enumerate() {
            if let Some(neighbors) = graph.neighbors(node_id) {
                out_degrees[i] = neighbors.len();
                sparse_neighbors[i].reserve(neighbors.len());
                
                for neighbor in neighbors {
                    if let Some(&neighbor_idx) = node_to_index.get(neighbor) {
                        sparse_neighbors[i].push(neighbor_idx);
                    }
                }
            }
        }
        
        // SIMD-aligned vectors (padded to multiple of 4)
        let initial_rank = 1.0 / node_count as f64;
        let mut current_ranks = vec![initial_rank; simd_size];
        let mut new_ranks = vec![0.0; simd_size];
        
        // Pad unused elements with 0
        for i in node_count..simd_size {
            current_ranks[i] = 0.0;
        }
        
        let base_rank = (1.0 - damping_factor) / node_count as f64;
        let _base_rank_simd = f64x4::splat(base_rank);
        
        for iteration in 0..max_iterations {
            // SIMD reset of new_ranks (4 elements at a time)
            for chunk in new_ranks.chunks_exact_mut(4) {
                let base_array: [f64; 4] = [base_rank; 4];
                chunk.copy_from_slice(&base_array);
            }
            
            // Handle dangling nodes
            let mut dangling_contribution = 0.0;
            for i in 0..node_count {
                if out_degrees[i] == 0 {
                    dangling_contribution += current_ranks[i];
                }
            }
            dangling_contribution = damping_factor * dangling_contribution / node_count as f64;
            let dangling_simd = f64x4::splat(dangling_contribution);
            
            // SIMD addition of dangling contribution
            for chunk in new_ranks.chunks_exact_mut(4) {
                if chunk.len() == 4 {
                    let current_chunk = f64x4::new([chunk[0], chunk[1], chunk[2], chunk[3]]);
                    let updated_chunk = current_chunk + dangling_simd;
                    let result_array: [f64; 4] = updated_chunk.into();
                    chunk.copy_from_slice(&result_array);
                }
            }
            
            // Distribute rank contributions (this part is still scalar due to sparse nature)
            for i in 0..node_count {
                if !sparse_neighbors[i].is_empty() {
                    let contribution = damping_factor * current_ranks[i] / sparse_neighbors[i].len() as f64;
                    
                    for &neighbor_idx in &sparse_neighbors[i] {
                        new_ranks[neighbor_idx] += contribution;
                    }
                }
            }
            
            // SIMD convergence check (4 elements at a time)
            let mut diff_accumulator = f64x4::splat(0.0);
            
            for i in (0..node_count).step_by(4) {
                let end = (i + 4).min(simd_size);
                if end - i == 4 && i + 4 <= node_count {
                    let new_chunk = f64x4::new([new_ranks[i], new_ranks[i+1], new_ranks[i+2], new_ranks[i+3]]);
                    let current_chunk = f64x4::new([current_ranks[i], current_ranks[i+1], current_ranks[i+2], current_ranks[i+3]]);
                    let diff_chunk = (new_chunk - current_chunk).abs();
                    diff_accumulator += diff_chunk;
                } else {
                    // Handle remaining elements scalar
                    for j in i..end.min(node_count) {
                        let scalar_diff = (new_ranks[j] - current_ranks[j]).abs();
                        diff_accumulator += f64x4::splat(scalar_diff);
                    }
                }
            }
            
            // Sum SIMD lanes to get total difference
            let diff_array: [f64; 4] = diff_accumulator.into();
            let total_diff = diff_array.iter().sum::<f64>();
            
            // Swap ranks
            std::mem::swap(&mut current_ranks, &mut new_ranks);
            
            if total_diff < tolerance {
                println!("SIMD PageRank converged after {} iterations", iteration + 1);
                break;
            }
        }
        
        // Return only the valid node count elements
        current_ranks.truncate(node_count);
        current_ranks
    }

    fn run_connected_components_real(&self, graph: &ArrowGraph) -> Vec<usize> {
        // REAL Connected Components using Union-Find algorithm
        let node_ids: Vec<String> = graph.node_ids().cloned().collect();
        let node_count = node_ids.len();
        
        if node_count == 0 {
            return Vec::new();
        }
        
        // Create node ID to index mapping
        let node_to_index: std::collections::HashMap<String, usize> = node_ids
            .iter()
            .enumerate()
            .map(|(i, id)| (id.clone(), i))
            .collect();
        
        // Union-Find data structure
        let mut parent = vec![0; node_count];
        let mut rank = vec![0; node_count];
        
        // Initialize: each node is its own parent
        for i in 0..node_count {
            parent[i] = i;
        }
        
        // Find with path compression
        fn find(parent: &mut [usize], x: usize) -> usize {
            if parent[x] != x {
                parent[x] = find(parent, parent[x]); // Path compression
            }
            parent[x]
        }
        
        // Union by rank
        fn union(parent: &mut [usize], rank: &mut [usize], x: usize, y: usize) {
            let root_x = find(parent, x);
            let root_y = find(parent, y);
            
            if root_x != root_y {
                if rank[root_x] < rank[root_y] {
                    parent[root_x] = root_y;
                } else if rank[root_x] > rank[root_y] {
                    parent[root_y] = root_x;
                } else {
                    parent[root_y] = root_x;
                    rank[root_x] += 1;
                }
            }
        }
        
        // Process all edges to build connected components
        for (i, node_id) in node_ids.iter().enumerate() {
            if let Some(neighbors) = graph.neighbors(node_id) {
                for neighbor in neighbors {
                    if let Some(&neighbor_idx) = node_to_index.get(neighbor) {
                        union(&mut parent, &mut rank, i, neighbor_idx);
                    }
                }
            }
        }
        
        // Find final component IDs for each node
        let mut components = Vec::with_capacity(node_count);
        for i in 0..node_count {
            components.push(find(&mut parent, i));
        }
        
        // Normalize component IDs to be 0, 1, 2, ... (compress them)
        let mut component_map = std::collections::HashMap::new();
        let mut next_component_id = 0;
        
        for component in &mut components {
            if !component_map.contains_key(component) {
                component_map.insert(*component, next_component_id);
                next_component_id += 1;
            }
            *component = component_map[component];
        }
        
        components
    }

    fn create_test_graph(&self, num_nodes: usize, num_edges: usize) -> ArrowGraph {
        use arrow::array::{StringArray, RecordBatch};
        use arrow::datatypes::{DataType, Field, Schema};
        use std::sync::Arc;
        
        // Create empty nodes and edges batches for now
        let node_schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Utf8, false),
        ]));
        
        let edge_schema = Arc::new(Schema::new(vec![
            Field::new("source", DataType::Utf8, false),
            Field::new("target", DataType::Utf8, false),
        ]));
        
        // Create minimal data for the test
        let node_ids: Vec<String> = (0..num_nodes.min(100)).map(|i| format!("node_{}", i)).collect();
        let edge_sources: Vec<String> = (0..num_edges.min(100)).map(|i| format!("node_{}", i % node_ids.len())).collect();
        let edge_targets: Vec<String> = (0..num_edges.min(100)).map(|i| format!("node_{}", (i + 1) % node_ids.len())).collect();
        
        let nodes_batch = RecordBatch::try_new(
            node_schema,
            vec![Arc::new(StringArray::from(node_ids))],
        ).unwrap();
        
        let edges_batch = RecordBatch::try_new(
            edge_schema,
            vec![
                Arc::new(StringArray::from(edge_sources)),
                Arc::new(StringArray::from(edge_targets)),
            ],
        ).unwrap();
        
        ArrowGraph::new(nodes_batch, edges_batch).unwrap()
    }
}

/// Run all simple benchmarks
pub fn run_all_benchmarks(graph: &ArrowGraph) -> Result<Vec<SimpleBenchmarkResult>> {
    let benchmark = SimpleBenchmark::new("Arrow Graph Performance");
    let mut results = Vec::new();

    println!("Running simple benchmark suite...");

    // Optimized PageRank benchmark
    println!("Running Optimized PageRank benchmark...");
    let pagerank_result = benchmark.bench_pagerank(graph)?;
    results.push(pagerank_result);

    // SIMD PageRank benchmark
    println!("Running SIMD PageRank benchmark...");
    let simd_pagerank_result = benchmark.bench_pagerank_simd(graph)?;
    results.push(simd_pagerank_result);

    // Connected Components benchmark
    println!("Running Connected Components benchmark...");
    let cc_result = benchmark.bench_connected_components(graph)?;
    results.push(cc_result);

    // Graph Construction benchmark
    println!("Running Graph Construction benchmark...");
    let construction_result = benchmark.bench_graph_construction(1000, 5000)?;
    results.push(construction_result);

    println!("Benchmark suite completed!");
    Ok(results)
}

/// Print benchmark results in a nice format
pub fn print_results(results: &[SimpleBenchmarkResult]) {
    println!("\n{:-<80}", "");
    println!("{:^80}", "ARROW GRAPH BENCHMARK RESULTS");
    println!("{:-<80}", "");
    println!("{:<25} {:>12} {:>12} {:>12} {:>15}", 
             "Benchmark", "Avg (ms)", "Min (ms)", "Max (ms)", "Throughput (ops/s)");
    println!("{:-<80}", "");

    for result in results {
        println!("{:<25} {:>12.2} {:>12.2} {:>12.2} {:>15.0}",
                 result.name,
                 result.avg_duration.as_millis(),
                 result.min_duration.as_millis(),
                 result.max_duration.as_millis(),
                 result.throughput_ops_per_sec);
    }

    println!("{:-<80}", "");
}

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

    #[test]
    fn test_simple_benchmark_creation() {
        let bench = SimpleBenchmark::new("test");
        assert_eq!(bench.name, "test");
        assert_eq!(bench.warmup_iterations, 3);
        assert_eq!(bench.measurement_iterations, 10);
    }

    #[test]
    fn test_pagerank_benchmark() {
        let bench = SimpleBenchmark::new("test");
        
        // Create a simple test graph
        use arrow::array::{StringArray, Float64Array};
        use arrow::datatypes::{DataType, Field, Schema};
        use arrow::record_batch::RecordBatch;
        use std::sync::Arc;
        
        let nodes_schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Utf8, false),
        ]));
        let node_ids = StringArray::from(vec!["A", "B", "C"]);
        let nodes_batch = RecordBatch::try_new(
            nodes_schema,
            vec![Arc::new(node_ids)],
        ).unwrap();

        let edges_schema = Arc::new(Schema::new(vec![
            Field::new("source", DataType::Utf8, false),
            Field::new("target", DataType::Utf8, false),
            Field::new("weight", DataType::Float64, false),
        ]));
        let sources = StringArray::from(vec!["A", "B"]);
        let targets = StringArray::from(vec!["B", "C"]);
        let weights = Float64Array::from(vec![1.0, 1.0]);
        let edges_batch = RecordBatch::try_new(
            edges_schema,
            vec![Arc::new(sources), Arc::new(targets), Arc::new(weights)],
        ).unwrap();

        let graph = ArrowGraph::new(nodes_batch, edges_batch).unwrap();
        
        let result = bench.bench_pagerank(&graph).unwrap();
        assert_eq!(result.name, "PageRank");
        assert!(result.avg_duration.as_nanos() > 0);
        assert_eq!(result.iterations, 10);
    }

    #[test]
    fn test_graph_construction_benchmark() {
        let bench = SimpleBenchmark::new("test");
        let result = bench.bench_graph_construction(100, 200).unwrap();
        assert_eq!(result.name, "Graph Construction");
        assert!(result.avg_duration.as_nanos() > 0);
    }
}