//! Comprehensive TDD tests for real data pipeline
#![cfg(feature = "disabled_tests")]

use crate::{PreprocessingError, Result, real_data_pipeline::*};

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

    #[test]
    fn test_tensor_creation() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![2, 3]);
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];

        let tensor = RealTensor::from_data(data, shape, device).unwrap();
        assert_eq!(tensor.shape().dims(), &[2, 3]);
        assert_eq!(tensor.numel(), 6);
    }

    #[test]
    fn test_tensor_zeros() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![3, 4]);

        let tensor = RealTensor::zeros(&shape, &device).unwrap();
        assert_eq!(tensor.shape().dims(), &[3, 4]);
        assert_eq!(tensor.numel(), 12);

        let data = tensor.to_vec().unwrap();
        assert!(data.iter().all(|&x| x == 0.0));
    }

    #[test]
    fn test_tensor_ones() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![2, 2]);

        let tensor = RealTensor::ones(&shape, &device).unwrap();
        let data = tensor.to_vec().unwrap();
        assert!(data.iter().all(|&x| x == 1.0));
    }

    #[test]
    fn test_tensor_randn() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![100]);

        let tensor = RealTensor::randn(&shape, &device).unwrap();
        let data = tensor.to_vec().unwrap();

        // Check that we have normal distribution characteristics
        let mean: f32 = data.iter().sum::<f32>() / data.len() as f32;
        let variance: f32 =
            data.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / data.len() as f32;

        // Should be close to N(0,1)
        assert!(mean.abs() < 0.3); // Allow some tolerance
        assert!(variance > 0.5 && variance < 2.0);
    }

    #[test]
    fn test_tensor_add() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![2, 2]);

        let a = RealTensor::ones(&shape, &device).unwrap();
        let b = RealTensor::ones(&shape, &device).unwrap();
        let c = a.add(&b).unwrap();

        let data = c.to_vec().unwrap();
        assert!(data.iter().all(|&x| x == 2.0));
    }

    #[test]
    fn test_tensor_sub() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![2, 2]);

        let a = RealTensor::ones(&shape, &device).unwrap();
        let b = RealTensor::ones(&shape, &device).unwrap();
        let c = a.sub(&b).unwrap();

        let data = c.to_vec().unwrap();
        assert!(data.iter().all(|&x| x == 0.0));
    }

    #[test]
    fn test_tensor_mul() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![2, 2]);

        let a =
            RealTensor::from_data(vec![2.0, 3.0, 4.0, 5.0], shape.clone(), device.clone()).unwrap();
        let b = RealTensor::from_data(vec![1.0, 2.0, 3.0, 4.0], shape, device).unwrap();
        let c = a.mul(&b).unwrap();

        let data = c.to_vec().unwrap();
        assert_eq!(data, vec![2.0, 6.0, 12.0, 20.0]);
    }

    #[test]
    fn test_tensor_div() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![2, 2]);

        let a = RealTensor::from_data(vec![4.0, 6.0, 8.0, 10.0], shape.clone(), device.clone())
            .unwrap();
        let b = RealTensor::from_data(vec![2.0, 2.0, 2.0, 2.0], shape, device).unwrap();
        let c = a.div(&b).unwrap();

        let data = c.to_vec().unwrap();
        assert_eq!(data, vec![2.0, 3.0, 4.0, 5.0]);
    }

    #[test]
    fn test_tensor_matmul() {
        let device = RealDevice::cpu();

        let a = RealTensor::from_data(
            vec![1.0, 2.0, 3.0, 4.0],
            RealShape::new(vec![2, 2]),
            device.clone(),
        )
        .unwrap();

        let b = RealTensor::from_data(vec![5.0, 6.0, 7.0, 8.0], RealShape::new(vec![2, 2]), device)
            .unwrap();

        let c = a.matmul(&b).unwrap();
        let data = c.to_vec().unwrap();

        // [1 2] [5 6]   [19 22]
        // [3 4] [7 8] = [43 50]
        assert_eq!(data, vec![19.0, 22.0, 43.0, 50.0]);
    }

    #[test]
    fn test_tensor_reshape() {
        let device = RealDevice::cpu();
        let original_shape = RealShape::new(vec![2, 3]);
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];

        let tensor = RealTensor::from_data(data, original_shape, device).unwrap();
        let reshaped = tensor.reshape(&[3, 2]).unwrap();

        assert_eq!(reshaped.shape().dims(), &[3, 2]);
        assert_eq!(reshaped.numel(), 6);

        let reshaped_data = reshaped.to_vec().unwrap();
        assert_eq!(reshaped_data, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
    }

    #[test]
    fn test_tensor_transpose() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![2, 3]);
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];

        let tensor = RealTensor::from_data(data, shape, device).unwrap();
        let transposed = tensor.transpose().unwrap();

        assert_eq!(transposed.shape().dims(), &[3, 2]);

        let transposed_data = transposed.to_vec().unwrap();
        // Original: [1 2 3]   Transposed: [1 4]
        //           [4 5 6]                [2 5]
        //                                 [3 6]
        assert_eq!(transposed_data, vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
    }

    #[test]
    fn test_tensor_mean() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![2, 3]);
        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];

        let tensor = RealTensor::from_data(data, shape, device).unwrap();
        let mean = tensor.mean(None).unwrap();

        let mean_data = mean.to_vec().unwrap();
        assert_eq!(mean_data, vec![3.5]); // (1+2+3+4+5+6)/6 = 3.5
    }

    #[test]
    fn test_tensor_std() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![4]);
        let data = vec![1.0, 2.0, 3.0, 4.0];

        let tensor = RealTensor::from_data(data, shape, device).unwrap();
        let std = tensor.std(None).unwrap();

        let std_data = std.to_vec().unwrap();
        // Standard deviation of [1,2,3,4] ≈ 1.29
        assert!((std_data[0] - 1.29).abs() < 0.1);
    }

    #[test]
    fn test_tensor_min_max() {
        let device = RealDevice::cpu();
        let shape = RealShape::new(vec![3]);
        let data = vec![3.0, 1.0, 4.0];

        let tensor = RealTensor::from_data(data, shape, device).unwrap();

        let min = tensor.min(None).unwrap();
        let max = tensor.max(None).unwrap();

        assert_eq!(min.to_vec().unwrap(), vec![1.0]);
        assert_eq!(max.to_vec().unwrap(), vec![4.0]);
    }

    #[test]
    fn test_tensor_device_transfer() {
        let cpu = RealDevice::cpu();
        let cuda = RealDevice::cuda(0);

        let shape = RealShape::new(vec![2, 2]);
        let tensor = RealTensor::ones(&shape, &cpu).unwrap();

        // Transfer to CUDA (will fallback to CPU if no GPU)
        let gpu_tensor = tensor.to_device(&cuda).unwrap();
        assert_eq!(gpu_tensor.device(), &cuda);

        // Transfer back to CPU
        let cpu_tensor = gpu_tensor.to_device(&cpu).unwrap();
        assert_eq!(cpu_tensor.device(), &cpu);

        let data = cpu_tensor.to_vec().unwrap();
        assert!(data.iter().all(|&x| x == 1.0));
    }

    #[test]
    fn test_tensor_error_cases() {
        let device = RealDevice::cpu();

        // Mismatched shapes for operations
        let a = RealTensor::ones(&RealShape::new(vec![2, 2]), &device).unwrap();
        let b = RealTensor::ones(&RealShape::new(vec![3, 3]), &device).unwrap();

        assert!(a.add(&b).is_err());
        assert!(a.mul(&b).is_err());

        // Invalid reshape
        let tensor = RealTensor::ones(&RealShape::new(vec![2, 3]), &device).unwrap();
        assert!(tensor.reshape(&[2, 2]).is_err()); // 6 elements can't fit in 2x2

        // Invalid matmul dimensions
        let a = RealTensor::ones(&RealShape::new(vec![2, 3]), &device).unwrap();
        let b = RealTensor::ones(&RealShape::new(vec![2, 3]), &device).unwrap(); // Should be 3x?
        assert!(a.matmul(&b).is_err());
    }
}

