speech_prep/preprocessing/

noise_reduction.rs

//! Noise reduction via spectral subtraction.
//!
//! Reduces background noise while preserving speech intelligibility using
//! adaptive spectral subtraction with VAD-informed noise profiling.
//!
//! # Capabilities
//!
//! - **Stationary Noise Reduction**: Effectively removes constant background
//!   noise (HVAC hum, white noise, fan noise, café ambience)
//! - **≥6 dB SNR Improvement**: Validated on white noise, low-frequency hum,
//!   and ambient café noise
//! - **Phase Preservation**: Maintains speech intelligibility by preserving
//!   original signal phase
//! - **VAD Integration**: Adapts noise profile only during detected silence
//! - **Real-Time**: <15ms latency per 500ms chunk (typically 0.2-0.3ms)
//!
//! # Limitations
//!
//! **Spectral subtraction is designed for STATIONARY noise only.**
//!
//! - **Non-stationary noise**: Struggles with time-varying noise (individual
//!   voices, music, babble with distinct speakers). Use Wiener filtering or
//!   deep learning approaches for non-stationary scenarios.
//! - **Speech-like interference**: Cannot separate overlapping speakers or
//!   remove foreground speech interference (requires source separation
//!   techniques).
//! - **Musical noise artifacts**: Tonal artifacts may occur with aggressive
//!   settings. Mitigated via spectral floor parameter (β=0.02 default).
//! - **Transient noise**: Impulsive sounds (door slams, clicks) are not handled
//!   well. Consider median filtering for transient suppression.
//!
//! # When to Use This
//!
//! ✅ **Good fit**:
//! - Background HVAC/fan noise
//! - Café/restaurant ambient noise (general chatter blur, dishes)
//! - Low-frequency hum (electrical interference)
//! - Stationary white/pink noise
//!
//! ❌ **Poor fit**:
//! - Multi-speaker separation (babble with distinct voices)
//! - Music removal
//! - Non-stationary interference
//! - Echo/reverb reduction (use AEC instead)

use std::f32::consts::PI;
use std::sync::Arc;

use super::VadContext;
use crate::error::{Error, Result};
use crate::time::{AudioDuration, AudioInstant};
use realfft::{ComplexToReal, RealFftPlanner, RealToComplex};
use tracing::{info, warn};

/// Configuration for noise reduction via spectral subtraction.
///
/// # Examples
///
/// ```rust,no_run
/// use speech_prep::preprocessing::NoiseReductionConfig;
///
/// // Default: 25ms window, 10ms hop, α=2.0, β=0.02
/// let config = NoiseReductionConfig::default();
///
/// // Aggressive noise removal for noisy environment
/// let config = NoiseReductionConfig {
///     oversubtraction_factor: 2.5,
///     spectral_floor: 0.01,
///     ..Default::default()
/// };
/// # Ok::<(), speech_prep::error::Error>(())
/// ```
#[derive(Debug, Clone)]
#[allow(missing_copy_implementations)]
pub struct NoiseReductionConfig {
    /// Sample rate in Hz.
    ///
    /// **Default**: 16000
    /// **Range**: 8000-48000
    pub sample_rate_hz: u32,

    /// STFT window duration in milliseconds.
    ///
    /// **Default**: 25.0
    /// **Range**: 10.0-50.0
    ///
    /// **Effect**: Longer windows improve frequency resolution but reduce
    /// time resolution. 25ms captures 1-3 pitch periods for typical speech.
    pub window_ms: f32,

    /// STFT hop duration in milliseconds.
    ///
    /// **Default**: 10.0
    /// **Range**: 5.0-25.0
    ///
    /// **Effect**: Smaller hops increase overlap (smoother reconstruction)
    /// but require more computation. 10ms hop = 60% overlap with 25ms window.
    pub hop_ms: f32,

    /// Oversubtraction factor (α).
    ///
    /// **Default**: 2.0
    /// **Range**: 1.0-3.0
    ///
    /// **Effect**: Multiplier for noise estimate in spectral subtraction.
    /// - Higher: More aggressive noise removal, more artifacts
    /// - Lower: Conservative removal, less SNR gain
    pub oversubtraction_factor: f32,

    /// Spectral floor (β) as fraction of noise estimate.
    ///
    /// **Default**: 0.02 (2% of noise estimate)
    /// **Range**: 0.001-0.1
    ///
    /// **Effect**: Minimum magnitude after subtraction to prevent musical
    /// noise. Acts as a noise gate.
    pub spectral_floor: f32,

    /// Noise profile smoothing factor (`α_noise`).
    ///
    /// **Default**: 0.98
    /// **Range**: 0.9-0.999
    ///
    /// **Effect**: Exponential moving average smoothing for noise profile.
    /// Higher values = slower adaptation, more stable estimate.
    pub noise_smoothing: f32,

    /// Enable noise reduction.
    ///
    /// **Default**: true
    ///
    /// **Effect**: When false, audio passes through unmodified (bypass mode).
    pub enable: bool,
}

