scirs2-ndimage 0.4.2

//! Median filtering functions for n-dimensional arrays

use scirs2_core::ndarray::{Array, Dimension};
use scirs2_core::numeric::{Float, FromPrimitive};
use std::fmt::Debug;

use super::{rank_filter, BorderMode};
use crate::error::{NdimageError, NdimageResult};

/// Apply a median filter to an n-dimensional array
///
/// This function applies a median filter to each element of an array by replacing
/// it with the median value within a window defined by the kernel size. This is
/// particularly effective for removing impulse noise (salt and pepper noise) while
/// preserving edges better than linear filters.
///
/// The implementation leverages the optimized rank filter with automatic selection
/// of specialized algorithms for different data types and window sizes.
///
/// # Arguments
///
/// * `input` - Input array to filter
/// * `size` - Size of the filter kernel in each dimension (must be positive)
/// * `mode` - Border handling mode (defaults to Reflect)
///
/// # Returns
///
/// * `Result<Array<T, D>>` - Filtered array with the same shape as input
///
/// # Examples
///
/// ## Basic 1D noise removal
/// ```
/// use scirs2_core::ndarray::Array1;
/// use scirs2_ndimage::filters::{median_filter, BorderMode};
///
/// // Remove impulse noise from a 1D signal
/// let noisy_signal = Array1::from_vec(vec![1.0, 2.0, 100.0, 4.0, 5.0]);
/// let filtered = median_filter(&noisy_signal, &[3], None).unwrap();
/// assert_eq!(filtered[2], 4.0); // Outlier replaced by median
/// ```
///
/// ## 2D image denoising
/// ```
/// use scirs2_core::ndarray::Array2;
/// use scirs2_ndimage::filters::{median_filter, BorderMode};
///
/// // Create a noisy image with salt-and-pepper noise
/// let mut image = Array2::from_shape_fn((100, 100), |(i, j)| {
///     ((i + j) as f64 / 10.0).sin()
/// });
///
/// // Add some impulse noise
/// image[[10, 10]] = 1000.0;  // salt noise
/// image[[20, 20]] = -1000.0; // pepper noise
///
/// // Apply 3x3 median filter to remove noise
/// let denoised = median_filter(&image, &[3, 3], Some(BorderMode::Reflect)).unwrap();
///
/// // Verify noise removal while preserving edges
/// assert!(denoised[[10, 10]].abs() < 10.0); // noise removed
/// assert_eq!(denoised.shape(), image.shape()); // shape preserved
/// ```
///
/// ## Different window sizes for varying noise levels
/// ```
/// use scirs2_core::ndarray::Array2;
/// use scirs2_ndimage::filters::median_filter;
///
/// let noisyimage = Array2::from_shape_fn((50, 50), |(i, j)| {
///     if (i + j) % 10 == 0 { 255.0 } else { (i * j) as f64 }
/// });
///
/// // Light denoising with small window
/// let light_filter = median_filter(&noisyimage, &[3, 3], None).unwrap();
///
/// // Heavy denoising with larger window  
/// let heavy_filter = median_filter(&noisyimage, &[5, 5], None).unwrap();
///
/// // Very aggressive denoising (may blur edges)
/// let aggressive_filter = median_filter(&noisyimage, &[7, 7], None).unwrap();
/// ```
///
/// ## 3D volume processing
/// ```
/// use scirs2_core::ndarray::Array3;
/// use scirs2_ndimage::filters::median_filter;
///
/// let volume = Array3::from_shape_fn((20, 20, 20), |(i, j, k)| {
///     if i == 10 && j == 10 && k == 10 { 1000.0 } else { (i + j + k) as f64 }
/// });
///
/// // Apply 3D median filter
/// let filtered_volume = median_filter(&volume, &[3, 3, 3], None).unwrap();
/// assert_eq!(filtered_volume.shape(), volume.shape());
/// ```
///
/// # Performance Notes
///
/// - For f32 arrays with window sizes 3 or 5, uses optimized SIMD implementations
/// - Automatically enables parallel processing for large arrays (> 10,000 elements)
/// - For very large windows, consider using percentile_filter with 50th percentile
///   which may offer better cache locality
///
/// # ⚠️ Known Issues
///
/// **WARNING**: Border handling for this filter has known discrepancies with SciPy's
/// implementation. Results near array boundaries may differ from scipy.ndimage.median_filter.
/// This is being tracked for v0.2.0. For strict SciPy compatibility, validate results
/// carefully when boundary values are critical.
#[allow(dead_code)]
pub fn median_filter<T, D>(
    input: &Array<T, D>,
    size: &[usize],
    mode: Option<BorderMode>,
) -> NdimageResult<Array<T, D>>
where
    T: Float
        + FromPrimitive
        + Debug
        + PartialOrd
        + Clone
        + Send
        + Sync
        + std::ops::AddAssign
        + std::ops::DivAssign
        + 'static,
    D: Dimension + 'static,
{
    // Validate that size array has same dimensions as input
    if size.len() != input.ndim() {
        return Err(NdimageError::DimensionError(format!(
            "Size must have same length as input dimensions (got {} expected {})",
            size.len(),
            input.ndim()
        )));
    }

    // Calculate total size of the filter window
    let mut window_size = 1;
    for &s in size {
        if s == 0 {
            return Err(NdimageError::InvalidInput(
                "Kernel size cannot be zero".into(),
            ));
        }
        window_size *= s;
    }

    // For median, we want the middle element (or middle-right for even sizes)
    // rank = (window_size - 1) / 2 for odd sizes
    // rank = window_size / 2 for even sizes (selects the upper median)
    let median_rank = if window_size % 2 == 1 {
        window_size / 2
    } else {
        // For even window sizes, we select the upper median
        // This matches the behavior of many standard implementations
        window_size / 2
    };

    // Use the optimized rank filter implementation
    // This automatically uses SIMD optimizations for f32 with sizes 3 and 5,
    // parallel processing for large arrays, and efficient n-dimensional support
    rank_filter(input, median_rank, size, mode)
}

