#!/usr/bin/env rust-script

//! Enhanced test for magnitude pruning with structured patterns and schedules
//! Tests the full implementation including channel, filter, and N:M pruning

use std::collections::HashMap;

/// Result type for this standalone test
type Result<T> = std::result::Result<T, Box<dyn std::error::Error>>;

/// Simple error type for testing
#[derive(Debug)]
struct TestError(String);

impl std::fmt::Display for TestError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Test error: {}", self.0)
    }
}

impl std::error::Error for TestError {}

// Include all the implementations (expanded from the previous test)

/// Configuration for magnitude-based pruning
#[derive(Debug, Clone, PartialEq)]
pub struct PruningConfig {
    /// Pruning strategy to apply
    pub strategy: PruningStrategy,
    /// Target sparsity ratio (0.0 to 1.0)
    pub sparsity_ratio: f64,
    /// Pruning granularity
    pub granularity: PruningGranularity,
    /// Optional pruning schedule for gradual pruning
    pub schedule: Option<PruningSchedule>,
    /// Random seed for deterministic pruning
    pub seed: Option<u64>,
}

/// Strategy for selecting weights to prune
#[derive(Debug, Clone, PartialEq)]
pub enum PruningStrategy {
    /// Global magnitude pruning across all parameters
    Global,
    /// Layer-wise magnitude pruning within each layer
    LayerWise,
    /// Structured pruning (channels, filters)
    Structured {
        /// Type of structured pattern
        pattern_type: StructuredPattern,
    },
}

/// Granularity of pruning operations
#[derive(Debug, Clone, PartialEq)]
pub enum PruningGranularity {
    /// Element-wise pruning (unstructured)
    Element,
    /// Channel-wise pruning
    Channel,
    /// Filter-wise pruning
    Filter,
    /// Block-wise pruning
    Block { height: usize, width: usize },
}

/// Structured pruning patterns
#[derive(Debug, Clone, PartialEq)]
pub enum StructuredPattern {
    /// Channel pruning
    Channel,
    /// Filter pruning
    Filter,
    /// N:M sparsity pattern
    NM { n: usize, m: usize },
}

/// Pruning schedule for gradual sparsification during training
#[derive(Debug, Clone, PartialEq)]
pub enum PruningSchedule {
    /// Linear schedule from initial to final sparsity
    Linear {
        /// Initial sparsity ratio
        initial_sparsity: f64,
        /// Final sparsity ratio  
        final_sparsity: f64,
        /// Duration in training steps
        duration_steps: usize,
    },
    /// Polynomial schedule
    Polynomial {
        /// Initial sparsity ratio
        initial_sparsity: f64,
        /// Final sparsity ratio
        final_sparsity: f64,
        /// Duration in training steps
        duration_steps: usize,
        /// Polynomial exponent
        exponent: f64,
    },
    /// Exponential schedule
    Exponential {
        /// Initial sparsity ratio
        initial_sparsity: f64,
        /// Final sparsity ratio
        final_sparsity: f64,
        /// Duration in training steps
        duration_steps: usize,
        /// Decay rate
        decay_rate: f64,
    },
}

/// Binary mask for indicating which weights to keep/prune
#[derive(Debug, Clone)]
pub struct PruningMask {
    /// Binary mask values (true = keep, false = prune)
    mask: Vec<bool>,
    /// Shape of the mask
    shape: Vec<usize>,
    /// Number of parameters kept
    kept_parameters: usize,
    /// Total number of parameters
    total_parameters: usize,
}

/// Statistics about pruning results
#[derive(Debug, Clone, PartialEq)]
pub struct PruningStats {
    /// Original parameter count
    pub original_params: usize,
    /// Parameters after pruning
    pub pruned_params: usize,
    /// Achieved sparsity ratio
    pub sparsity_ratio: f64,
    /// Compression ratio
    pub compression_ratio: f64,
    /// Memory savings (in bytes, estimated)
    pub memory_savings: usize,
}

/// Main struct for magnitude-based pruning operations
#[derive(Debug)]
pub struct MagnitudePruning;

/// Simple tensor for testing
#[derive(Debug, Clone)]
pub struct DenseTensor {
    pub data: Vec<f32>,
    pub shape: Vec<usize>,
}

impl DenseTensor {
    pub fn new(data: Vec<f32>, shape: Vec<usize>) -> Self {
        assert_eq!(data.len(), shape.iter().product::<usize>());
        Self { data, shape }
    }
    
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
    