#[cfg(test)]
mod data_loader_tests {
    use super::*;
    use std::path::PathBuf;

    #[tokio::test]
    async fn test_data_loader_creation() {
        let config = DataLoaderConfig::default();
        let loader = RealDataLoader::new(config);

        assert_eq!(loader.config().batch_size, 32);
        assert!(!loader.config().shuffle);
    }

    #[tokio::test]
    async fn test_data_loader_from_memory() {
        let device = RealDevice::cpu();

        // Create sample data
        let mut samples = Vec::new();
        for i in 0..100 {
            let data = vec![i as f32; 10];
            let tensor =
                RealTensor::from_data(data, RealShape::new(vec![10]), device.clone()).unwrap();
            samples.push(tensor);
        }

        let mut config = DataLoaderConfig::default();
        config.batch_size = 10;

        let mut loader = RealDataLoader::new(config);
        loader.load_from_memory(samples).await.unwrap();

        assert_eq!(loader.len(), 100);
        assert_eq!(loader.num_batches(), 10);
    }

    #[tokio::test]
    async fn test_data_loader_iteration() {
        let device = RealDevice::cpu();

        // Create sample data
        let mut samples = Vec::new();
        for i in 0..20 {
            let data = vec![i as f32; 5];
            let tensor =
                RealTensor::from_data(data, RealShape::new(vec![5]), device.clone()).unwrap();
            samples.push(tensor);
        }

        let mut config = DataLoaderConfig::default();
        config.batch_size = 5;

        let mut loader = RealDataLoader::new(config);
        loader.load_from_memory(samples).await.unwrap();

        let mut batch_count = 0;
        while let Some(batch) = loader.next_batch().await.unwrap() {
            assert_eq!(batch.len(), 5); // batch_size
            batch_count += 1;
        }

        assert_eq!(batch_count, 4); // 20 samples / 5 batch_size = 4 batches
    }