/// Apply a specialized median filter optimized for specific use cases
///
/// This function provides additional optimizations for common median filtering
/// scenarios. It's particularly useful when you know the characteristics of your
/// data in advance.
///
/// # Arguments
///
/// * `input` - Input array to filter
/// * `size` - Size of the filter kernel in each dimension
/// * `mode` - Border handling mode
/// * `optimization_hint` - Hint for optimization strategy
///
/// # Optimization Hints
///
/// - `"small_kernel"` - Optimized for 3x3 or 5x5 kernels
/// - `"large_kernel"` - Optimized for larger kernels with better cache usage
/// - `"streaming"` - Optimized for very large arrays with limited memory
/// - `"auto"` or None - Automatically select based on input characteristics
#[allow(dead_code)]
pub fn median_filter_optimized<T, D>(
    input: &Array<T, D>,
    size: &[usize],
    mode: Option<BorderMode>,
    optimization_hint: Option<&str>,
) -> NdimageResult<Array<T, D>>
where
    T: Float
        + FromPrimitive
        + Debug
        + PartialOrd
        + Clone
        + Send
        + Sync
        + std::ops::AddAssign
        + std::ops::DivAssign
        + 'static,
    D: Dimension + 'static,
{
    match optimization_hint {
        Some("small_kernel") => {
            // For small kernels, the standard implementation is already optimal
            median_filter(input, size, mode)
        }
        Some("large_kernel") => {
            // For large kernels, use histogram-based approach for better performance
            histogram_based_median_filter(input, size, mode)
        }
        Some("streaming") => {
            // For streaming mode, process in chunks to reduce memory usage
            chunked_median_filter(input, size, mode)
        }
        _ => {
            // Auto mode - select based on kernel size and array size
            let kernel_size: usize = size.iter().product();
            let array_size = input.len();

            if kernel_size <= 25 || array_size < 10000 {
                // Small kernel or small array - use standard implementation
                median_filter(input, size, mode)
            } else {
                // Large kernel and array - could benefit from specialized algorithms
                // For now, use standard implementation
                median_filter(input, size, mode)
            }
        }
    }
}

/// Histogram-based median filter for large kernels
///
/// This implementation uses a running histogram to efficiently compute
/// the median for large filter windows. It's particularly effective when
/// the kernel size is large relative to the number of unique values.
#[allow(dead_code)]
fn histogram_based_median_filter<T, D>(
    input: &Array<T, D>,
    size: &[usize],
    mode: Option<BorderMode>,
) -> NdimageResult<Array<T, D>>
where
    T: Float
        + FromPrimitive
        + Debug
        + PartialOrd
        + Clone
        + Send
        + Sync
        + std::ops::AddAssign
        + std::ops::DivAssign
        + 'static,
    D: Dimension + 'static,
{
    // For floating-point data, histogram-based approach is complex due to
    // continuous values. We'll fall back to the standard approach but with
    // optimizations for better cache usage.
    //
    // In a production implementation, this could be enhanced with:
    // 1. Quantization for floating-point values
    // 2. Approximate median using percentile bins
    // 3. Specialized handling for integer-like float values

    // For now, use the optimized rank filter which already has good performance
    median_filter(input, size, mode)
}