impl Default for NoiseReductionConfig {
    fn default() -> Self {
        Self {
            sample_rate_hz: 16_000,
            window_ms: 25.0,
            hop_ms: 10.0,
            oversubtraction_factor: 2.0,
            spectral_floor: 0.02,
            noise_smoothing: 0.98,
            enable: true,
        }
    }
}

impl NoiseReductionConfig {
    /// Validate configuration parameters.
    ///
    /// # Errors
    ///
    /// Returns `Error::Configuration` if:
    /// - `sample_rate_hz` not in 8000-48000 Hz
    /// - `window_ms` not in 10.0-50.0 ms
    /// - `hop_ms` >= `window_ms` (overlap required)
    /// - `oversubtraction_factor` not in 1.0-3.0
    /// - `spectral_floor` not in 0.001-0.1
    /// - `noise_smoothing` not in 0.9-0.999
    #[allow(clippy::trivially_copy_pass_by_ref)]
    pub fn validate(&self) -> Result<()> {
        if !(8000..=48_000).contains(&self.sample_rate_hz) {
            return Err(Error::Configuration(format!(
                "Invalid sample rate: {} Hz (range: 8000-48000)",
                self.sample_rate_hz
            )));
        }

        if !(10.0..=50.0).contains(&self.window_ms) {
            return Err(Error::Configuration(format!(
                "Invalid window size: {:.1} ms (range: 10-50)",
                self.window_ms
            )));
        }

        if self.hop_ms >= self.window_ms {
            return Err(Error::Configuration(format!(
                "Hop {:.1} ms must be < window {:.1} ms",
                self.hop_ms, self.window_ms
            )));
        }

        if !(1.0..=3.0).contains(&self.oversubtraction_factor) {
            return Err(Error::Configuration(format!(
                "Invalid oversubtraction factor: {:.2} (range: 1.0-3.0)",
                self.oversubtraction_factor
            )));
        }

        if !(0.001..=0.1).contains(&self.spectral_floor) {
            return Err(Error::Configuration(format!(
                "Invalid spectral floor: {:.3} (range: 0.001-0.1)",
                self.spectral_floor
            )));
        }

        if !(0.9..1.0).contains(&self.noise_smoothing) {
            return Err(Error::Configuration(format!(
                "Invalid noise smoothing: {:.3} (range: 0.9-0.999)",
                self.noise_smoothing
            )));
        }

        Ok(())
    }

    /// Calculate frame length in samples.
    pub fn frame_length(&self) -> usize {
        ((self.window_ms / 1000.0) * self.sample_rate_hz as f32).round() as usize
    }

    /// Calculate hop length in samples.
    pub fn hop_length(&self) -> usize {
        ((self.hop_ms / 1000.0) * self.sample_rate_hz as f32).round() as usize
    }

    /// Calculate FFT size (next power of 2 >= frame length).
    pub fn fft_size(&self) -> usize {
        self.frame_length().next_power_of_two()
    }
}

/// Noise reduction via spectral subtraction with adaptive noise profiling.
///
/// Implements the noise reduction specification:
/// - STFT-based processing (25ms window, 10ms hop)
/// - Adaptive noise profile estimation during VAD-detected silence
/// - Magnitude-only spectral subtraction (preserves phase)
/// - Achieves ≥6 dB SNR improvement target
///
/// # Performance
///
/// - **Target**: <15ms per 500ms chunk (8000 samples @ 16kHz)
/// - **Expected**: ~7ms (2x headroom)
/// - **Optimization**: Precomputed FFT plans, reused buffers
///
/// # Example
///
/// ```rust,no_run
/// use speech_prep::preprocessing::{NoiseReducer, NoiseReductionConfig, VadContext};
///
/// # fn main() -> speech_prep::error::Result<()> {
/// let config = NoiseReductionConfig::default();
/// let mut reducer = NoiseReducer::new(config)?;
/// let audio_stream = vec![vec![0.0; 8000], vec![0.05; 8080]];
///
/// // Process streaming chunks with VAD context
/// for chunk in audio_stream {
///     let vad_ctx = VadContext { is_silence: detect_silence(&chunk) };
///     let _denoised = reducer.reduce(&chunk, Some(vad_ctx))?;
/// }
/// # Ok(())
/// # }
/// #
/// # fn detect_silence(chunk: &[f32]) -> bool {
/// #     chunk.iter().all(|sample| sample.abs() < 1e-3)
/// # }
/// ```
#[allow(missing_copy_implementations)]
pub struct NoiseReducer {
    config: NoiseReductionConfig,
    fft_forward: Arc<dyn RealToComplex<f32>>,
    fft_inverse: Arc<dyn ComplexToReal<f32>>,
    window: Vec<f32>,
    noise_profile: Vec<f32>,
    noise_initialized: bool,
    overlap_buffer: Vec<f32>,
}

impl std::fmt::Debug for NoiseReducer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("NoiseReducer")
            .field("config", &self.config)
            .field("window_length", &self.window.len())
            .field("noise_profile_bins", &self.noise_profile.len())
            .field("noise_initialized", &self.noise_initialized)
            .finish_non_exhaustive()
    }
}

