//! Tests for ring-based AllReduce implementation
//!
//! These tests verify the ring-based AllReduce algorithm implementation
//! for efficient gradient synchronization across multiple GPUs.

use rtx_distributed::{Backend, MultiGpuTrainer, ProcessGroup, ReduceOp, Result, WorldInfo};
use rtx_tensor::{Device, Tensor};
use std::time::Duration;

/// Test ring AllReduce basic functionality
#[tokio::test]
async fn test_ring_allreduce_basic() -> Result<()> {
    let world_size = 4;
    let rank = 0;

    let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;

    // Enable communication overlap to use ring AllReduce
    trainer.scaling_optimizer.communication_overlap_enabled = true;

    // Create test gradients with known values
    let shape_dims = vec![1000];
    let gradient1 = Tensor::full(
        &shape_dims,
        2.0,
        &Device::cuda(0).unwrap_or(Device::default()),
    )?;
    let gradient2 = Tensor::full(
        &shape_dims,
        4.0,
        &Device::cuda(0).unwrap_or(Device::default()),
    )?;
    let mut gradients = vec![gradient1, gradient2];

    // Perform ring AllReduce synchronization
    trainer.synchronize_gradients(&mut gradients).await?;

    // Verify the gradients were processed correctly
    let grad1_data = gradients[0].data()?;
    let grad2_data = gradients[1].data()?;

    println!(
        "Grad1 first value: {}, last value: {}",
        grad1_data[0],
        grad1_data[grad1_data.len() - 1]
    );
    println!(
        "Grad2 first value: {}, last value: {}",
        grad2_data[0],
        grad2_data[grad2_data.len() - 1]
    );

    // Check that all values within each tensor are consistent
    for (i, &value) in grad1_data.iter().enumerate() {
        if (value - grad1_data[0]).abs() >= 1e-6 {
            println!(
                "Inconsistent value at index {}: {} vs {}",
                i, value, grad1_data[0]
            );
        }
        assert!(
            (value - grad1_data[0]).abs() < 1e-6,
            "Ring AllReduce should produce consistent values"
        );
    }

    for (i, &value) in grad2_data.iter().enumerate() {
        if (value - grad2_data[0]).abs() >= 1e-6 {
            println!(
                "Inconsistent value at index {}: {} vs {}",
                i, value, grad2_data[0]
            );
        }
        assert!(
            (value - grad2_data[0]).abs() < 1e-6,
            "Ring AllReduce should produce consistent values"
        );
    }

    Ok(())
}

/// Test ring AllReduce with different buffer sizes
#[tokio::test]
async fn test_ring_allreduce_buffer_sizes() -> Result<()> {
    let world_size = 4;
    let rank = 0;

    let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
    trainer.scaling_optimizer.communication_overlap_enabled = true;

    // Test different buffer sizes
    let buffer_sizes = vec![512, 1024, 2048, 4096];

    for buffer_size in buffer_sizes {
        trainer.scaling_optimizer.ring_buffer_size = buffer_size;

        // Create gradients that will be chunked based on buffer size
        let shape_dims = vec![buffer_size * 2]; // Larger than buffer to test chunking
        let gradient = Tensor::full(
            &shape_dims,
            1.0,
            &Device::cuda(0).unwrap_or(Device::default()),
        )?;
        let mut gradients = vec![gradient];

        let start_time = std::time::Instant::now();
        trainer.synchronize_gradients(&mut gradients).await?;
        let duration = start_time.elapsed();

        println!(
            "Ring AllReduce with buffer size {}: {:.2}ms",
            buffer_size,
            duration.as_millis()
        );

        // Verify correctness
        let data = gradients[0].data()?;
        for &value in &data {
            assert!(
                (value - data[0]).abs() < 1e-6,
                "Values should be consistent after ring AllReduce"
            );
        }
    }

    Ok(())
}

