//! Noise Suppression — Wiener filter with minimum statistics noise floor.
//!
//! ## Algorithm
//!
//! Frequency-domain Wiener filter:
//!
//! 1. **Analysis** — 480-sample frame zero-padded to 1024, Hann-windowed,
//!    transformed with a correct radix-2 DIT complex FFT.
//!
//! 2. **Noise floor** — per-bin minimum statistics tracker (Martin 2001).
//!    Updated only in noise-dominated bins (SNR < VAD_SNR_THRESHOLD).
//!    Bias correction factor 1.5 accounts for minimum-statistics
//!    underestimation.
//!
//! 3. **Wiener gain** — G(k) = max(SNR(k)/(SNR(k)+1), GAIN_FLOOR).
//!    Floor at −20 dB prevents musical noise artefacts.
//!
//! 4. **Synthesis** — gain-weighted spectrum → IFFT → overlap-add.
//!
//! ## Realtime safety
//!
//! All buffers pre-allocated. No heap allocation in the hot path.

#![allow(clippy::needless_range_loop)]

use super::super::FRAME_SAMPLES;

const NS_FFT: usize = 1024;
const NS_BINS: usize = NS_FFT / 2 + 1;
/// Wiener gain floor (−20 dB).
const GAIN_FLOOR: f32 = 0.1;
/// Noise PSD smoothing (per-frame IIR).
const NOISE_ALPHA: f32 = 0.98;
/// Bias correction for minimum-statistics underestimation.
const BIAS: f32 = 1.5;
/// Bins with SNR below this are treated as noise-only.
const VAD_SNR_THRESHOLD: f32 = 1.5;

/// Wiener filter noise suppressor.
pub struct NoiseSuppressor {
    /// Per-bin noise PSD estimate.
    noise_psd: Box<[f32; NS_BINS]>,
    /// Overlap-add tail from the previous frame.
    ola_tail: Box<[f32; FRAME_SAMPLES]>,
    /// Hann window (NS_FFT length).
    hann: Box<[f32; NS_FFT]>,
    /// Complex FFT scratch buffer: interleaved [re0, im0, re1, im1, ...].
    /// Length = 2 * NS_FFT.
    fft_buf: Vec<f32>,
    enabled: bool,
    frame_count: u32,
}

impl NoiseSuppressor {
    /// Construct a noise suppressor with the P1 default estimator state.
    pub fn new() -> Self {
        let mut hann = Box::new([0.0_f32; NS_FFT]);
        for (i, h) in hann.iter_mut().enumerate() {
            *h = 0.5 * (1.0 - (2.0 * std::f32::consts::PI * i as f32 / NS_FFT as f32).cos());
        }
        Self {
            noise_psd: Box::new([1e-6_f32; NS_BINS]),
            ola_tail: Box::new([0.0_f32; FRAME_SAMPLES]),
            hann,
            fft_buf: vec![0.0_f32; 2 * NS_FFT],
            enabled: true,
            frame_count: 0,
        }
    }

    /// Enable or disable noise suppression. When disabled `process` is a no-op.
    pub fn set_enabled(&mut self, enabled: bool) {
        self.enabled = enabled;
    }