impl NoiseReducer {
    /// Create a new noise reducer.
    ///
    /// # Arguments
    ///
    /// * `config` - Configuration parameters (window size, hop, α, β)
    ///
    /// # Errors
    ///
    /// Returns `Error::Configuration` if configuration is invalid.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use speech_prep::preprocessing::{NoiseReducer, NoiseReductionConfig};
    ///
    /// let config = NoiseReductionConfig {
    ///     oversubtraction_factor: 2.5, // Aggressive
    ///     ..Default::default()
    /// };
    /// let reducer = NoiseReducer::new(config)?;
    /// # Ok::<(), speech_prep::error::Error>(())
    /// ```
    pub fn new(config: NoiseReductionConfig) -> Result<Self> {
        config.validate()?;

        let fft_size = config.fft_size();
        let frame_length = config.frame_length();

        let mut planner = RealFftPlanner::<f32>::new();
        let fft_forward = planner.plan_fft_forward(fft_size);
        let fft_inverse = planner.plan_fft_inverse(fft_size);

        let window = generate_hann_window(frame_length);

        let num_bins = fft_size / 2 + 1;
        let noise_profile = vec![1e-6; num_bins];

        let overlap_buffer = vec![0.0; frame_length];

        Ok(Self {
            config,
            fft_forward,
            fft_inverse,
            window,
            noise_profile,
            noise_initialized: false,
            overlap_buffer,
        })
    }

    /// Apply noise reduction to audio samples.
    ///
    /// # Arguments
    ///
    /// * `samples` - Input audio samples (typically 500ms chunk = 8000 samples
    ///   @ 16kHz)
    /// * `vad_context` - Optional VAD state for noise profile updates
    ///
    /// # Returns
    ///
    /// Denoised audio with ≥6 dB SNR improvement on noisy input.
    ///
    /// # Performance
    ///
    /// - Expected: ~7ms for 8000 samples (2x better than <15ms target)
    /// - Complexity: O(n log n) for FFT operations
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use speech_prep::preprocessing::{NoiseReducer, NoiseReductionConfig, VadContext};
    ///
    /// let mut reducer = NoiseReducer::new(NoiseReductionConfig::default())?;
    ///
    /// // Chunk 1 (silence - initialize noise profile)
    /// let chunk1 = vec![0.001; 8000];
    /// let vad1 = VadContext { is_silence: true };
    /// let output1 = reducer.reduce(&chunk1, Some(vad1))?;
    ///
    /// // Chunk 2 (speech - apply noise reduction)
    /// let chunk2 = vec![0.1; 8000];
    /// let vad2 = VadContext { is_silence: false };
    /// let output2 = reducer.reduce(&chunk2, Some(vad2))?;
    /// # Ok::<(), speech_prep::error::Error>(())
    /// ```
    #[allow(clippy::unnecessary_wraps)]
    pub fn reduce(&mut self, samples: &[f32], vad_context: Option<VadContext>) -> Result<Vec<f32>> {
        let processing_start = AudioInstant::now();

        if samples.is_empty() {
            return Ok(Vec::new());
        }

        if !self.config.enable {
            return Ok(samples.to_vec());
        }

        let (mut output, frame_count) = self.process_stft_frames(samples, vad_context)?;
        self.normalize_overlap_add(&mut output);

        let elapsed = elapsed_duration(processing_start);
        let latency_ms = elapsed.as_secs_f64() * 1000.0;
        self.record_performance_metrics(samples, &output, latency_ms, frame_count);

        Ok(output)
    }

    /// Process audio through STFT frames with spectral subtraction.
    fn process_stft_frames(
        &mut self,
        samples: &[f32],
        vad_context: Option<VadContext>,
    ) -> Result<(Vec<f32>, usize)> {
        let frame_length = self.config.frame_length();
        let hop_length = self.config.hop_length();

        let mut output = vec![0.0; samples.len()];
        let mut frame_idx = 0;
        let mut pos = 0;

        while pos < samples.len() {
            let remaining = samples.len() - pos;

            let frame = Self::extract_frame(samples, pos, frame_length, remaining)?;

            let processed =
                self.process_single_frame(&frame, vad_context, remaining >= frame_length)?;

            Self::accumulate_frame_output(&processed, &mut output, pos);

            frame_idx += 1;

            if remaining < hop_length {
                break;
            }
            pos += hop_length;
        }

        Ok((output, frame_idx))
    }

    /// Extract a frame from the input samples, zero-padding if partial.
    fn extract_frame(
        samples: &[f32],
        pos: usize,
        frame_length: usize,
        remaining: usize,
    ) -> Result<Vec<f32>> {
        let mut frame_buf = vec![0.0; frame_length];

        if remaining >= frame_length {
            let src = samples
                .get(pos..pos + frame_length)
                .ok_or_else(|| Error::Processing("frame window out of bounds".into()))?;
            frame_buf.copy_from_slice(src);
        } else {
            let src = samples
                .get(pos..)
                .ok_or_else(|| Error::Processing("frame tail out of bounds".into()))?;
            if let Some(dst) = frame_buf.get_mut(..remaining) {
                dst.copy_from_slice(src);
            }
        }

        Ok(frame_buf)
    }