    pub fn data(&self) -> &[f32] {
        &self.data
    }
}

impl PruningConfig {
    /// Create global magnitude pruning configuration
    pub fn global(sparsity_ratio: f64) -> Self {
        Self {
            strategy: PruningStrategy::Global,
            sparsity_ratio,
            granularity: PruningGranularity::Element,
            schedule: None,
            seed: None,
        }
    }
    
    /// Create layer-wise magnitude pruning configuration
    pub fn layerwise(sparsity_ratio: f64) -> Self {
        Self {
            strategy: PruningStrategy::LayerWise,
            sparsity_ratio,
            granularity: PruningGranularity::Element,
            schedule: None,
            seed: None,
        }
    }
    
    /// Create structured magnitude pruning configuration
    pub fn structured(sparsity_ratio: f64, pattern: StructuredPattern) -> Self {
        let granularity = match pattern {
            StructuredPattern::Channel => PruningGranularity::Channel,
            StructuredPattern::Filter => PruningGranularity::Filter,
            StructuredPattern::NM { .. } => PruningGranularity::Element,
        };
        
        Self {
            strategy: PruningStrategy::Structured {
                pattern_type: pattern,
            },
            sparsity_ratio,
            granularity,
            schedule: None,
            seed: None,
        }
    }
    
    /// Add pruning schedule
    pub fn with_schedule(mut self, schedule: PruningSchedule) -> Self {
        self.schedule = Some(schedule);
        self
    }
    
    /// Validate configuration parameters
    pub fn validate(&self) -> Result<()> {
        if !(0.0..=1.0).contains(&self.sparsity_ratio) {
            return Err(Box::new(TestError("Sparsity ratio must be between 0.0 and 1.0".into())));
        }
        Ok(())
    }
}

impl PruningSchedule {
    /// Create linear pruning schedule
    pub fn linear(duration_steps: usize, final_sparsity: f64) -> Self {
        Self::Linear {
            initial_sparsity: 0.0,
            final_sparsity,
            duration_steps,
        }
    }
    
    /// Create polynomial pruning schedule
    pub fn polynomial(duration_steps: usize, exponent: f64) -> Self {
        Self::Polynomial {
            initial_sparsity: 0.0,
            final_sparsity: 0.9, // Default to 90% sparsity
            duration_steps,
            exponent,
        }
    }
    
    /// Create exponential pruning schedule
    pub fn exponential(duration_steps: usize, decay_rate: f64) -> Self {
        Self::Exponential {
            initial_sparsity: 0.0,
            final_sparsity: 0.9, // Default to 90% sparsity
            duration_steps,
            decay_rate,
        }
    }
    
    /// Compute sparsity ratio at given training step
    pub fn compute_sparsity(&self, current_step: usize) -> f64 {
        match self {
            Self::Linear { initial_sparsity, final_sparsity, duration_steps } => {
                if current_step >= *duration_steps {
                    return *final_sparsity;
                }
                let progress = current_step as f64 / *duration_steps as f64;
                initial_sparsity + progress * (final_sparsity - initial_sparsity)
            }
            Self::Polynomial { initial_sparsity, final_sparsity, duration_steps, exponent } => {
                if current_step >= *duration_steps {
                    return *final_sparsity;
                }
                let progress = current_step as f64 / *duration_steps as f64;
                let poly_progress = progress.powf(*exponent);
                initial_sparsity + poly_progress * (final_sparsity - initial_sparsity)
            }
            Self::Exponential { initial_sparsity, final_sparsity, duration_steps, decay_rate } => {
                if current_step >= *duration_steps {
                    return *final_sparsity;
                }
                let progress = current_step as f64 / *duration_steps as f64;
                let exp_progress = 1.0 - (-decay_rate * progress).exp();
                initial_sparsity + exp_progress * (final_sparsity - initial_sparsity)
            }
        }
    }
}

impl PruningMask {
    /// Create new pruning mask
    pub fn new(mask: Vec<bool>, shape: Vec<usize>) -> Self {
        let total_parameters = shape.iter().product::<usize>();
        let kept_parameters = mask.iter().filter(|&&x| x).count();
        
        Self {
            mask,
            shape,
            kept_parameters,
            total_parameters,
        }
    }
    
    /// Get mask values
    pub fn mask(&self) -> &[bool] {
        &self.mask
    }
    
    /// Get mask shape
    pub fn shape(&self) -> &[usize] {
        &self.shape
    }
    