    /// Process one 10 ms capture frame in-place. Realtime-safe.
    pub fn process(&mut self, frame: &mut [f32; FRAME_SAMPLES]) {
        if !self.enabled {
            return;
        }
        self.frame_count = self.frame_count.saturating_add(1);

        // Build complex analysis buffer: real = windowed frame, imag = 0.
        // Zero-pad from FRAME_SAMPLES to NS_FFT.
        for i in 0..NS_FFT {
            let re = if i < FRAME_SAMPLES {
                frame[i] * self.hann[i]
            } else {
                0.0
            };
            self.fft_buf[2 * i] = re;
            self.fft_buf[2 * i + 1] = 0.0;
        }

        // Forward FFT.
        fft_complex_forward(&mut self.fft_buf, NS_FFT);

        // Compute per-bin power spectrum from complex output.
        let mut power = [0.0_f32; NS_BINS];
        for k in 0..NS_BINS {
            let re = self.fft_buf[2 * k];
            let im = self.fft_buf[2 * k + 1];
            power[k] = re * re + im * im;
        }

        // Cold-start: accumulate noise floor for 20 frames without suppression.
        if self.frame_count <= 20 {
            for k in 0..NS_BINS {
                self.noise_psd[k] =
                    self.noise_psd[k] * NOISE_ALPHA + power[k] * (1.0 - NOISE_ALPHA);
            }
            return;
        }

        // Compute Wiener gain and update noise floor.
        let mut gain = [0.0_f32; NS_BINS];
        for k in 0..NS_BINS {
            let noise = self.noise_psd[k] * BIAS;
            let snr = ((power[k] - noise) / noise.max(1e-10)).max(0.0);
            gain[k] = (snr / (snr + 1.0)).max(GAIN_FLOOR);
            // Update noise PSD only in noise-dominated bins.
            if snr < VAD_SNR_THRESHOLD {
                self.noise_psd[k] =
                    self.noise_psd[k] * NOISE_ALPHA + power[k] * (1.0 - NOISE_ALPHA);
            }
        }

        // Apply gain to the complex spectrum.
        // Bins 0..NS_BINS are the positive-frequency half.
        // Mirror to the negative-frequency half (conjugate symmetry).
        for k in 0..NS_BINS {
            self.fft_buf[2 * k] *= gain[k];
            self.fft_buf[2 * k + 1] *= gain[k];
        }
        // Mirror: bin k maps to bin NS_FFT - k.
        for k in 1..NS_BINS - 1 {
            let mirror = NS_FFT - k;
            self.fft_buf[2 * mirror] = self.fft_buf[2 * k];
            self.fft_buf[2 * mirror + 1] = -self.fft_buf[2 * k + 1]; // conjugate
        }

        // Inverse FFT.
        fft_complex_inverse(&mut self.fft_buf, NS_FFT);

        // Overlap-add: output = IFFT real part + previous tail.
        for i in 0..FRAME_SAMPLES {
            frame[i] = self.fft_buf[2 * i] + self.ola_tail[i];
        }
        // Save tail for next frame.
        for i in 0..FRAME_SAMPLES {
            self.ola_tail[i] = if i + FRAME_SAMPLES < NS_FFT {
                self.fft_buf[2 * (i + FRAME_SAMPLES)]
            } else {
                0.0
            };
        }
    }

    /// Reset all state.
    pub fn reset(&mut self) {
        self.noise_psd.fill(1e-6);
        self.ola_tail.fill(0.0);
        self.fft_buf.fill(0.0);
        self.frame_count = 0;
    }
}

impl Default for NoiseSuppressor {
    fn default() -> Self {
        Self::new()
    }
}

// ── Correct radix-2 DIT complex FFT ──────────────────────────────────────
//
// Buffer layout: interleaved [re0, im0, re1, im1, ..., re_{n-1}, im_{n-1}].
// Length of buf must be 2*n where n is a power of 2.

fn fft_complex_forward(buf: &mut [f32], n: usize) {
    debug_assert_eq!(buf.len(), 2 * n);
    debug_assert!(n.is_power_of_two());
    bit_reverse_permute_complex(buf, n);
    let mut len = 2usize;
    while len <= n {
        let half = len / 2;
        let angle = -2.0 * std::f32::consts::PI / len as f32;
        let (wre, wim) = (angle.cos(), angle.sin());
        let mut start = 0;
        while start < n {
            let (mut cur_re, mut cur_im) = (1.0_f32, 0.0_f32);
            for j in 0..half {
                let u_re = buf[2 * (start + j)];
                let u_im = buf[2 * (start + j) + 1];
                let v_re = buf[2 * (start + j + half)];
                let v_im = buf[2 * (start + j + half) + 1];
                // twiddle * v
                let tv_re = v_re * cur_re - v_im * cur_im;
                let tv_im = v_re * cur_im + v_im * cur_re;
                buf[2 * (start + j)] = u_re + tv_re;
                buf[2 * (start + j) + 1] = u_im + tv_im;
                buf[2 * (start + j + half)] = u_re - tv_re;
                buf[2 * (start + j + half) + 1] = u_im - tv_im;
                // advance twiddle
                let new_re = cur_re * wre - cur_im * wim;
                let new_im = cur_re * wim + cur_im * wre;
                cur_re = new_re;
                cur_im = new_im;
            }
            start += len;
        }
        len *= 2;
    }
}