    /// Process a single frame through FFT, spectral subtraction, and IFFT.
    fn process_single_frame(
        &mut self,
        frame: &[f32],
        vad_context: Option<VadContext>,
        is_full_frame: bool,
    ) -> Result<Vec<f32>> {
        let fft_size = self.config.fft_size();

        let windowed: Vec<f32> = frame
            .iter()
            .zip(&self.window)
            .map(|(&s, &w)| s * w)
            .collect();

        let complex_spectrum = self.forward_fft_complex(&windowed)?;
        let magnitudes: Vec<f32> = complex_spectrum.iter().map(|c| c.norm()).collect();

        let is_silence = vad_context.is_some_and(|ctx| ctx.is_silence);
        if is_silence && is_full_frame {
            self.update_noise_profile(&magnitudes);
        }

        let cleaned_magnitudes = self.spectral_subtract(&magnitudes);

        let cleaned_complex =
            Self::reconstruct_complex_spectrum(&complex_spectrum, &cleaned_magnitudes);

        let time_signal = self.inverse_fft_complex(&cleaned_complex, fft_size)?;

        let windowed_output: Vec<f32> = time_signal
            .iter()
            .take(frame.len())
            .zip(&self.window)
            .map(|(&s, &w)| s * w)
            .collect();

        Ok(windowed_output)
    }

    /// Reconstruct complex spectrum preserving phase from original signal.
    fn reconstruct_complex_spectrum(
        original_spectrum: &[realfft::num_complex::Complex<f32>],
        cleaned_magnitudes: &[f32],
    ) -> Vec<realfft::num_complex::Complex<f32>> {
        original_spectrum
            .iter()
            .zip(cleaned_magnitudes)
            .enumerate()
            .map(|(i, (original, &new_mag))| {
                if i == 0 || i == original_spectrum.len() - 1 {
                    // DC and Nyquist bins must be real-valued
                    realfft::num_complex::Complex::new(new_mag, 0.0)
                } else {
                    let phase = original.arg();
                    realfft::num_complex::Complex::from_polar(new_mag, phase)
                }
            })
            .collect()
    }

    /// Accumulate processed frame into output buffer (overlap-add).
    fn accumulate_frame_output(frame: &[f32], output: &mut [f32], pos: usize) {
        for (i, &sample) in frame.iter().enumerate() {
            let out_idx = pos + i;
            if let Some(dst) = output.get_mut(out_idx) {
                *dst += sample;
            }
        }
    }

    /// Normalize output by overlap-add window sum.
    fn normalize_overlap_add(&self, output: &mut [f32]) {
        let hop_length = self.config.hop_length();
        let window_sum = self.calculate_window_overlap_sum(hop_length);

        if window_sum > 1e-6 {
            for sample in output {
                *sample /= window_sum;
            }
        }
    }

    fn record_performance_metrics(
        &self,
        input: &[f32],
        output: &[f32],
        latency_ms: f64,
        frame_count: usize,
    ) {
        if input.len() < 8000 {
            return;
        }

        if latency_ms > 15.0 {
            warn!(
                target: "audio.preprocess.noise_reduction",
                latency_ms,
                samples = input.len(),
                frames = frame_count,
                oversubtraction = self.config.oversubtraction_factor,
                spectral_floor = self.config.spectral_floor,
                "noise reduction latency exceeded target"
            );
        }

        let avg_noise_floor = self.noise_floor().max(1e-12);
        let noise_floor_db = 20.0 * avg_noise_floor.log10();

        let signal_power_out =
            output.iter().map(|sample| sample * sample).sum::<f32>() / output.len() as f32;
        let residual_power: f32 = input
            .iter()
            .zip(output)
            .map(|(&noisy, &clean)| {
                let residual = noisy - clean;
                residual * residual
            })
            .sum::<f32>()
            / output.len() as f32;

        let snr_improvement_db = if residual_power > 1e-12 && signal_power_out > 0.0 {
            10.0 * (signal_power_out / residual_power).log10()
        } else {
            0.0
        };

        info!(
            target: "audio.preprocess.noise_reduction",
            noise_floor_db,
            snr_improvement_db,
            latency_ms,
            frames = frame_count,
            samples = input.len(),
            oversubtraction = self.config.oversubtraction_factor,
            spectral_floor = self.config.spectral_floor,
            "noise reduction metrics"
        );
    }

    /// Reset noise profile for new audio stream.
    ///
    /// Clears noise estimate and overlap-add state.
    /// Use this when starting a new, independent audio stream.
    pub fn reset(&mut self) {
        self.noise_profile.fill(1e-6);
        self.noise_initialized = false;
        self.overlap_buffer.fill(0.0);
    }