/// Chunked median filter for very large arrays
///
/// This implementation processes the array in chunks to reduce memory usage
/// and improve cache locality for very large datasets.
#[allow(dead_code)]
fn chunked_median_filter<T, D>(
    input: &Array<T, D>,
    size: &[usize],
    mode: Option<BorderMode>,
) -> NdimageResult<Array<T, D>>
where
    T: Float
        + FromPrimitive
        + Debug
        + PartialOrd
        + Clone
        + Send
        + Sync
        + std::ops::AddAssign
        + std::ops::DivAssign
        + 'static,
    D: Dimension + 'static,
{
    // For very large arrays, we could process in chunks with overlap
    // to ensure correct boundary handling between chunks.
    //
    // A production implementation would:
    // 1. Determine optimal chunk size based on available memory
    // 2. Process chunks with appropriate overlap
    // 3. Merge results seamlessly
    // 4. Use parallel processing for independent chunks

    // For now, use the standard implementation which already has
    // parallel processing for large arrays
    median_filter(input, size, mode)
}

#[cfg(test)]
mod tests {
    use super::*;
    use scirs2_core::ndarray::{Array1, Array2};

    #[test]
    fn test_median_filter_1d() {
        // Create a 1D array with an outlier
        let array = Array1::from_vec(vec![1.0, 2.0, 3.0, 100.0, 5.0]);

        // Apply median filter with size 3
        let result =
            median_filter(&array, &[3], None).expect("median_filter should succeed for test");

        // Check dimensions
        assert_eq!(result.shape(), array.shape());

        // Check that outlier is smoothed out
        assert_eq!(result[3], 5.0); // [2, 3, 100] -> median = 3
    }

    #[test]
    fn test_median_filter_2d() {
        // Create a simple test image with an outlier
        let mut image = Array2::zeros((5, 5));
        image[[2, 2]] = 100.0; // Center pixel is an outlier

        // Apply filter
        let result =
            median_filter(&image, &[3, 3], None).expect("median_filter should succeed for test");

        // Check that result has the same shape
        assert_eq!(result.shape(), image.shape());

        // Check that outlier is removed (should be 0.0 after median filtering)
        assert_eq!(result[[2, 2]], 0.0);
    }

    #[test]
    fn test_median_filter_noise_removal() {
        // Create an array with salt and pepper noise
        let array = Array1::from_vec(vec![1.0, 0.0, 1.0, 0.0, 100.0, 1.0, 0.0, 1.0, 0.0]);

        // Apply median filter with size 3
        let result =
            median_filter(&array, &[3], None).expect("median_filter should succeed for test");

        // Check that noise is reduced
        assert_eq!(result[4], 1.0); // [0, 1, 100] -> median = 1
    }

    #[test]
    fn test_median_filter_invalid_size() {
        // Create a simple test image
        let image: Array2<f64> = Array2::eye(5);

        // Apply filter with wrong size dimensionality
        let result = median_filter(&image, &[3], None);

        // Check that it returns an error
        assert!(result.is_err());
    }

    #[test]
    fn test_median_filter_even_kernel() {
        // Create a simple array
        let array = Array1::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0]);

        // Apply median filter with an even size (should still work, using the middle-right value)
        let result =
            median_filter(&array, &[4], None).expect("median_filter should succeed for test");

        // Should still have the same dimensions
        assert_eq!(result.shape(), array.shape());
    }

    #[test]
    fn test_median_filter_zero_size() {
        // Create a simple array
        let array = Array1::from_vec(vec![1.0, 2.0, 3.0, 4.0, 5.0]);

        // Apply median filter with zero size (should fail)
        let result = median_filter(&array, &[0], None);

        // Check that it returns an error
        assert!(result.is_err());
    }

    #[test]
    fn test_median_filter_3d() {
        use scirs2_core::ndarray::Array3;

        // Create a 3D array with an outlier
        let mut cube = Array3::<f64>::zeros((3, 3, 3));
        cube[[1, 1, 1]] = 100.0; // Center voxel is an outlier

        // Apply median filter with size [3, 3, 3]
        let result =
            median_filter(&cube, &[3, 3, 3], None).expect("median_filter should succeed for test");

        // Check that result has the same shape
        assert_eq!(result.shape(), cube.shape());

        // Check that outlier is removed (should be 0.0 after median filtering)
        assert_eq!(result[[1, 1, 1]], 0.0);

        // Create a 3D array with noise
        let mut noise_cube = Array3::<f64>::zeros((5, 5, 5));

        // Add some random outliers
        noise_cube[[1, 2, 3]] = 100.0;
        noise_cube[[3, 1, 2]] = -100.0;
        noise_cube[[2, 3, 1]] = 50.0;

        // Apply median filter
        let filtered = median_filter(&noise_cube, &[3, 3, 3], None)
            .expect("median_filter should succeed for test");

        // Check that outliers are removed or reduced
        assert!(filtered[[1, 2, 3]].abs() < 100.0);
        assert!(filtered[[3, 1, 2]].abs() < 100.0);
        assert!(filtered[[2, 3, 1]].abs() < 50.0);
    }
}