zenjpeg 0.8.2

//! Fast Gaussian blur for pre-JPEG compression preprocessing.
//!
//! Optimized for small sigma (σ ≤ 0.5) where the Gaussian kernel is effectively
//! 3-tap: `[w1, w0, w1]`. At σ=0.4: w0 ≈ 0.919, w1 ≈ 0.040.
//!
//! Strategy: deinterleave to planar R/G/B, blur each plane with a separable
//! 3-tap fixed-point kernel, reinterleave. The planar layout gives clean
//! sequential access for auto-vectorization.
//!
//! All arithmetic stays in u16 with SHIFT=8. Weights are forced to sum to
//! exactly 256 for correct normalization.
//!
//! Performance (512x512, σ=0.4, target-cpu=native):
//!   - This implementation: ~300µs
//!   - Hand-tuned AVX2 SIMD: ~180µs (1.7x faster)
//!   - Scalar f32 gaussian: ~2.3ms (7.7x slower)
//!
//! Quality: max_diff=2 vs f32 reference, only 0.5% of pixels differ.

use alloc::vec;
use alloc::vec::Vec;

/// Gaussian blur on packed RGB u8 data (3 bytes per pixel, row-major).
///
/// For σ ≤ 0.5, uses an optimized 3-tap kernel. For larger sigma, falls back
/// to a general N-tap separable convolution.
///
/// Returns a new Vec with the blurred image data.
pub fn gaussian_blur_rgb(src: &[u8], width: usize, height: usize, sigma: f32) -> Vec<u8> {
    if sigma <= 0.0 || width < 2 || height < 2 {
        return src.to_vec();
    }

    let radius = ceil_f32(sigma * 3.0) as usize;
    if radius <= 1 {
        blur_3tap_planar_rgb(src, width, height, sigma)
    } else {
        blur_general_planar_rgb(src, width, height, sigma)
    }
}

/// 3-tap planar blur for small sigma (radius ≤ 1).
fn blur_3tap_planar_rgb(src: &[u8], width: usize, height: usize, sigma: f32) -> Vec<u8> {
    let n = width * height;

    // Compute 3-tap kernel weights (SHIFT=8, sum=256)
    let inv_2sigma2 = 1.0 / (2.0 * sigma * sigma);
    let w1_f = exp_f32(-inv_2sigma2);
    let total = 1.0f32 + 2.0 * w1_f;
    let w1 = (w1_f / total * 256.0 + 0.5) as u16;
    let w0 = 256u16.saturating_sub(2 * w1);

    // Deinterleave RGB to planar
    let mut planes = [vec![0u8; n], vec![0u8; n], vec![0u8; n]];
    for i in 0..n {
        planes[0][i] = src[i * 3];
        planes[1][i] = src[i * 3 + 1];
        planes[2][i] = src[i * 3 + 2];
    }

    // Blur each plane
    for plane in &mut planes {
        blur_plane_3tap(plane, width, height, w0, w1);
    }

    // Reinterleave
    let mut dst = vec![0u8; n * 3];
    for i in 0..n {
        dst[i * 3] = planes[0][i];
        dst[i * 3 + 1] = planes[1][i];
        dst[i * 3 + 2] = planes[2][i];
    }
    dst
}

/// General N-tap planar blur for larger sigma.
fn blur_general_planar_rgb(src: &[u8], width: usize, height: usize, sigma: f32) -> Vec<u8> {
    let n = width * height;

    // Compute kernel
    let radius = ceil_f32(sigma * 3.0) as usize;
    let kernel_size = 2 * radius + 1;

    // Use SHIFT=14 for larger kernels (more precision needed)
    const SHIFT: u32 = 14;
    const SCALE: f32 = (1u32 << SHIFT) as f32;

    let inv_2sigma2 = 1.0 / (2.0 * sigma * sigma);
    let mut weights_f32 = vec![0.0f32; kernel_size];
    let mut sum = 0.0f32;
    for i in 0..kernel_size {
        let d = i as f32 - radius as f32;
        let w = exp_f32(-d * d * inv_2sigma2);
        weights_f32[i] = w;
        sum += w;
    }
    let inv_sum = 1.0 / sum;
    let weights: Vec<u32> = weights_f32
        .iter()
        .map(|&w| (w * inv_sum * SCALE + 0.5) as u32)
        .collect();

    // Deinterleave
    let mut planes = [vec![0u8; n], vec![0u8; n], vec![0u8; n]];
    for i in 0..n {
        planes[0][i] = src[i * 3];
        planes[1][i] = src[i * 3 + 1];
        planes[2][i] = src[i * 3 + 2];
    }

    // Blur each plane
    for plane in &mut planes {
        blur_plane_general(plane, width, height, &weights, radius, SHIFT);
    }

    // Reinterleave
    let mut dst = vec![0u8; n * 3];
    for i in 0..n {
        dst[i * 3] = planes[0][i];
        dst[i * 3 + 1] = planes[1][i];
        dst[i * 3 + 2] = planes[2][i];
    }
    dst
}