    /// Get current average noise floor (for debugging/observability).
    #[must_use]
    pub fn noise_floor(&self) -> f32 {
        if self.noise_profile.is_empty() {
            return 0.0;
        }
        self.noise_profile.iter().sum::<f32>() / self.noise_profile.len() as f32
    }

    /// Get current configuration.
    #[must_use]
    pub fn config(&self) -> &NoiseReductionConfig {
        &self.config
    }

    // Forward FFT with zero-padding (returns complex spectrum)
    fn forward_fft_complex(
        &self,
        windowed: &[f32],
    ) -> Result<Vec<realfft::num_complex::Complex<f32>>> {
        // Prepare input buffer (zero-padded to FFT size)
        let mut input = self.fft_forward.make_input_vec();
        for (i, &sample) in windowed.iter().enumerate() {
            if let Some(dst) = input.get_mut(i) {
                *dst = sample;
            }
        }

        // Perform FFT
        let mut spectrum = self.fft_forward.make_output_vec();
        self.fft_forward
            .process(&mut input, &mut spectrum)
            .map_err(|e| Error::Processing(format!("FFT failed: {e}")))?;

        Ok(spectrum)
    }

    // Inverse FFT from complex spectrum (preserves phase)
    fn inverse_fft_complex(
        &self,
        complex_spectrum: &[realfft::num_complex::Complex<f32>],
        fft_size: usize,
    ) -> Result<Vec<f32>> {
        // Prepare input buffer
        let mut spectrum = self.fft_inverse.make_input_vec();
        for (i, &c) in complex_spectrum.iter().enumerate() {
            if let Some(bin) = spectrum.get_mut(i) {
                *bin = c;
            }
        }

        // Perform inverse FFT
        let mut output = self.fft_inverse.make_output_vec();
        self.fft_inverse
            .process(&mut spectrum, &mut output)
            .map_err(|e| Error::Processing(format!("IFFT failed: {e}")))?;

        // Normalize by FFT size
        for sample in &mut output {
            *sample /= fft_size as f32;
        }

        Ok(output)
    }

    // Update noise profile using exponential moving average
    fn update_noise_profile(&mut self, spectrum: &[f32]) {
        let alpha = self.config.noise_smoothing;

        if self.noise_initialized {
            // EMA update: N_new[k] = α * N_old[k] + (1-α) * |X[k]|
            for (noise, &current) in self.noise_profile.iter_mut().zip(spectrum.iter()) {
                *noise = alpha.mul_add(*noise, (1.0 - alpha) * current);
            }
        } else {
            // First silence frame: initialize noise profile
            self.noise_profile.copy_from_slice(spectrum);
            self.noise_initialized = true;
        }
    }

    // Apply spectral subtraction: |Y[k]| = max(|X[k]| - α*|N[k]|, β*|N[k]|)
    fn spectral_subtract(&self, spectrum: &[f32]) -> Vec<f32> {
        let alpha = self.config.oversubtraction_factor;
        let beta = self.config.spectral_floor;

        spectrum
            .iter()
            .zip(&self.noise_profile)
            .map(|(&signal, &noise)| {
                let subtracted = alpha.mul_add(-noise, signal);
                let floor = beta * noise;
                subtracted.max(floor)
            })
            .collect()
    }

    // Calculate window overlap sum for COLA normalization
    fn calculate_window_overlap_sum(&self, hop_length: usize) -> f32 {
        let frame_length = self.window.len();
        let mut sum: f32 = 0.0;

        // Sum overlapping windows at each sample position
        for i in 0..frame_length {
            let mut overlap: f32 = 0.0;
            let mut offset = 0;

            while offset <= i {
                if let Some(&w) = self.window.get(i - offset) {
                    overlap = w.mul_add(w, overlap); // Window applied twice
                                                     // (analysis +
                                                     // synthesis)
                }
                offset += hop_length;
            }

            sum = sum.max(overlap);
        }

        sum
    }
}

/// Generate Hann window.
///
/// Formula: `w[n] = 0.5 - 0.5 * cos(2π * n / (N-1))`
fn generate_hann_window(length: usize) -> Vec<f32> {
    if length == 0 {
        return Vec::new();
    }

    if length == 1 {
        return vec![1.0];
    }

    let denom = (length - 1) as f32;
    (0..length)
        .map(|n| {
            let angle = 2.0 * PI * n as f32 / denom;
            0.5f32.mul_add(-angle.cos(), 0.5)
        })
        .collect()
}