/// Test ring AllReduce communication overlap
#[tokio::test]
async fn test_ring_allreduce_overlap() -> Result<()> {
    let world_size = 8;
    let rank = 0;

    let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;

    // Create large gradients to test overlap effectiveness
    let shape_dims = vec![100_000];
    let gradient = Tensor::ones(&shape_dims, &Device::cuda(0).unwrap_or(Device::default()))?;
    let mut gradients = vec![gradient.clone(), gradient.clone(), gradient];

    // Test without overlap
    trainer.scaling_optimizer.communication_overlap_enabled = false;
    let start_no_overlap = std::time::Instant::now();
    trainer.synchronize_gradients(&mut gradients).await?;
    let time_no_overlap = start_no_overlap.elapsed();

    // Reset gradients
    let gradient = Tensor::ones(&shape_dims, &Device::cuda(0).unwrap_or(Device::default()))?;
    gradients = vec![gradient.clone(), gradient.clone(), gradient];

    // Test with overlap
    trainer.scaling_optimizer.communication_overlap_enabled = true;
    let start_with_overlap = std::time::Instant::now();
    trainer.synchronize_gradients(&mut gradients).await?;
    let time_with_overlap = start_with_overlap.elapsed();

    println!(
        "Without overlap: {:.2}ms, With overlap: {:.2}ms",
        time_no_overlap.as_millis(),
        time_with_overlap.as_millis()
    );

    // With proper implementation, overlap should be faster or similar
    // For simulation, we just verify both complete successfully
    assert!(
        time_no_overlap.as_millis() < 5000,
        "Sync should complete in reasonable time"
    );
    assert!(
        time_with_overlap.as_millis() < 5000,
        "Overlapped sync should complete in reasonable time"
    );

    Ok(())
}

/// Test ring AllReduce with gradient compression
#[tokio::test]
async fn test_ring_allreduce_compression() -> Result<()> {
    let world_size = 4;
    let rank = 0;

    let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
    trainer.scaling_optimizer.communication_overlap_enabled = true;

    // Create gradients for compression testing
    let shape_dims = vec![10_000];
    let gradient = Tensor::full(
        &shape_dims,
        3.14159,
        &Device::cuda(0).unwrap_or(Device::default()),
    )?;
    let mut gradients = vec![gradient];

    // Test without compression
    trainer.scaling_optimizer.gradient_compression = false;
    let start_uncompressed = std::time::Instant::now();
    trainer.synchronize_gradients(&mut gradients).await?;
    let time_uncompressed = start_uncompressed.elapsed();

    // Reset gradient
    let gradient = Tensor::full(
        &shape_dims,
        3.14159,
        &Device::cuda(0).unwrap_or(Device::default()),
    )?;
    gradients = vec![gradient];

    // Test with compression
    trainer.scaling_optimizer.gradient_compression = true;
    let start_compressed = std::time::Instant::now();
    trainer.synchronize_gradients(&mut gradients).await?;
    let time_compressed = start_compressed.elapsed();

    println!(
        "Uncompressed: {:.2}ms, Compressed: {:.2}ms",
        time_uncompressed.as_millis(),
        time_compressed.as_millis()
    );

    // Verify correctness is maintained with compression
    let data = gradients[0].data()?;
    for &value in &data {
        // With compression, there might be small precision loss
        assert!(
            (value - data[0]).abs() < 1e-4,
            "Compression should maintain reasonable precision"
        );
    }

    Ok(())
}

/// Test ring AllReduce error handling and recovery
#[tokio::test]
async fn test_ring_allreduce_error_handling() -> Result<()> {
    let world_size = 4;
    let rank = 0;

    let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
    trainer.scaling_optimizer.communication_overlap_enabled = true;

    // Test with empty gradients
    let mut empty_gradients = Vec::new();
    let result = trainer.synchronize_gradients(&mut empty_gradients).await;
    assert!(
        result.is_ok(),
        "Empty gradients should be handled gracefully"
    );

    // Test with very large gradients (potential memory issues)
    let large_shape = vec![10_000_000]; // 10M elements
    let large_gradient_result =
        Tensor::ones(&large_shape, &Device::cuda(0).unwrap_or(Device::default()));

    if let Ok(large_gradient) = large_gradient_result {
        let mut large_gradients = vec![large_gradient];
        let result = trainer.synchronize_gradients(&mut large_gradients).await;
        // Should either succeed or fail gracefully
        match result {
            Ok(_) => println!("Large gradient sync succeeded"),
            Err(e) => println!("Large gradient sync failed gracefully: {}", e),
        }
    }

    Ok(())
}