fn fft_complex_inverse(buf: &mut [f32], n: usize) {
    // Conjugate input.
    for k in 0..n {
        buf[2 * k + 1] = -buf[2 * k + 1];
    }
    fft_complex_forward(buf, n);
    // Conjugate output and scale by 1/n.
    let scale = 1.0 / n as f32;
    for k in 0..n {
        buf[2 * k] *= scale;
        buf[2 * k + 1] = -buf[2 * k + 1] * scale;
    }
}

fn bit_reverse_permute_complex(buf: &mut [f32], n: usize) {
    let bits = n.trailing_zeros() as usize;
    for i in 0..n {
        let j = reverse_bits(i, bits);
        if j > i {
            buf.swap(2 * i, 2 * j);
            buf.swap(2 * i + 1, 2 * j + 1);
        }
    }
}

fn reverse_bits(mut x: usize, bits: usize) -> usize {
    let mut r = 0usize;
    for _ in 0..bits {
        r = (r << 1) | (x & 1);
        x >>= 1;
    }
    r
}

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

    #[test]
    fn fft_roundtrip() {
        // FFT then IFFT of a known signal should recover the original.
        let mut buf = vec![0.0_f32; 2 * 8];
        // Input: [1, 2, 3, 4, 0, 0, 0, 0] (real only)
        for i in 0..4 {
            buf[2 * i] = (i + 1) as f32;
        }
        let original: Vec<f32> = buf.iter().step_by(2).take(8).copied().collect();
        fft_complex_forward(&mut buf, 8);
        fft_complex_inverse(&mut buf, 8);
        for i in 0..8 {
            assert!(
                (buf[2 * i] - original[i]).abs() < 1e-4,
                "roundtrip failed at {i}: got {} expected {}",
                buf[2 * i],
                original[i]
            );
        }
    }

    #[test]
    fn ns_reduces_stationary_noise() {
        let mut ns = NoiseSuppressor::new();
        let mut rng: u32 = 0xDEAD_BEEF;
        let noise_frame = |rng: &mut u32| -> [f32; FRAME_SAMPLES] {
            let mut f = [0.0_f32; FRAME_SAMPLES];
            for s in f.iter_mut() {
                *rng = rng.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
                *s = (*rng as i32 as f32) / i32::MAX as f32 * 0.05;
            }
            f
        };
        // Warm up noise floor (20 frames cold-start + 10 more to converge).
        for _ in 0..30 {
            let mut frame = noise_frame(&mut rng);
            ns.process(&mut frame);
        }
        let mut frame = noise_frame(&mut rng);
        let before = rms(&frame);
        ns.process(&mut frame);
        let after = rms(&frame);
        assert!(
            after < before * 0.8,
            "NS did not suppress noise: before={before:.4} after={after:.4}"
        );
    }

    #[test]
    fn ns_disabled_is_passthrough() {
        let mut ns = NoiseSuppressor::new();
        ns.set_enabled(false);
        let mut frame = [0.1_f32; FRAME_SAMPLES];
        ns.process(&mut frame);
        assert!(frame.iter().all(|&s| (s - 0.1).abs() < 1e-6));
    }

    #[test]
    fn reset_clears_state() {
        let mut ns = NoiseSuppressor::new();
        for _ in 0..30 {
            let mut frame = [0.05_f32; FRAME_SAMPLES];
            ns.process(&mut frame);
        }
        ns.reset();
        assert_eq!(ns.frame_count, 0);
        assert!(ns.ola_tail.iter().all(|&s| s == 0.0));
    }

    fn rms(frame: &[f32]) -> f32 {
        (frame.iter().map(|s| s * s).sum::<f32>() / frame.len() as f32).sqrt()
    }
}