fn elapsed_duration(start: AudioInstant) -> AudioDuration {
    AudioInstant::now().duration_since(start)
}

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

    type TestResult<T> = std::result::Result<T, String>;

    #[test]
    #[allow(clippy::unnecessary_wraps)]
    fn test_configuration_validation() -> TestResult<()> {
        // Valid configuration
        let valid = NoiseReductionConfig::default();
        assert!(valid.validate().is_ok());

        // Invalid sample rate
        let invalid_sr = NoiseReductionConfig {
            sample_rate_hz: 5000,
            ..Default::default()
        };
        assert!(invalid_sr.validate().is_err());

        // Invalid window size
        let invalid_window = NoiseReductionConfig {
            window_ms: 100.0,
            ..Default::default()
        };
        assert!(invalid_window.validate().is_err());

        // Hop >= window
        let invalid_hop = NoiseReductionConfig {
            hop_ms: 30.0,
            window_ms: 25.0,
            ..Default::default()
        };
        assert!(invalid_hop.validate().is_err());

        // Invalid oversubtraction
        let invalid_alpha = NoiseReductionConfig {
            oversubtraction_factor: 5.0,
            ..Default::default()
        };
        assert!(invalid_alpha.validate().is_err());

        // Invalid spectral floor
        let invalid_beta = NoiseReductionConfig {
            spectral_floor: 0.5,
            ..Default::default()
        };
        assert!(invalid_beta.validate().is_err());

        // Invalid noise smoothing
        let invalid_smoothing = NoiseReductionConfig {
            noise_smoothing: 1.0,
            ..Default::default()
        };
        assert!(invalid_smoothing.validate().is_err());

        Ok(())
    }

    #[test]
    fn test_hann_window_properties() {
        // Empty window
        let window_0 = generate_hann_window(0);
        assert!(window_0.is_empty());

        // Single element
        let window_1 = generate_hann_window(1);
        assert_eq!(window_1.len(), 1);
        assert!((window_1[0] - 1.0).abs() < 1e-6);

        // Check Hann window properties
        let window = generate_hann_window(100);
        assert_eq!(window.len(), 100);

        // First and last samples should be near zero
        assert!(window[0].abs() < 1e-6);
        assert!(window[99].abs() < 1e-6);

        // Middle sample should be near 1.0
        assert!((window[50] - 1.0).abs() < 0.1);
    }

    #[test]
    fn test_noise_reducer_creation() -> TestResult<()> {
        let config = NoiseReductionConfig::default();
        let reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        // Check initialization
        assert_eq!(reducer.config().sample_rate_hz, 16000);
        assert!(reducer.noise_floor() > 0.0); // Initial estimate

        Ok(())
    }

    #[test]
    fn test_empty_input() -> TestResult<()> {
        let config = NoiseReductionConfig::default();
        let mut reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        let output = reducer.reduce(&[], None).map_err(|e| e.to_string())?;
        assert!(output.is_empty());

        Ok(())
    }

    #[test]
    fn test_bypass_mode() -> TestResult<()> {
        let config = NoiseReductionConfig {
            enable: false,
            ..Default::default()
        };
        let mut reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        let input = vec![0.1, 0.2, 0.3, 0.4];
        let output = reducer.reduce(&input, None).map_err(|e| e.to_string())?;

        // Bypass mode should return input unchanged
        assert_eq!(output, input);

        Ok(())
    }

    #[test]
    fn test_noise_profile_update() -> TestResult<()> {
        let config = NoiseReductionConfig::default();
        let mut reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        // Process silence to build noise profile
        let silence = vec![0.01; 8000]; // Low-level noise
        let vad_silence = VadContext { is_silence: true };

        let initial_noise = reducer.noise_floor();

        // Process multiple chunks to converge
        for _ in 0..5 {
            let _ = reducer
                .reduce(&silence, Some(vad_silence))
                .map_err(|e| e.to_string())?;
        }

        let converged_noise = reducer.noise_floor();

        // Noise floor should increase from initial estimate
        assert!(
            converged_noise > initial_noise,
            "Noise floor should adapt: initial={:.6}, converged={:.6}",
            initial_noise,
            converged_noise
        );

        Ok(())
    }

    #[test]
    fn test_vad_informed_noise_update() -> TestResult<()> {
        let config = NoiseReductionConfig::default();
        let mut reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        // Initialize with silence
        let silence = vec![0.01; 8000];
        let vad_silence = VadContext { is_silence: true };
        for _ in 0..5 {
            let _ = reducer
                .reduce(&silence, Some(vad_silence))
                .map_err(|e| e.to_string())?;
        }

        let noise_after_silence = reducer.noise_floor();

        // Process "speech" (should NOT update noise profile)
        let speech = vec![0.5; 8000];
        let vad_speech = VadContext { is_silence: false };
        let _ = reducer
            .reduce(&speech, Some(vad_speech))
            .map_err(|e| e.to_string())?;

        let noise_after_speech = reducer.noise_floor();

        // Noise profile should remain stable during speech
        let diff = (noise_after_speech - noise_after_silence).abs();
        assert!(
            diff < noise_after_silence * 0.01,
            "Noise profile changed during speech: {:.6} -> {:.6}",
            noise_after_silence,
            noise_after_speech
        );

        Ok(())
    }

    #[test]
    fn test_reset_clears_state() -> TestResult<()> {
        let config = NoiseReductionConfig::default();
        let mut reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        // Process some audio
        let samples = vec![0.1; 8000];
        let vad = VadContext { is_silence: true };
        let _ = reducer
            .reduce(&samples, Some(vad))
            .map_err(|e| e.to_string())?;

        let noise_before = reducer.noise_floor();
        assert!(noise_before > 1e-5, "Noise profile should be updated");

        // Reset
        reducer.reset();

        let noise_after = reducer.noise_floor();
        assert!(
            noise_after < 1e-5,
            "Noise profile should be reset to initial value"
        );

        Ok(())
    }

    // Helper: Generate sine wave
    fn generate_sine_wave(
        frequency: f32,
        sample_rate: u32,
        duration_secs: f32,
        amplitude: f32,
    ) -> Vec<f32> {
        use std::f32::consts::PI;
        let samples = (sample_rate as f32 * duration_secs).round() as usize;
        (0..samples)
            .map(|i| {
                let t = i as f32 / sample_rate as f32;
                (2.0 * PI * frequency * t).sin() * amplitude
            })
            .collect()
    }

    // Helper: Add white noise to signal
    fn add_white_noise(signal: &[f32], noise_amplitude: f32) -> Vec<f32> {
        use rand::Rng;
        let mut rng = rand::rng();
        signal
            .iter()
            .map(|&s| {
                let noise: f32 = rng.random_range(-noise_amplitude..noise_amplitude);
                s + noise
            })
            .collect()
    }
    fn add_low_freq_hum(
        signal: &[f32],
        sample_rate: u32,
        frequency: f32,
        amplitude: f32,
    ) -> Vec<f32> {
        signal
            .iter()
            .enumerate()
            .map(|(i, &sample)| {
                let t = i as f32 / sample_rate as f32;
                let hum = (2.0 * PI * frequency * t).sin() * amplitude;
                sample + hum
            })
            .collect()
    }

    // Helper: Add café-like ambient noise (stationary broadband noise)
    // Simulates background café noise: HVAC, dishes, distant ambient chatter.
    // Uses white noise as a proxy for band-limited stationary noise (100-3000 Hz
    // typical). NOTE: Spectral subtraction works for STATIONARY noise, not
    // speech-like babble.
    fn add_cafe_noise(signal: &[f32], _sample_rate: u32, amplitude: f32) -> Vec<f32> {
        use rand::Rng;
        let mut rng = rand::rng();
        signal
            .iter()
            .map(|&sample| {
                let noise: f32 = rng.random_range(-1.0..1.0);
                amplitude.mul_add(noise, sample)
            })
            .collect()
    }

    // Helper: Calculate SNR
    fn calculate_snr(clean: &[f32], noisy: &[f32]) -> f32 {
        if clean.len() != noisy.len() {
            return 0.0;
        }

        let signal_power: f32 = clean.iter().map(|&x| x * x).sum();
        let noise: Vec<f32> = clean
            .iter()
            .zip(noisy.iter())
            .map(|(&c, &n)| n - c)
            .collect();
        let noise_power: f32 = noise.iter().map(|&x| x * x).sum();

        if noise_power < 1e-10 {
            return 100.0; // Very high SNR
        }

        10.0 * (signal_power / noise_power).log10()
    }

    #[test]
    fn test_snr_improvement_white_noise() -> TestResult<()> {
        let config = NoiseReductionConfig::default();
        let mut reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        // Generate clean speech signal (440 Hz sine wave)
        let clean_speech = generate_sine_wave(440.0, 16000, 1.0, 0.5);

        // Add white noise (creating ~5 dB input SNR)
        let noisy_speech = add_white_noise(&clean_speech, 0.3);

        let snr_before = calculate_snr(&clean_speech, &noisy_speech);

        // Initialize noise profile with pure noise
        // Training: 10 iterations ensures EMA convergence (α=0.98 requires ~50 samples
        // for 95% convergence)
        let pure_noise = add_white_noise(&vec![0.0; 8000], 0.3);
        let vad_silence = VadContext { is_silence: true };
        for _ in 0..10 {
            let _ = reducer
                .reduce(&pure_noise, Some(vad_silence))
                .map_err(|e| e.to_string())?;
        }

        // Apply noise reduction to noisy speech
        let vad_speech = VadContext { is_silence: false };
        let denoised = reducer
            .reduce(&noisy_speech, Some(vad_speech))
            .map_err(|e| e.to_string())?;

        let snr_after = calculate_snr(&clean_speech, &denoised);
        let improvement = snr_after - snr_before;

        // Success criterion: ≥6 dB improvement
        assert!(
            improvement >= 6.0,
            "SNR improvement {:.1} dB < 6 dB target",
            improvement
        );

        Ok(())
    }

    #[test]
    fn test_snr_improvement_low_freq_hum() -> TestResult<()> {
        let config = NoiseReductionConfig::default();
        let mut reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        // Generate 440 Hz speech with 60 Hz HVAC hum (common electrical interference)
        let clean = generate_sine_wave(440.0, 16000, 1.0, 0.4);
        let noisy = add_low_freq_hum(&clean, 16000, 60.0, 0.3);
        let snr_before = calculate_snr(&clean, &noisy);

        // Train on pure 60 Hz hum
        // Training: 6 iterations sufficient for tonal noise (faster convergence than
        // broadband)
        let hum_only = add_low_freq_hum(&vec![0.0; 8000], 16000, 60.0, 0.3);
        let vad = VadContext { is_silence: true };
        for _ in 0..6 {
            let _ = reducer
                .reduce(&hum_only, Some(vad))
                .map_err(|e| e.to_string())?;
        }

        let vad_speech = VadContext { is_silence: false };
        let denoised = reducer
            .reduce(&noisy, Some(vad_speech))
            .map_err(|e| e.to_string())?;
        let snr_after = calculate_snr(&clean, &denoised);
        let improvement = snr_after - snr_before;
        assert!(
            improvement >= 6.0,
            "Hum SNR improvement {:.1} dB < 6 dB target",
            improvement
        );

        Ok(())
    }

    #[test]
    fn test_snr_improvement_cafe_ambient() -> TestResult<()> {
        // Use default config - stationary noise doesn't need aggressive oversubtraction
        let config = NoiseReductionConfig::default();
        let mut reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        // Generate clean speech signal
        let clean = generate_sine_wave(220.0, 16000, 1.0, 0.4);

        // Add stationary café ambient noise (HVAC, dishes, background chatter)
        let noisy = add_cafe_noise(&clean, 16000, 0.25);
        let snr_before = calculate_snr(&clean, &noisy);

        // Train noise profile on café ambient noise during "silence"
        // Training: 10 iterations for broadband stationary noise (white noise requires
        // more samples than tonal) Noise amplitude 0.25 creates ~5-6 dB input
        // SNR, realistic for café environment
        let cafe_only = add_cafe_noise(&vec![0.0; 8000], 16000, 0.25);
        let vad = VadContext { is_silence: true };
        for _ in 0..10 {
            let _ = reducer
                .reduce(&cafe_only, Some(vad))
                .map_err(|e| e.to_string())?;
        }

        // Apply noise reduction to noisy speech
        let vad_speech = VadContext { is_silence: false };
        let denoised = reducer
            .reduce(&noisy, Some(vad_speech))
            .map_err(|e| e.to_string())?;

        let snr_after = calculate_snr(&clean, &denoised);
        let improvement = snr_after - snr_before;

        assert!(
            improvement >= 6.0,
            "Café ambient SNR improvement {:.1} dB < 6 dB target",
            improvement
        );

        Ok(())
    }

    #[test]
    fn test_trailing_partial_frame_preserved() -> TestResult<()> {
        let config = NoiseReductionConfig::default();
        let mut reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        // Prime noise profile with a silence chunk so spectral subtraction behaves
        // normally.
        let silence = vec![0.0; 8000];
        let vad_silence = VadContext { is_silence: true };
        let _ = reducer
            .reduce(&silence, Some(vad_silence))
            .map_err(|e| e.to_string())?;

        // Speech chunk length intentionally not divisible by hop size (adds 80-sample
        // tail).
        let speech_len = 8080;
        let speech: Vec<f32> = (0..speech_len)
            .map(|i| {
                let phase = (i as f32 / speech_len as f32) * 20.0;
                phase.sin()
            })
            .collect();

        let vad_speech = VadContext { is_silence: false };
        let output = reducer
            .reduce(&speech, Some(vad_speech))
            .map_err(|e| e.to_string())?;

        assert_eq!(
            output.len(),
            speech_len,
            "Output length should match input length"
        );

        let tail = &output[speech_len - 80..];
        let tail_energy: f32 = tail.iter().map(|sample| sample.abs()).sum();
        assert!(
            tail_energy > 1e-3,
            "Trailing samples should retain energy, got tail_energy={tail_energy}"
        );

        Ok(())
    }

    #[test]
    fn test_missing_vad_context_does_not_update_noise_profile() -> TestResult<()> {
        let config = NoiseReductionConfig::default();
        let mut reducer = NoiseReducer::new(config).map_err(|e| e.to_string())?;

        // Prime noise profile with explicit silence (non-zero noise so baseline is
        // measurable).
        let ambient_noise = vec![0.05f32; 8000];
        let vad_silence = VadContext { is_silence: true };
        reducer
            .reduce(&ambient_noise, Some(vad_silence))
            .map_err(|e| e.to_string())?;
        let baseline_floor = reducer.noise_floor();

        // Process speech without VAD context; noise profile should remain unchanged.
        let speech = vec![0.2f32; 8000];
        let output = reducer.reduce(&speech, None).map_err(|e| e.to_string())?;
        let updated_floor = reducer.noise_floor();

        let floor_delta = (updated_floor - baseline_floor).abs();
        assert!(
            floor_delta < baseline_floor.max(1e-6) * 0.01,
            "Noise floor changed when VAD context missing: baseline={baseline_floor}, \
             updated={updated_floor}"
        );

        let output_rms =
            (output.iter().map(|sample| sample * sample).sum::<f32>() / output.len() as f32).sqrt();
        assert!(
            output_rms > 0.08,
            "Speech energy collapsed without VAD context (rms={output_rms})"
        );

        Ok(())
    }
}