//! Comprehensive tests for magnitude-based unstructured pruning
//!
//! Tests cover:
//! - Global and layer-wise magnitude pruning
//! - Gradual pruning with different schedules
#![cfg(feature = "disabled_tests")]
//! - Top-k sparsity patterns
//! - Gradient-based importance (SNIP)
//! - Random and structured-random pruning
//! - Fine-tuning with sparse gradients
//! - Sparse-to-sparse training continuation

use rtx_autograd::{no_grad, tensor_with_grad};
use rtx_compress::{
    Result,
    pruning::{
        GradualPruningSchedule, PruningCriterion, PruningMask, SparsityPattern, UnstructuredPruner,
        UnstructuredPruningConfig,
    },
};
use rtx_tensor::{Device, Tensor};
use std::collections::HashMap;

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

    fn create_test_tensor(shape: &[usize], seed: u64) -> Result<Tensor> {
        let device = Device::try_default()?;
        let mut data = Vec::new();
        let total_elements: usize = shape.iter().product();

        // Create deterministic pseudo-random data
        for i in 0..total_elements {
            let val = (((i as u64).wrapping_mul(seed) % 1000) as f32 / 1000.0 - 0.5) * 2.0;
            data.push(val);
        }

        Tensor::from_slice(&data, shape, &device)
    }

    fn create_test_model(seed: u64) -> Result<HashMap<String, Tensor>> {
        let mut model = HashMap::new();

        model.insert(
            "conv1.weight".to_string(),
            create_test_tensor(&[64, 3, 3, 3], seed)?,
        );
        model.insert(
            "conv1.bias".to_string(),
            create_test_tensor(&[64], seed + 1)?,
        );
        model.insert(
            "conv2.weight".to_string(),
            create_test_tensor(&[128, 64, 3, 3], seed + 2)?,
        );
        model.insert(
            "conv2.bias".to_string(),
            create_test_tensor(&[128], seed + 3)?,
        );
        model.insert(
            "fc1.weight".to_string(),
            create_test_tensor(&[256, 512], seed + 4)?,
        );
        model.insert(
            "fc1.bias".to_string(),
            create_test_tensor(&[256], seed + 5)?,
        );
        model.insert(
            "fc2.weight".to_string(),
            create_test_tensor(&[10, 256], seed + 6)?,
        );
        model.insert("fc2.bias".to_string(), create_test_tensor(&[10], seed + 7)?);

        Ok(model)
    }

    fn calculate_sparsity(tensor: &Tensor) -> Result<f32> {
        let values = tensor.to_vec::<f32>()?;
        let zero_count = values.iter().filter(|&&x| x.abs() < 1e-8).count();
        Ok(zero_count as f32 / values.len() as f32)
    }

    #[test]
    fn test_global_magnitude_pruning() -> Result<()> {
        let config = UnstructuredPruningConfig::new(
            PruningCriterion::GlobalMagnitude,
            0.8, // 80% sparsity
        );

        let pruner = UnstructuredPruner::new(config)?;
        let model = create_test_model(42)?;

        let pruned_model = pruner.prune_model(&model)?;

        // Check that all tensors maintain their shape
        for (name, original) in &model {
            let pruned = &pruned_model[name];
            assert_eq!(
                original.shape(),
                pruned.shape(),
                "Shape mismatch for {}",
                name
            );
        }

        // Calculate overall sparsity
        let mut total_elements = 0;
        let mut total_zeros = 0;

        for (_, tensor) in &pruned_model {
            let values = tensor.to_vec::<f32>()?;
            total_elements += values.len();
            total_zeros += values.iter().filter(|&&x| x.abs() < 1e-8).count();
        }

        let overall_sparsity = total_zeros as f32 / total_elements as f32;
        assert!(
            (overall_sparsity - 0.8).abs() < 0.01,
            "Expected ~80% sparsity, got {:.2}%",
            overall_sparsity * 100.0
        );

        Ok(())
    }

    #[test]
    fn test_layer_wise_magnitude_pruning() -> Result<()> {
        let mut layer_sparsities = HashMap::new();
        layer_sparsities.insert("conv1.weight".to_string(), 0.5);
        layer_sparsities.insert("conv2.weight".to_string(), 0.7);
        layer_sparsities.insert("fc1.weight".to_string(), 0.9);
        layer_sparsities.insert("fc2.weight".to_string(), 0.6);

        let config = UnstructuredPruningConfig::new_layer_wise(
            PruningCriterion::LocalMagnitude,
            layer_sparsities,
        );

        let pruner = UnstructuredPruner::new(config)?;
        let model = create_test_model(42)?;

        let pruned_model = pruner.prune_model(&model)?;

        // Check individual layer sparsities
        let conv1_sparsity = calculate_sparsity(&pruned_model["conv1.weight"])?;
        let conv2_sparsity = calculate_sparsity(&pruned_model["conv2.weight"])?;
        let fc1_sparsity = calculate_sparsity(&pruned_model["fc1.weight"])?;
        let fc2_sparsity = calculate_sparsity(&pruned_model["fc2.weight"])?;

        assert!(
            (conv1_sparsity - 0.5).abs() < 0.02,
            "Conv1 sparsity: {:.3}",
            conv1_sparsity
        );
        assert!(
            (conv2_sparsity - 0.7).abs() < 0.02,
            "Conv2 sparsity: {:.3}",
            conv2_sparsity
        );
        assert!(
            (fc1_sparsity - 0.9).abs() < 0.02,
            "FC1 sparsity: {:.3}",
            fc1_sparsity
        );
        assert!(
            (fc2_sparsity - 0.6).abs() < 0.02,
            "FC2 sparsity: {:.3}",
            fc2_sparsity
        );

        Ok(())
    }

    #[test]
    fn test_gradual_pruning_schedule() -> Result<()> {
        let schedule = GradualPruningSchedule::polynomial(
            0.0,  // Start with no pruning
            0.9,  // End with 90% sparsity
            1000, // Over 1000 steps
            3.0,  // Cubic schedule
        );

        let config = UnstructuredPruningConfig::new_with_schedule(
            PruningCriterion::GlobalMagnitude,
            schedule,
        );

        let mut pruner = UnstructuredPruner::new(config)?;
        let tensor = create_test_tensor(&[128, 256], 42)?;

        // Test gradual pruning over multiple steps
        let mut sparsities = Vec::new();
        let steps = [0, 100, 300, 600, 999];

        for &step in &steps {
            pruner.set_current_step(step);
            let pruned = pruner.prune_tensor(&tensor, "test")?;
            let sparsity = calculate_sparsity(&pruned)?;
            sparsities.push(sparsity);
        }

        // Sparsity should increase monotonically
        for i in 1..sparsities.len() {
            assert!(
                sparsities[i] >= sparsities[i - 1],
                "Sparsity should increase: step {} = {:.3}, step {} = {:.3}",
                steps[i - 1],
                sparsities[i - 1],
                steps[i],
                sparsities[i]
            );
        }

        // Final sparsity should be close to target
        assert!(
            (sparsities.last().unwrap() - 0.9).abs() < 0.05,
            "Final sparsity should be ~90%, got {:.1}%",
            sparsities.last().unwrap() * 100.0
        );

        Ok(())
    }

    #[test]
    fn test_snip_gradient_based_pruning() -> Result<()> {
        let config = UnstructuredPruningConfig::new(
            PruningCriterion::SNIP, // Single-shot Network Pruning
            0.8,
        );

        let mut pruner = UnstructuredPruner::new(config)?;

        // Create tensor with gradients for SNIP
        let tensor = create_test_tensor(&[64, 128], 42)?;
        let gradient = create_test_tensor(&[64, 128], 123)?; // Simulate gradient

        let pruned = pruner.prune_tensor_with_gradient(&tensor, "test", &gradient)?;
        let sparsity = calculate_sparsity(&pruned)?;

        assert!(
            (sparsity - 0.8).abs() < 0.02,
            "SNIP sparsity should be ~80%, got {:.1}%",
            sparsity * 100.0
        );

        // Verify that SNIP preserves weights with high gradient * weight product
        let tensor_vals = tensor.to_vec::<f32>()?;
        let grad_vals = gradient.to_vec::<f32>()?;
        let pruned_vals = pruned.to_vec::<f32>()?;

        let mut preserved_high_importance = 0;
        let mut total_high_importance = 0;

        for ((w, g), p) in tensor_vals
            .iter()
            .zip(grad_vals.iter())
            .zip(pruned_vals.iter())
        {
            let importance = w.abs() * g.abs();
            if importance > 0.1 {
                // Arbitrary threshold for "high importance"
                total_high_importance += 1;
                if p.abs() > 1e-8 {
                    preserved_high_importance += 1;
                }
            }
        }

        if total_high_importance > 0 {
            let preservation_rate = preserved_high_importance as f32 / total_high_importance as f32;
            assert!(
                preservation_rate > 0.3,
                "SNIP should preserve high-importance weights better: {:.1}%",
                preservation_rate * 100.0
            );
        }

        Ok(())
    }

    #[test]
    fn test_top_k_sparsity_pattern() -> Result<()> {
        let config = UnstructuredPruningConfig::new_with_pattern(
            PruningCriterion::GlobalMagnitude,
            0.75,                              // 75% sparsity
            SparsityPattern::TopK { k: 1000 }, // Keep top 1000 weights per tensor
        );

        let pruner = UnstructuredPruner::new(config)?;
        let tensor = create_test_tensor(&[50, 100], 42)?; // 5000 elements total

        let pruned = pruner.prune_tensor(&tensor, "test")?;
        let values = pruned.to_vec::<f32>()?;

        // Count non-zero values
        let non_zero_count = values.iter().filter(|&&x| x.abs() > 1e-8).count();

        // Should keep approximately k values (allowing for ties)
        assert!(
            non_zero_count >= 900 && non_zero_count <= 1100,
            "Top-K should keep ~1000 values, got {}",
            non_zero_count
        );

        Ok(())
    }

    #[test]
    fn test_random_pruning() -> Result<()> {
        let config = UnstructuredPruningConfig::new(
            PruningCriterion::Random { seed: 42 },
            0.6, // 60% sparsity
        );

        let pruner = UnstructuredPruner::new(config)?;
        let tensor = create_test_tensor(&[32, 64], 42)?;

        let pruned1 = pruner.prune_tensor(&tensor, "test1")?;
        let pruned2 = pruner.prune_tensor(&tensor, "test2")?;

        // Both should have same sparsity
        let sparsity1 = calculate_sparsity(&pruned1)?;
        let sparsity2 = calculate_sparsity(&pruned2)?;

        assert!((sparsity1 - 0.6).abs() < 0.05);
        assert!((sparsity2 - 0.6).abs() < 0.05);

        // Should produce identical results (deterministic)
        let vals1 = pruned1.to_vec::<f32>()?;
        let vals2 = pruned2.to_vec::<f32>()?;

        let differences = vals1
            .iter()
            .zip(vals2.iter())
            .filter(|(&a, &b)| (a - b).abs() > 1e-6)
            .count();

        assert_eq!(
            differences, 0,
            "Random pruning should be deterministic with same seed"
        );

        Ok(())
    }

    #[test]
    fn test_structured_random_pruning() -> Result<()> {
        let config = UnstructuredPruningConfig::new_with_pattern(
            PruningCriterion::Random { seed: 42 },
            0.75,                                     // 75% sparsity
            SparsityPattern::Block { block_size: 4 }, // 4x4 blocks
        );

        let pruner = UnstructuredPruner::new(config)?;
        let tensor = create_test_tensor(&[32, 32], 42)?; // Square tensor for easy blocking

        let pruned = pruner.prune_tensor(&tensor, "test")?;
        let sparsity = calculate_sparsity(&pruned)?;

        assert!(
            (sparsity - 0.75).abs() < 0.1,
            "Block sparsity: {:.2}",
            sparsity
        );

        // Verify block structure (simplified check)
        let values = pruned.to_vec::<f32>()?;
        let mut block_zero_count = 0;
        let mut total_blocks = 0;

        // Check 4x4 blocks
        for block_row in 0..8 {
            // 32/4 = 8
            for block_col in 0..8 {
                let mut block_is_zero = true;
                for i in 0..4 {
                    for j in 0..4 {
                        let row = block_row * 4 + i;
                        let col = block_col * 4 + j;
                        let idx = row * 32 + col;
                        if values[idx].abs() > 1e-8 {
                            block_is_zero = false;
                            break;
                        }
                    }
                    if !block_is_zero {
                        break;
                    }
                }
                if block_is_zero {
                    block_zero_count += 1;
                }
                total_blocks += 1;
            }
        }

        let block_sparsity = block_zero_count as f32 / total_blocks as f32;
        // Block pruning should create some fully zero blocks
        assert!(
            block_sparsity > 0.1,
            "Should have some zero blocks: {:.2}",
            block_sparsity
        );

        Ok(())
    }

    #[test]
    fn test_magnitude_vs_gradient_pruning() -> Result<()> {
        let tensor = create_test_tensor(&[16, 32], 42)?;
        let gradient = create_test_tensor(&[16, 32], 123)?;

        // Magnitude-based pruning
        let mag_config = UnstructuredPruningConfig::new(PruningCriterion::GlobalMagnitude, 0.5);
        let mag_pruner = UnstructuredPruner::new(mag_config)?;
        let mag_pruned = mag_pruner.prune_tensor(&tensor, "test")?;

        // Gradient-based pruning (SNIP)
        let snip_config = UnstructuredPruningConfig::new(PruningCriterion::SNIP, 0.5);
        let mut snip_pruner = UnstructuredPruner::new(snip_config)?;
        let snip_pruned = snip_pruner.prune_tensor_with_gradient(&tensor, "test", &gradient)?;

        // Both should achieve similar sparsity
        let mag_sparsity = calculate_sparsity(&mag_pruned)?;
        let snip_sparsity = calculate_sparsity(&snip_pruned)?;

        assert!((mag_sparsity - 0.5).abs() < 0.05);
        assert!((snip_sparsity - 0.5).abs() < 0.05);

        // But should preserve different weights
        let mag_vals = mag_pruned.to_vec::<f32>()?;
        let snip_vals = snip_pruned.to_vec::<f32>()?;

        let different_preserved = mag_vals
            .iter()
            .zip(snip_vals.iter())
            .filter(|(&m, &s)| (m.abs() > 1e-8) != (s.abs() > 1e-8))
            .count();

        assert!(
            different_preserved > 100,
            "Magnitude and SNIP should preserve different weights: {} differences",
            different_preserved
        );

        Ok(())
    }

    #[test]
    fn test_pruning_mask_generation() -> Result<()> {
        let config = UnstructuredPruningConfig::new_with_mask_generation(
            PruningCriterion::LocalMagnitude,
            0.7,
            true, // Generate masks
        );

        let pruner = UnstructuredPruner::new(config)?;
        let tensor = create_test_tensor(&[24, 48], 42)?;

        let result = pruner.prune_tensor_with_mask(&tensor, "test")?;

        // Check that mask is binary
        let mask_values = result.mask.to_vec::<f32>()?;
        for &val in &mask_values {
            assert!(
                val == 0.0 || val == 1.0,
                "Mask should be binary, found {}",
                val
            );
        }

        // Check that pruned tensor equals original tensor masked
        let original_values = tensor.to_vec::<f32>()?;
        let pruned_values = result.pruned_tensor.to_vec::<f32>()?;

        for ((orig, pruned), &mask) in original_values
            .iter()
            .zip(pruned_values.iter())
            .zip(mask_values.iter())
        {
            if mask == 0.0 {
                assert_eq!(*pruned, 0.0, "Masked weights should be zero");
            } else {
                assert_eq!(*pruned, *orig, "Unmasked weights should be preserved");
            }
        }

        Ok(())
    }

    #[test]
    fn test_fine_tuning_mask_persistence() -> Result<()> {
        let config = UnstructuredPruningConfig::new_with_mask_generation(
            PruningCriterion::GlobalMagnitude,
            0.8,
            true,
        );

        let pruner = UnstructuredPruner::new(config)?;
        let tensor = create_test_tensor(&[16, 32], 42)?;

        let result = pruner.prune_tensor_with_mask(&tensor, "test")?;

        // Simulate fine-tuning by modifying the pruned tensor
        let updated_values: Vec<f32> = result
            .pruned_tensor
            .to_vec::<f32>()?
            .iter()
            .map(|&x| if x != 0.0 { x * 1.1 } else { x })
            .collect();

        let updated_tensor = Tensor::from_slice(
            &updated_values,
            tensor.shape().dims(),
            &Device::try_default()?,
        )?;

        // Apply mask again to ensure sparsity is maintained
        let remasked = updated_tensor.mul(&result.mask)?;
        let final_sparsity = calculate_sparsity(&remasked)?;

        assert!(
            (final_sparsity - 0.8).abs() < 0.05,
            "Sparsity should be maintained after fine-tuning: {:.2}",
            final_sparsity
        );

        Ok(())
    }

    #[test]
    fn test_threshold_based_pruning() -> Result<()> {
        let config = UnstructuredPruningConfig::new(
            PruningCriterion::Threshold { threshold: 0.1 },
            0.0, // Sparsity determined by threshold
        );

        let pruner = UnstructuredPruner::new(config)?;

        // Create tensor with known distribution
        let device = Device::try_default()?;
        let values: Vec<f32> = (0..1000)
            .map(|i| {
                if i < 200 {
                    0.05
                }
                // Small values (should be pruned)
                else if i < 800 {
                    0.15
                }
                // Medium values (should be kept)
                else {
                    0.25
                } // Large values (should be kept)
            })
            .collect();

        let tensor = Tensor::from_slice(&values, &[25, 40], &device)?;
        let pruned = pruner.prune_tensor(&tensor, "test")?;

        let pruned_vals = pruned.to_vec::<f32>()?;
        let kept_count = pruned_vals.iter().filter(|&&x| x.abs() > 1e-8).count();

        // Should keep ~800 values (those >= 0.1)
        assert!(
            kept_count >= 750 && kept_count <= 850,
            "Threshold pruning should keep ~800 values, got {}",
            kept_count
        );

        Ok(())
    }

    #[test]
    fn test_multi_step_gradual_pruning() -> Result<()> {
        let schedule = GradualPruningSchedule::exponential(
            0.2,  // Start at 20% sparsity
            0.95, // End at 95% sparsity
            500,  // Over 500 steps
        );

        let config = UnstructuredPruningConfig::new_with_schedule(
            PruningCriterion::GlobalMagnitude,
            schedule,
        );

        let mut pruner = UnstructuredPruner::new(config)?;
        let mut current_tensor = create_test_tensor(&[64, 64], 42)?;

        let mut sparsity_progression = Vec::new();
        let checkpoints = [0, 50, 125, 250, 375, 499];

        for &step in &checkpoints {
            pruner.set_current_step(step);
            current_tensor = pruner.prune_tensor(&current_tensor, "test")?;
            let sparsity = calculate_sparsity(&current_tensor)?;
            sparsity_progression.push((step, sparsity));
        }

        // Verify progression
        for i in 1..sparsity_progression.len() {
            let (prev_step, prev_sparsity) = sparsity_progression[i - 1];
            let (curr_step, curr_sparsity) = sparsity_progression[i];

            assert!(
                curr_sparsity >= prev_sparsity,
                "Sparsity should increase: step {} = {:.3}, step {} = {:.3}",
                prev_step,
                prev_sparsity,
                curr_step,
                curr_sparsity
            );
        }

        // Final sparsity should be close to target
        let (_, final_sparsity) = sparsity_progression.last().unwrap();
        assert!(
            (final_sparsity - 0.95).abs() < 0.03,
            "Final sparsity should be ~95%, got {:.2}%",
            final_sparsity * 100.0
        );

        Ok(())
    }

    #[test]
    fn test_layerwise_sensitivity_adjustment() -> Result<()> {
        // Create layer sensitivities (higher values = more sensitive = less pruning)
        let mut sensitivities = HashMap::new();
        sensitivities.insert("conv1.weight".to_string(), 0.9); // High sensitivity
        sensitivities.insert("conv2.weight".to_string(), 0.5); // Medium sensitivity
        sensitivities.insert("fc1.weight".to_string(), 0.1); // Low sensitivity

        let config = UnstructuredPruningConfig::new_with_sensitivity(
            PruningCriterion::LocalMagnitude,
            0.7, // Base 70% sparsity
            sensitivities,
        );

        let pruner = UnstructuredPruner::new(config)?;
        let model = create_test_model(42)?;

        let pruned_model = pruner.prune_model(&model)?;

        // High sensitivity layer should have lower sparsity
        let conv1_sparsity = calculate_sparsity(&pruned_model["conv1.weight"])?;
        let conv2_sparsity = calculate_sparsity(&pruned_model["conv2.weight"])?;
        let fc1_sparsity = calculate_sparsity(&pruned_model["fc1.weight"])?;

        assert!(
            conv1_sparsity < conv2_sparsity,
            "High sensitivity layer should have lower sparsity: {:.2} vs {:.2}",
            conv1_sparsity,
            conv2_sparsity
        );

        assert!(
            conv2_sparsity < fc1_sparsity,
            "Medium sensitivity should have lower sparsity than low sensitivity: {:.2} vs {:.2}",
            conv2_sparsity,
            fc1_sparsity
        );

        Ok(())
    }

    #[test]
    fn test_sparse_to_sparse_continuation() -> Result<()> {
        // Start with already sparse tensor
        let mut initial_values = vec![0.0f32; 1024];
        for i in (0..1024).step_by(3) {
            // Every 3rd element is non-zero
            initial_values[i] = ((i % 100) as f32) / 100.0;
        }

        let device = Device::try_default()?;
        let sparse_tensor = Tensor::from_slice(&initial_values, &[32, 32], &device)?;
        let initial_sparsity = calculate_sparsity(&sparse_tensor)?;

        assert!(
            initial_sparsity > 0.6,
            "Initial tensor should be sparse: {:.2}",
            initial_sparsity
        );

        // Further prune the already sparse tensor
        let config = UnstructuredPruningConfig::new(
            PruningCriterion::GlobalMagnitude,
            0.9, // Target 90% sparsity
        );

        let pruner = UnstructuredPruner::new(config)?;
        let further_pruned = pruner.prune_tensor(&sparse_tensor, "test")?;
        let final_sparsity = calculate_sparsity(&further_pruned)?;

        assert!(
            final_sparsity > initial_sparsity,
            "Further pruning should increase sparsity: {:.2} -> {:.2}",
            initial_sparsity,
            final_sparsity
        );

        assert!(
            (final_sparsity - 0.9).abs() < 0.05,
            "Should reach target sparsity: {:.2}",
            final_sparsity
        );

        Ok(())
    }

    #[test]
    fn test_pruning_statistics_collection() -> Result<()> {
        let config = UnstructuredPruningConfig::new_with_stats_collection(
            PruningCriterion::GlobalMagnitude,
            0.8,
            true, // Collect detailed statistics
        );

        let pruner = UnstructuredPruner::new(config)?;
        let model = create_test_model(42)?;

        let _pruned_model = pruner.prune_model(&model)?;
        let stats = pruner.get_statistics();

        assert!(stats.total_parameters > 0);
        assert!(stats.pruned_parameters > 0);
        assert!(stats.compression_ratio > 1.0);
        assert!(!stats.layer_statistics.is_empty());
        assert!(stats.memory_saved_mb > 0.0);

        // Check that statistics make sense
        let expected_remaining = stats.total_parameters - stats.pruned_parameters;
        assert_eq!(expected_remaining, stats.remaining_parameters);

        let expected_sparsity = stats.pruned_parameters as f32 / stats.total_parameters as f32;
        assert!((stats.overall_sparsity - expected_sparsity).abs() < 0.01);

        Ok(())
    }
}