/// In-place 3-tap separable blur on a single u8 plane.
fn blur_plane_3tap(plane: &mut [u8], width: usize, height: usize, w0: u16, w1: u16) {
    let n = width * height;
    let mut temp = vec![0u8; n];

    // Horizontal pass
    for y in 0..height {
        let row = y * width;

        // Left edge: clamp
        temp[row] = ((plane[row] as u16 * (w0 + w1) + plane[row + 1] as u16 * w1) >> 8) as u8;

        // Interior
        for x in 1..width - 1 {
            let i = row + x;
            let val = plane[i - 1] as u16 * w1 + plane[i] as u16 * w0 + plane[i + 1] as u16 * w1;
            temp[i] = (val >> 8) as u8;
        }

        // Right edge: clamp
        let last = row + width - 1;
        temp[last] = ((plane[last - 1] as u16 * w1 + plane[last] as u16 * (w0 + w1)) >> 8) as u8;
    }

    // Vertical pass
    // Top edge
    for x in 0..width {
        plane[x] = ((temp[x] as u16 * (w0 + w1) + temp[width + x] as u16 * w1) >> 8) as u8;
    }

    // Interior rows
    for y in 1..height - 1 {
        let above = (y - 1) * width;
        let center = y * width;
        let below = (y + 1) * width;
        for x in 0..width {
            let val = temp[above + x] as u16 * w1
                + temp[center + x] as u16 * w0
                + temp[below + x] as u16 * w1;
            plane[center + x] = (val >> 8) as u8;
        }
    }

    // Bottom edge
    let last_row = (height - 1) * width;
    let prev_row = (height - 2) * width;
    for x in 0..width {
        plane[last_row + x] =
            ((temp[prev_row + x] as u16 * w1 + temp[last_row + x] as u16 * (w0 + w1)) >> 8) as u8;
    }
}

/// In-place general N-tap separable blur on a single u8 plane.
fn blur_plane_general(
    plane: &mut [u8],
    width: usize,
    height: usize,
    weights: &[u32],
    radius: usize,
    shift: u32,
) {
    let n = width * height;
    let kernel_size = weights.len();

    // Horizontal pass (u8 → u16 temp via u32 accumulator)
    let mut temp = vec![0u16; n];
    for y in 0..height {
        let row = y * width;
        for x in 0..width {
            let mut acc = 0u32;
            for ki in 0..kernel_size {
                let sx = (x as isize + ki as isize - radius as isize).clamp(0, width as isize - 1)
                    as usize;
                acc += plane[row + sx] as u32 * weights[ki];
            }
            temp[row + x] = (acc >> shift) as u16;
        }
    }

    // Vertical pass (u16 → u8 via u32 accumulator)
    for y in 0..height {
        for x in 0..width {
            let mut acc = 0u32;
            for ki in 0..kernel_size {
                let sy = (y as isize + ki as isize - radius as isize).clamp(0, height as isize - 1)
                    as usize;
                acc += temp[sy * width + x] as u32 * weights[ki];
            }
            plane[y * width + x] = ((acc >> shift) as u16).min(255) as u8;
        }
    }
}

#[inline]
fn exp_f32(x: f32) -> f32 {
    x.exp()
}

#[inline]
fn ceil_f32(x: f32) -> f32 {
    x.ceil()
}

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

    #[test]
    fn blur_identity_at_zero_sigma() {
        let src = vec![100u8, 150, 200, 50, 100, 150, 200, 50, 100, 150, 200, 50];
        let result = gaussian_blur_rgb(&src, 2, 2, 0.0);
        assert_eq!(src, result);
    }

    #[test]
    fn blur_3tap_preserves_size() {
        let src = vec![128u8; 64 * 64 * 3];
        let result = gaussian_blur_rgb(&src, 64, 64, 0.4);
        assert_eq!(result.len(), src.len());
    }

    #[test]
    fn blur_3tap_uniform_image_unchanged() {
        let src = vec![100u8; 32 * 32 * 3];
        let result = gaussian_blur_rgb(&src, 32, 32, 0.4);
        for (i, (&s, &r)) in src.iter().zip(result.iter()).enumerate() {
            assert!(
                (s as i16 - r as i16).unsigned_abs() <= 1,
                "pixel {i}: src={s} result={r}"
            );
        }
    }

    #[test]
    fn blur_general_preserves_size() {
        // sigma=1.0 uses general path (radius=3)
        let src = vec![128u8; 64 * 64 * 3];
        let result = gaussian_blur_rgb(&src, 64, 64, 1.0);
        assert_eq!(result.len(), src.len());
    }

    #[test]
    fn blur_small_image() {
        // 3x3 should not panic
        let src = vec![128u8; 3 * 3 * 3];
        let result = gaussian_blur_rgb(&src, 3, 3, 0.4);
        assert_eq!(result.len(), src.len());
    }
}