    /// Get sparsity ratio
    pub fn sparsity_ratio(&self) -> f64 {
        1.0 - (self.kept_parameters as f64 / self.total_parameters as f64)
    }
    
    /// Apply mask to tensor (element-wise multiplication)
    pub fn apply(&self, tensor: &DenseTensor) -> Result<DenseTensor> {
        if tensor.shape != self.shape {
            return Err(Box::new(TestError("Tensor and mask shapes must match".into())));
        }
        
        let masked_data: Vec<f32> = tensor.data.iter()
            .zip(self.mask.iter())
            .map(|(&value, &keep)| if keep { value } else { 0.0 })
            .collect();
        
        Ok(DenseTensor::new(masked_data, tensor.shape.clone()))
    }
}

impl PruningStats {
    /// Create new pruning statistics
    pub fn new(original_params: usize, pruned_params: usize) -> Self {
        let sparsity_ratio = 1.0 - (pruned_params as f64 / original_params as f64);
        let compression_ratio = if pruned_params == 0 {
            f64::INFINITY
        } else {
            original_params as f64 / pruned_params as f64
        };
        
        // Estimate memory savings (assuming 4 bytes per f32)
        let memory_savings = (original_params - pruned_params) * 4;
        
        Self {
            original_params,
            pruned_params,
            sparsity_ratio,
            compression_ratio,
            memory_savings,
        }
    }
}

impl MagnitudePruning {
    /// Apply magnitude-based pruning to tensor
    pub fn prune(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
        config.validate()?;
        
        match config.strategy {
            PruningStrategy::Global => {
                Self::prune_global(tensor, config)
            }
            PruningStrategy::LayerWise => {
                Self::prune_layerwise(tensor, config)
            }
            PruningStrategy::Structured { ref pattern_type } => {
                Self::prune_structured(tensor, config, pattern_type)
            }
        }
    }
    
    /// Create pruning mask without modifying the original tensor
    pub fn create_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        config.validate()?;
        