    #[tokio::test]
    async fn test_data_loader_shuffle() {
        let device = RealDevice::cpu();

        // Create sample data with unique identifiers
        let mut samples = Vec::new();
        for i in 0..10 {
            let data = vec![i as f32; 1];
            let tensor =
                RealTensor::from_data(data, RealShape::new(vec![1]), device.clone()).unwrap();
            samples.push(tensor);
        }

        let mut config = DataLoaderConfig::default();
        config.batch_size = 1;
        config.shuffle = true;

        let mut loader = RealDataLoader::new(config);
        loader.load_from_memory(samples).await.unwrap();

        // Get first epoch
        let mut first_epoch = Vec::new();
        while let Some(batch) = loader.next_batch().await.unwrap() {
            let data = batch[0].to_vec().unwrap();
            first_epoch.push(data[0] as i32);
        }

        // Get second epoch
        let mut second_epoch = Vec::new();
        while let Some(batch) = loader.next_batch().await.unwrap() {
            let data = batch[0].to_vec().unwrap();
            second_epoch.push(data[0] as i32);
        }

        // Epochs should have same elements but potentially different order
        first_epoch.sort();
        second_epoch.sort();
        assert_eq!(first_epoch, second_epoch);

        // Note: Due to randomness, we can't guarantee order is different,
        // but this tests the shuffle functionality works
    }

    #[tokio::test]
    async fn test_data_loader_validation() {
        let device = RealDevice::cpu();

        // Create valid data
        let mut samples = Vec::new();
        for i in 0..10 {
            let data = vec![i as f32; 3];
            let tensor =
                RealTensor::from_data(data, RealShape::new(vec![3]), device.clone()).unwrap();
            samples.push(tensor);
        }

        let config = DataLoaderConfig::default();
        let mut loader = RealDataLoader::new(config);
        loader.load_from_memory(samples).await.unwrap();

        let validation_report = loader.validate_data().await.unwrap();
        assert!(validation_report.is_valid);
        assert_eq!(validation_report.total_samples, 10);
        assert_eq!(validation_report.errors.len(), 0);
    }