/// Test ring AllReduce with multiple communication streams
#[tokio::test]
async fn test_ring_allreduce_multi_stream() -> Result<()> {
    let world_size = 4;
    let rank = 0;

    let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
    trainer.scaling_optimizer.communication_overlap_enabled = true;

    // Test different numbers of communication streams
    let stream_counts = vec![1, 2, 4, 8];

    for stream_count in stream_counts {
        trainer.scaling_optimizer.communication_streams = stream_count;
        trainer.scaling_optimizer.compute_streams = stream_count * 2;

        // Create multiple gradients to utilize multiple streams
        let shape_dims = vec![5000];
        let gradients_count = stream_count;
        let mut gradients = Vec::new();

        for i in 0..gradients_count {
            let gradient = Tensor::full(
                &shape_dims,
                (i + 1) as f32,
                &Device::cuda(0).unwrap_or(Device::default()),
            )?;
            gradients.push(gradient);
        }

        let start_time = std::time::Instant::now();
        trainer.synchronize_gradients(&mut gradients).await?;
        let duration = start_time.elapsed();

        println!(
            "Ring AllReduce with {} streams: {:.2}ms",
            stream_count,
            duration.as_millis()
        );

        // Verify all gradients were processed correctly
        for (i, gradient) in gradients.iter().enumerate() {
            let data = gradient.data()?;
            for &value in &data {
                assert!(
                    (value - data[0]).abs() < 1e-6,
                    "Gradient {} values should be consistent",
                    i
                );
            }
        }
    }

    Ok(())
}

/// Benchmark ring AllReduce performance vs sequential AllReduce
#[tokio::test]
async fn benchmark_ring_vs_sequential_allreduce() -> Result<()> {
    let world_size = 8;
    let rank = 0;

    // Test different gradient sizes
    let sizes = vec![1_000, 10_000, 100_000];

    for size in sizes {
        let mut trainer = MultiGpuTrainer::new(world_size, rank).await?;
        let shape_dims = vec![size];

        // Benchmark sequential AllReduce
        trainer.scaling_optimizer.communication_overlap_enabled = false;
        let gradient = Tensor::ones(&shape_dims, &Device::cuda(0).unwrap_or(Device::default()))?;
        let mut gradients = vec![gradient];

        let start_sequential = std::time::Instant::now();
        trainer.synchronize_gradients(&mut gradients).await?;
        let time_sequential = start_sequential.elapsed();

        // Benchmark ring AllReduce
        trainer.scaling_optimizer.communication_overlap_enabled = true;
        let gradient = Tensor::ones(&shape_dims, &Device::cuda(0).unwrap_or(Device::default()))?;
        let mut gradients = vec![gradient];

        let start_ring = std::time::Instant::now();
        trainer.synchronize_gradients(&mut gradients).await?;
        let time_ring = start_ring.elapsed();

        let speedup = time_sequential.as_micros() as f64 / time_ring.as_micros() as f64;

        println!("Size: {} elements", size);
        println!("  Sequential: {:.2}ms", time_sequential.as_millis());
        println!("  Ring:       {:.2}ms", time_ring.as_millis());
        println!("  Speedup:    {:.2}x", speedup);

        // Both should complete in reasonable time
        assert!(
            time_sequential.as_millis() < 1000,
            "Sequential should be reasonably fast"
        );
        assert!(
            time_ring.as_millis() < 1000,
            "Ring should be reasonably fast"
        );
    }

    Ok(())
}

/// Verify that the overlapped (ring AllReduce) gradient synchronization path
/// produces the same result as the sequential path.
///
/// Previously, ring_allreduce() contained its own simulation that bypassed the
/// ProcessGroup dispatch, meaning it would NOT use NCCL/RNCCL even when compiled in.
/// This test ensures both paths delegate to the same ProcessGroup allreduce logic.
#[tokio::test]
async fn test_ring_allreduce_matches_sequential_path() -> Result<()> {
    let world_size = 4;
    let input_value = 3.0_f32;
    let shape_dims = vec![64usize];
    let device = Device::cuda(0).unwrap_or(Device::default());

    // Sequential path
    let mut trainer_seq = MultiGpuTrainer::new(world_size, 0).await?;
    trainer_seq.scaling_optimizer.communication_overlap_enabled = false;
    let mut grad_seq = vec![Tensor::full(&shape_dims, input_value, &device)?];
    trainer_seq.synchronize_gradients(&mut grad_seq).await?;
    let seq_data = grad_seq[0].data()?;

    // Ring AllReduce (overlapped) path
    let mut trainer_ring = MultiGpuTrainer::new(world_size, 0).await?;
    trainer_ring.scaling_optimizer.communication_overlap_enabled = true;
    let mut grad_ring = vec![Tensor::full(&shape_dims, input_value, &device)?];
    trainer_ring.synchronize_gradients(&mut grad_ring).await?;
    let ring_data = grad_ring[0].data()?;

    assert_eq!(seq_data.len(), ring_data.len(), "output lengths must match");
    for (i, (&s, &r)) in seq_data.iter().zip(ring_data.iter()).enumerate() {
        assert!(
            (s - r).abs() < 1e-5,
            "element {i}: sequential={s}, ring={r} -- overlapped path must match sequential"
        );
    }
    Ok(())
}