        match config.strategy {
            PruningStrategy::Global => {
                Self::create_global_mask(tensor, config)
            }
            PruningStrategy::LayerWise => {
                Self::create_layerwise_mask(tensor, config)
            }
            PruningStrategy::Structured { ref pattern_type } => {
                Self::create_structured_mask(tensor, config, pattern_type)
            }
        }
    }
    
    /// Compute importance scores based on magnitude
    pub fn compute_importance_scores(tensor: &DenseTensor) -> Result<Vec<f32>> {
        let scores: Vec<f32> = tensor.data.iter().map(|&x| x.abs()).collect();
        Ok(scores)
    }
    
    /// Analyze pruning results
    pub fn analyze_pruning(original: &DenseTensor, pruned: &DenseTensor) -> Result<PruningStats> {
        if original.shape != pruned.shape {
            return Err(Box::new(TestError("Tensor shapes must match for analysis".into())));
        }
        
        let original_nonzeros = original.data.iter().filter(|&&x| x != 0.0).count();
        let pruned_nonzeros = pruned.data.iter().filter(|&&x| x != 0.0).count();
        
        Ok(PruningStats::new(original_nonzeros, pruned_nonzeros))
    }
    
    fn prune_global(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
        let mask = Self::create_global_mask(tensor, config)?;
        mask.apply(tensor)
    }
    
    fn prune_layerwise(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
        let mask = Self::create_layerwise_mask(tensor, config)?;
        mask.apply(tensor)
    }
    
    fn prune_structured(
        tensor: &DenseTensor,
        config: &PruningConfig,
        pattern: &StructuredPattern,
    ) -> Result<DenseTensor> {
        let mask = Self::create_structured_mask(tensor, config, pattern)?;
        mask.apply(tensor)
    }
    
    fn create_global_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        let total_elements = tensor.data.len();
        let elements_to_keep = ((1.0 - config.sparsity_ratio) * total_elements as f64).round() as usize;
        
        // Compute importance scores (absolute magnitudes)
        let importance_scores = Self::compute_importance_scores(tensor)?;
        
        // Create (index, score) pairs for sorting
        let mut indexed_scores: Vec<(usize, f32)> = importance_scores.iter()
            .enumerate()
            .map(|(i, &score)| (i, score))
            .collect();
        
        // Sort by importance score in descending order (keep highest magnitude values)
        indexed_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
        
        // Create mask: true for kept elements, false for pruned elements
        let mut mask_data = vec![false; total_elements];
        for i in 0..elements_to_keep {
            if i < indexed_scores.len() {
                let (index, _) = indexed_scores[i];
                mask_data[index] = true;
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
    
    fn create_layerwise_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        // For layer-wise pruning, we currently support 2D tensors (treating each row as a layer)
        if tensor.shape.len() != 2 {
            return Err(Box::new(TestError("Layer-wise pruning currently only supports 2D tensors".into())));
        }
        
        let rows = tensor.shape[0];
        let cols = tensor.shape[1];
        let elements_per_row = cols;
        let elements_to_keep_per_row = ((1.0 - config.sparsity_ratio) * elements_per_row as f64).round() as usize;
        
        let mut mask_data = vec![false; tensor.data.len()];
        
        // Process each row independently
        for row in 0..rows {
            let row_start = row * cols;
            let row_end = row_start + cols;
            
            // Get magnitudes for this row
            let row_scores: Vec<(usize, f32)> = (row_start..row_end)
                .map(|i| (i, tensor.data[i].abs()))
                .collect();
            
            // Sort by magnitude in descending order
            let mut sorted_scores = row_scores;
            sorted_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
            
            // Keep the top elements_to_keep_per_row elements in this row
            for i in 0..elements_to_keep_per_row.min(sorted_scores.len()) {
                let (original_index, _) = sorted_scores[i];
                mask_data[original_index] = true;
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
    
    fn create_structured_mask(
        tensor: &DenseTensor,
        config: &PruningConfig,
        pattern: &StructuredPattern,
    ) -> Result<PruningMask> {
        match pattern {
            StructuredPattern::Channel => {
                Self::create_channel_pruning_mask(tensor, config)
            }
            StructuredPattern::Filter => {
                Self::create_filter_pruning_mask(tensor, config)
            }
            StructuredPattern::NM { n, m } => {
                Self::create_nm_pruning_mask(tensor, config, *n, *m)
            }
        }
    }
    
    fn create_channel_pruning_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        // Channel pruning for 2D tensors (each row is a channel)
        if tensor.shape.len() != 2 {
            return Err(Box::new(TestError("Channel pruning currently only supports 2D tensors".into())));
        }
        
        let rows = tensor.shape[0];
        let cols = tensor.shape[1];
        let channels_to_keep = ((1.0 - config.sparsity_ratio) * rows as f64).round() as usize;
        
        // Calculate channel importance (sum of absolute values per row)
        let mut channel_scores: Vec<(usize, f32)> = Vec::new();
        
        for row in 0..rows {
            let row_start = row * cols;
            let row_end = row_start + cols;
            let row_sum: f32 = tensor.data[row_start..row_end]
                .iter()
                .map(|&x| x.abs())
                .sum();
            channel_scores.push((row, row_sum));
        }
        
        // Sort by importance in descending order
        channel_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
        
        // Create mask: keep top channels, zero out others
        let mut mask_data = vec![false; tensor.data.len()];
        for i in 0..channels_to_keep {
            if i < channel_scores.len() {
                let (channel_idx, _) = channel_scores[i];
                let row_start = channel_idx * cols;
                let row_end = row_start + cols;
                for j in row_start..row_end {
                    mask_data[j] = true;
                }
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
    
    fn create_filter_pruning_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
        // Filter pruning for 2D tensors (each column is a filter)
        if tensor.shape.len() != 2 {
            return Err(Box::new(TestError("Filter pruning currently only supports 2D tensors".into())));
        }
        
        let rows = tensor.shape[0];
        let cols = tensor.shape[1];
        let filters_to_keep = ((1.0 - config.sparsity_ratio) * cols as f64).round() as usize;
        
        // Calculate filter importance (sum of absolute values per column)
        let mut filter_scores: Vec<(usize, f32)> = Vec::new();
        
        for col in 0..cols {
            let col_sum: f32 = (0..rows)
                .map(|row| tensor.data[row * cols + col].abs())
                .sum();
            filter_scores.push((col, col_sum));
        }
        
        // Sort by importance in descending order
        filter_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
        
        // Create mask: keep top filters, zero out others
        let mut mask_data = vec![false; tensor.data.len()];
        for i in 0..filters_to_keep {
            if i < filter_scores.len() {
                let (filter_idx, _) = filter_scores[i];
                for row in 0..rows {
                    mask_data[row * cols + filter_idx] = true;
                }
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
    
    fn create_nm_pruning_mask(tensor: &DenseTensor, config: &PruningConfig, n: usize, m: usize) -> Result<PruningMask> {
        // N:M sparsity: keep N elements out of every M consecutive elements
        if n == 0 || m == 0 {
            return Err(Box::new(TestError("N and M must be non-zero for N:M sparsity".into())));
        }
        if n > m {
            return Err(Box::new(TestError("N cannot be greater than M in N:M sparsity".into())));
        }
        
        let total_elements = tensor.data.len();
        let mut mask_data = vec![false; total_elements];
        
        // Process in groups of M elements
        for group_start in (0..total_elements).step_by(m) {
            let group_end = (group_start + m).min(total_elements);
            let group_size = group_end - group_start;
            let elements_to_keep = n.min(group_size);
            
            // Get magnitudes for this group
            let mut group_scores: Vec<(usize, f32)> = (group_start..group_end)
                .map(|i| (i, tensor.data[i].abs()))
                .collect();
            
            // Sort by magnitude in descending order
            group_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
            
            // Keep top N elements in this group
            for i in 0..elements_to_keep {
                let (original_index, _) = group_scores[i];
                mask_data[original_index] = true;
            }
        }
        
        Ok(PruningMask::new(mask_data, tensor.shape.clone()))
    }
}

/// Test utilities
mod test_utils {
    use super::*;
    
    /// Create a simple 2D test tensor with known values
    pub fn create_test_tensor_2d() -> DenseTensor {
        // 3x4 tensor with specific values for predictable pruning results
        let data = vec![
            0.1, 0.9, 0.2, 0.8,   // Row 0: sum = 2.0
            0.3, 0.7, 0.4, 0.6,   // Row 1: sum = 2.0
            0.05, 0.95, 0.15, 0.85 // Row 2: sum = 2.0
        ];
        DenseTensor::new(data, vec![3, 4])
    }
    
    /// Create tensor with different channel magnitudes
    pub fn create_channel_test_tensor() -> DenseTensor {
        let data = vec![
            0.1, 0.1, 0.1, 0.1,   // Row 0: sum = 0.4 (weakest channel)
            0.5, 0.5, 0.5, 0.5,   // Row 1: sum = 2.0 (strongest channel)  
            0.2, 0.2, 0.2, 0.2    // Row 2: sum = 0.8 (middle channel)
        ];
        DenseTensor::new(data, vec![3, 4])
    }
    
    /// Create tensor with different filter magnitudes
    pub fn create_filter_test_tensor() -> DenseTensor {
        let data = vec![
            0.1, 0.5, 0.2, 0.05,  // Col sums: 0.6, 1.5, 0.8, 0.35
            0.2, 0.6, 0.3, 0.15, 
            0.3, 0.4, 0.3, 0.15   
        ];
        DenseTensor::new(data, vec![3, 4])
    }
}

fn main() -> Result<()> {
    println!("Running enhanced magnitude pruning tests...");
    
    // Test 1: Channel Pruning
    println!("\n=== Test 1: Channel Pruning ===");
    let tensor = test_utils::create_channel_test_tensor();
    println!("Channel test tensor: {:?}", tensor);
    
    // Channel-wise magnitude sums: Row0=0.4, Row1=2.0, Row2=0.8
    // With 33% sparsity, should keep 2 out of 3 channels (keep strongest 2)
    let config = PruningConfig::structured(0.33, StructuredPattern::Channel);
    let pruned = MagnitudePruning::prune(&tensor, &config)?;
    println!("Channel pruned result: {:?}", pruned.data);
    
    // Should keep Row 1 (sum=2.0) and Row 2 (sum=0.8), prune Row 0 (sum=0.4)
    let expected_data = vec![
        0.0, 0.0, 0.0, 0.0,   // Row 0 pruned
        0.5, 0.5, 0.5, 0.5,   // Row 1 kept
        0.2, 0.2, 0.2, 0.2    // Row 2 kept
    ];
    assert_eq!(pruned.data, expected_data);
    println!("✓ Channel pruning passed");
    
    // Test 2: Filter Pruning
    println!("\n=== Test 2: Filter Pruning ===");
    let tensor = test_utils::create_filter_test_tensor();
    println!("Filter test tensor: {:?}", tensor);
    
    // Column-wise magnitude sums: Col0=0.6, Col1=1.5, Col2=0.8, Col3=0.35
    // With 50% sparsity, should keep 2 out of 4 filters (keep strongest 2: Col1, Col2)
    let config = PruningConfig::structured(0.5, StructuredPattern::Filter);
    let pruned = MagnitudePruning::prune(&tensor, &config)?;
    println!("Filter pruned result: {:?}", pruned.data);
    
    // Should keep Col 1 (sum=1.5) and Col 2 (sum=0.8), prune Col 0 and Col 3
    let expected_data = vec![
        0.0, 0.5, 0.2, 0.0,   
        0.0, 0.6, 0.3, 0.0, 
        0.0, 0.4, 0.3, 0.0   
    ];
    assert_eq!(pruned.data, expected_data);
    println!("✓ Filter pruning passed");
    
    // Test 3: N:M Sparsity (2:4 pattern)
    println!("\n=== Test 3: N:M Sparsity (2:4) ===");
    let tensor = DenseTensor::new(vec![0.1, 0.9, 0.2, 0.8, 0.3, 0.7, 0.4, 0.6], vec![8]);
    println!("N:M test tensor: {:?}", tensor.data);
    
    let config = PruningConfig::structured(0.0, StructuredPattern::NM { n: 2, m: 4 }); // Keep 2 out of every 4
    let pruned = MagnitudePruning::prune(&tensor, &config)?;
    println!("2:4 pruned result: {:?}", pruned.data);
    
    // Group 1 [0.1, 0.9, 0.2, 0.8]: keep top 2 [0.9, 0.8]
    // Group 2 [0.3, 0.7, 0.4, 0.6]: keep top 2 [0.7, 0.6]
    let expected_data = vec![0.0, 0.9, 0.0, 0.8, 0.0, 0.7, 0.0, 0.6];
    assert_eq!(pruned.data, expected_data);
    println!("✓ N:M sparsity (2:4) passed");
    
    // Test 4: Pruning Schedules
    println!("\n=== Test 4: Pruning Schedules ===");
    
    // Test linear schedule
    let linear_schedule = PruningSchedule::linear(100, 0.8);
    let sparsity_0 = linear_schedule.compute_sparsity(0);
    let sparsity_50 = linear_schedule.compute_sparsity(50);
    let sparsity_100 = linear_schedule.compute_sparsity(100);
    let sparsity_150 = linear_schedule.compute_sparsity(150);
    
    println!("Linear schedule sparsity at steps [0, 50, 100, 150]: [{:.2}, {:.2}, {:.2}, {:.2}]", 
             sparsity_0, sparsity_50, sparsity_100, sparsity_150);
    assert!((sparsity_0 - 0.0).abs() < f64::EPSILON);
    assert!((sparsity_50 - 0.4).abs() < f64::EPSILON);
    assert!((sparsity_100 - 0.8).abs() < f64::EPSILON);
    assert!((sparsity_150 - 0.8).abs() < f64::EPSILON);
    
    // Test polynomial schedule  
    let poly_schedule = PruningSchedule::polynomial(100, 2.0);
    let poly_sparsity_50 = poly_schedule.compute_sparsity(50);
    println!("Polynomial schedule sparsity at step 50: {:.3}", poly_sparsity_50);
    assert!(poly_sparsity_50 < 0.5); // Should be less than linear at midpoint
    
    // Test exponential schedule
    let exp_schedule = PruningSchedule::exponential(100, 2.0);
    let exp_sparsity_50 = exp_schedule.compute_sparsity(50);
    println!("Exponential schedule sparsity at step 50: {:.3}", exp_sparsity_50);
    assert!(exp_sparsity_50 > 0.3 && exp_sparsity_50 < 0.7); // Should be in reasonable range
    
    println!("✓ Pruning schedules passed");
    
    // Test 5: Advanced Statistics
    println!("\n=== Test 5: Advanced Statistics ===");
    let original = test_utils::create_test_tensor_2d();
    let config = PruningConfig::global(0.75); // 75% sparsity
    let pruned = MagnitudePruning::prune(&original, &config)?;
    
    let stats = MagnitudePruning::analyze_pruning(&original, &pruned)?;
    println!("Advanced stats: {:?}", stats);
    
    assert_eq!(stats.original_params, 12);
    assert_eq!(stats.pruned_params, 3); // Keep 25% = 3 elements
    assert!((stats.sparsity_ratio - 0.75).abs() < 0.01);
    assert!((stats.compression_ratio - 4.0).abs() < 0.1);
    
    println!("✓ Advanced statistics passed");
    
    println!("\n🎉 All enhanced tests passed! Complete magnitude pruning implementation is working correctly.");
    
    Ok(())
}