    #[tokio::test]
    async fn test_data_loader_statistics() {
        let device = RealDevice::cpu();

        // Create sample data
        let mut samples = Vec::new();
        for i in 0..50 {
            let data = vec![i as f32; 2];
            let tensor =
                RealTensor::from_data(data, RealShape::new(vec![2]), device.clone()).unwrap();
            samples.push(tensor);
        }

        let mut config = DataLoaderConfig::default();
        config.batch_size = 10;

        let mut loader = RealDataLoader::new(config);
        loader.load_from_memory(samples).await.unwrap();

        // Process some batches
        for _ in 0..3 {
            let _ = loader.next_batch().await.unwrap();
        }

        let stats = loader.statistics().await;
        assert_eq!(stats.total_samples, 50);
        assert_eq!(stats.batches_processed, 3);
        assert!(stats.processing_time_ms > 0.0);
    }

    #[tokio::test]
    async fn test_data_loader_error_cases() {
        let config = DataLoaderConfig::default();
        let mut loader = RealDataLoader::new(config);

        // Try to get batch without loading data
        assert!(loader.next_batch().await.is_err());

        // Try to validate without data
        assert!(loader.validate_data().await.is_err());
    }
}

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

    #[test]
    fn test_device_creation() {
        let cpu = RealDevice::cpu();
        assert_eq!(cpu, RealDevice::cuda(0).unwrap_or(Device::default()));

        let cuda = RealDevice::cuda(0);
        assert_eq!(cuda, RealDevice::Cuda(0));
    }

    #[test]
    fn test_device_properties() {
        let cpu = RealDevice::cpu();
        assert!(!cpu.is_cuda());
        assert_eq!(cpu.device_id(), None);

        let cuda = RealDevice::cuda(1);
        assert!(cuda.is_cuda());
        assert_eq!(cuda.device_id(), Some(1));
    }

    #[test]
    fn test_device_display() {
        let cpu = RealDevice::cpu();
        assert_eq!(format!("{}", cpu), "cpu");

        let cuda = RealDevice::cuda(2);
        assert_eq!(format!("{}", cuda), "cuda:2");
    }
}

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

    #[test]
    fn test_shape_creation() {
        let shape = RealShape::new(vec![2, 3, 4]);
        assert_eq!(shape.dims(), &[2, 3, 4]);
        assert_eq!(shape.ndim(), 3);
        assert_eq!(shape.numel(), 24);
    }

    #[test]
    fn test_shape_properties() {
        let scalar_shape = RealShape::new(vec![]);
        assert!(scalar_shape.is_scalar());
        assert!(!scalar_shape.is_vector());
        assert!(!scalar_shape.is_matrix());

        let vector_shape = RealShape::new(vec![5]);
        assert!(!vector_shape.is_scalar());
        assert!(vector_shape.is_vector());
        assert!(!vector_shape.is_matrix());

        let matrix_shape = RealShape::new(vec![3, 4]);
        assert!(!matrix_shape.is_scalar());
        assert!(!matrix_shape.is_vector());
        assert!(matrix_shape.is_matrix());
    }

    #[test]
    fn test_shape_validation() {
        let valid_shape = RealShape::new(vec![2, 3, 4]);
        assert!(valid_shape.validate().is_ok());

        let empty_shape = RealShape::new(vec![2, 0, 4]);
        assert!(empty_shape.validate().is_err());

        let zero_shape = RealShape::new(vec![]);
        assert!(zero_shape.validate().is_ok()); // Scalars are valid
    }

    #[test]
    fn test_shape_conversions() {
        let shape1 = RealShape::from(vec![1, 2, 3]);
        assert_eq!(shape1.dims(), &[1, 2, 3]);

        let shape2 = RealShape::from(&[4, 5, 6][..]);
        assert_eq!(shape2.dims(), &[4, 5, 6]);

        let shape3 = RealShape::from([7, 8, 9]);
        assert_eq!(shape3.dims(), &[7, 8, 9]);
    }
}

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

    #[tokio::test]
    async fn test_tensor_data_loader_integration() {
        let device = RealDevice::cpu();

        // Create tensors for training
        let mut training_data = Vec::new();
        for i in 0..100 {
            let features = vec![i as f32, (i * 2) as f32, (i * 3) as f32];
            let tensor =
                RealTensor::from_data(features, RealShape::new(vec![3]), device.clone()).unwrap();
            training_data.push(tensor);
        }

        // Set up data loader
        let mut config = DataLoaderConfig::default();
        config.batch_size = 10;
        config.shuffle = true;

        let mut loader = RealDataLoader::new(config);
        loader.load_from_memory(training_data).await.unwrap();

        // Simulate training loop
        let mut total_processed = 0;
        while let Some(batch) = loader.next_batch().await.unwrap() {
            assert_eq!(batch.len(), 10);

            // Process each tensor in batch
            for tensor in batch {
                assert_eq!(tensor.shape().dims(), &[3]);
                let data = tensor.to_vec().unwrap();
                assert_eq!(data.len(), 3);

                // Verify data relationships
                assert_eq!(data[1], data[0] * 2.0);
                assert_eq!(data[2], data[0] * 3.0);
            }

            total_processed += 10;
        }

        assert_eq!(total_processed, 100);
    }

    #[tokio::test]
    async fn test_tensor_operations_pipeline() {
        let device = RealDevice::cpu();

        // Create input data
        let input_shape = RealShape::new(vec![4, 4]);
        let input_data: Vec<f32> = (0..16).map(|x| x as f32).collect();
        let input = RealTensor::from_data(input_data, input_shape, device.clone()).unwrap();

        // Apply sequence of operations (simulating a preprocessing pipeline)

        // 1. Normalize (subtract mean, divide by std)
        let mean = input.mean(None).unwrap();
        let std = input.std(None).unwrap();

        let mean_broadcast = RealTensor::from_data(
            vec![mean.to_vec().unwrap()[0]; 16],
            RealShape::new(vec![4, 4]),
            device.clone(),
        )
        .unwrap();

        let std_broadcast = RealTensor::from_data(
            vec![std.to_vec().unwrap()[0]; 16],
            RealShape::new(vec![4, 4]),
            device.clone(),
        )
        .unwrap();

        let normalized = input
            .sub(&mean_broadcast)
            .unwrap()
            .div(&std_broadcast)
            .unwrap();

        // 2. Reshape to vector
        let flattened = normalized.reshape(&[16]).unwrap();

        // 3. Apply linear transformation (matrix multiplication)
        let weights = RealTensor::randn(&RealShape::new(vec![16, 8]), &device).unwrap();
        let features = weights
            .transpose()
            .unwrap()
            .matmul(&flattened.reshape(&[16, 1]).unwrap())
            .unwrap();

        // 4. Apply activation (ReLU approximation)
        let zero = RealTensor::zeros(&features.shape(), &device).unwrap();
        let activated = features.max_elementwise(&zero).unwrap();

        // Verify final shape and properties
        assert_eq!(activated.shape().dims(), &[8, 1]);
        let final_data = activated.to_vec().unwrap();
        assert!(final_data.iter().all(|&x| x >= 0.0)); // ReLU ensures non-negative
    }

    #[tokio::test]
    async fn test_performance_benchmark() {
        let device = RealDevice::cpu();

        // Create large tensors for performance testing
        let size = 1000;
        let shape = RealShape::new(vec![size, size]);

        let start = std::time::Instant::now();

        let a = RealTensor::randn(&shape, &device).unwrap();
        let b = RealTensor::randn(&shape, &device).unwrap();

        let creation_time = start.elapsed();

        let start = std::time::Instant::now();
        let c = a.add(&b).unwrap();
        let addition_time = start.elapsed();

        let start = std::time::Instant::now();
        let d = a.mul(&b).unwrap();
        let multiplication_time = start.elapsed();

        println!("Performance benchmark results:");
        println!("Tensor creation ({}x{}): {:?}", size, size, creation_time);
        println!("Addition: {:?}", addition_time);
        println!("Multiplication: {:?}", multiplication_time);

        // Basic sanity checks
        assert_eq!(c.shape().dims(), &[size, size]);
        assert_eq!(d.shape().dims(), &[size, size]);
        assert!(creation_time.as_millis() < 5000); // Should complete within 5 seconds
        assert!(addition_time.as_millis() < 1000);
        assert!(multiplication_time.as_millis() < 1000);
    }
}