audiofp 0.3.9

//! Mel filterbank: triangular filters spaced on the perceptual mel scale.
//!
//! [`MelFilterBank`] holds an `(n_mels, n_fft/2 + 1)` matrix of triangle
//! weights; [`MelFilterBank::log_mel`] dots a magnitude spectrum (squared
//! to power) into one log-mel frame.
//!
//! Filters are slaney-normalised — each triangle has unit area in the
//! linear-frequency domain — so log-mel output magnitudes are stable
//! across `n_mels` choices and match `librosa.feature.melspectrogram`'s
//! defaults.

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

use libm::{expf, log2f, logf, powf};

/// Selects how hertz are mapped to the mel scale.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
pub enum MelScale {
    /// HTK formula `mel = 2595 · log10(1 + hz/700)`. Single closed-form
    /// expression, slightly different from Slaney above 1 kHz.
    Htk,

    /// Slaney's auditory-toolbox mapping: linear below 1 kHz, log above.
    /// This is `librosa`'s default and the right choice for most music
    /// applications.
    Slaney,
}

const SLANEY_F_SP: f32 = 200.0 / 3.0;
const SLANEY_MIN_LOG_HZ: f32 = 1000.0;
/// log(6.4) / 27 ≈ 0.068751965 — precomputed once instead of calling
/// `libm::logf` + division every time `hz_to_mel` / `mel_to_hz` is called.
const SLANEY_LOGSTEP: f32 = 0.068_751_97_f32;
// log2 is cheaper than log10 on libm softfloat, so we multiply
// `core::f32::consts::LOG10_2 * log2f(x)` instead of calling `log10f`.
const SLANEY_MIN_LOG_MEL: f32 = SLANEY_MIN_LOG_HZ / SLANEY_F_SP;

impl MelScale {
    #[inline]
    fn hz_to_mel(self, hz: f32) -> f32 {
        match self {
            MelScale::Htk => 2595.0 * core::f32::consts::LOG10_2 * log2f(1.0 + hz / 700.0),
            MelScale::Slaney => {
                if hz < SLANEY_MIN_LOG_HZ {
                    hz / SLANEY_F_SP
                } else {
                    SLANEY_MIN_LOG_MEL + logf(hz / SLANEY_MIN_LOG_HZ) / SLANEY_LOGSTEP
                }
            }
        }
    }

    #[inline]
    fn mel_to_hz(self, mel: f32) -> f32 {
        match self {
            MelScale::Htk => 700.0 * (powf(10.0, mel / 2595.0) - 1.0),
            MelScale::Slaney => {
                if mel < SLANEY_MIN_LOG_MEL {
                    SLANEY_F_SP * mel
                } else {
                    SLANEY_MIN_LOG_HZ * expf(SLANEY_LOGSTEP * (mel - SLANEY_MIN_LOG_MEL))
                }
            }
        }
    }
}

/// One mel band in CSR (Compressed Sparse Row) form: the triangle
/// weights for bins `[start_bin .. start_bin + weights.len()]`.
/// All other bins have weight 0 and are skipped in the hot path.
#[derive(Clone, Debug)]
struct MelBand {
    start_bin: usize,
    weights: Vec<f32>,
}

/// A precomputed triangular mel filterbank.
///
/// # Example
///
/// ```
/// use audiofp::dsp::mel::{MelFilterBank, MelScale};
///
/// // 128 mels covering 0–11025 Hz at sr=22050, n_fft=2048.
/// let fb = MelFilterBank::new(128, 2048, 22_050, 0.0, 11_025.0, MelScale::Slaney);
/// assert_eq!(fb.n_mels, 128);
/// assert_eq!(fb.n_bins(), 1025);
/// ```
#[derive(Clone, Debug)]
pub struct MelFilterBank {
    /// Number of mel bands (rows of the matrix).
    pub n_mels: usize,
    /// FFT length the upstream STFT uses; bin count is `n_fft / 2 + 1`.
    pub n_fft: usize,
    /// Sample rate of the audio fed to the upstream STFT.
    pub sr: u32,
    /// Lowest frequency (Hz) covered by the filterbank.
    pub fmin: f32,
    /// Highest frequency (Hz) covered by the filterbank.
    pub fmax: f32,
    /// Mel scale convention used to lay out filter centres.
    pub scale: MelScale,

    /// Sparse (CSR) representation of each mel band. Only stores the
    /// non-zero weight range per band so `log_mel_from_power` iterates
    /// ~20-40 bins instead of all `n_bins` (513+). This is the sole
    /// hot-path representation; the dense matrix is reconstructed
    /// on-demand via the `matrix()` getter.
    sparse: Vec<MelBand>,
}

impl MelFilterBank {
    /// Build a filterbank.
    ///
    /// # Panics
    ///
    /// Panics if `n_mels == 0`, `n_fft < 2`, `n_fft` is not even, or
    /// `fmin >= fmax`.
    ///
    /// `fmin = 0` is accepted — both the Slaney and HTK mel scales
    /// handle 0 Hz without hitting `log(0)` (Slaney's linear branch
    /// covers `hz < 1000`, HTK's formula evaluates `log10(1 + 0) = 0`).
    /// The first filter simply starts at 0 Hz. Note that
    /// [`HaitsmaConfig`](crate::classical::HaitsmaConfig) independently
    /// requires `fmin > 0` because its log-spaced band edges use
    /// `powf(fmax / fmin, …)`, which is undefined for `fmin = 0`; that
    /// restriction is Haitsma-specific and does not apply here.
    #[must_use]
    pub fn new(
        n_mels: usize,
        n_fft: usize,
        sr: u32,
        fmin: f32,
        fmax: f32,
        scale: MelScale,
    ) -> Self {
        Self::try_new(n_mels, n_fft, sr, fmin, fmax, scale).expect("invalid MelFilterBank config")
    }

    /// Fallible constructor — returns [`AfpError::Config`](crate::AfpError::Config) on invalid
    /// parameters instead of panicking.
    ///
    /// # Errors
    ///
    /// - `n_mels == 0`
    /// - `n_fft < 2` or odd
    /// - `fmin < 0`
    /// - `fmin >= fmax`
    pub fn try_new(
        n_mels: usize,
        n_fft: usize,
        sr: u32,
        fmin: f32,
        fmax: f32,
        scale: MelScale,
    ) -> crate::Result<Self> {
        if n_mels == 0 {
            return Err(crate::AfpError::Config("n_mels must be > 0".into()));
        }
        if n_fft < 2 || !n_fft.is_multiple_of(2) {
            return Err(crate::AfpError::Config(
                "n_fft must be even and >= 2".into(),
            ));
        }
        if fmin < 0.0 || fmin.is_nan() {
            return Err(crate::AfpError::Config("fmin must be >= 0".into()));
        }
        if fmin >= fmax || fmax.is_nan() {
            return Err(crate::AfpError::Config(
                "fmin must be strictly less than fmax".into(),
            ));
        }

        let n_bins = n_fft / 2 + 1;

        // Mel-spaced centre points, including the left and right "skirts".
        let mel_min = scale.hz_to_mel(fmin);
        let mel_max = scale.hz_to_mel(fmax);
        let n_points = n_mels + 2;
        let mut hz_points = Vec::with_capacity(n_points);
        for k in 0..n_points {
            let mel = mel_min + (mel_max - mel_min) * k as f32 / (n_points - 1) as f32;
            hz_points.push(scale.mel_to_hz(mel));
        }

        // FFT bin frequencies in Hz: bin b corresponds to b * sr / n_fft.
        let bin_hz = sr as f32 / n_fft as f32;

        // Build CSR (sparse) representation directly: for each band,
        // analytically determine the non-zero bin range and compute only
        // those weights. This avoids allocating the full dense matrix
        // (n_mels × n_bins, e.g. 512 KB for typical configs).
        let mut sparse = Vec::with_capacity(n_mels);
        for k in 0..n_mels {
            let left = hz_points[k];
            let centre = hz_points[k + 1];
            let right = hz_points[k + 2];
            // Slaney normalisation: unit area in linear frequency.
            let norm = 2.0 / (right - left).max(1e-10);

            // Analytical bin range where the triangle is non-zero:
            // weight > 0 iff left < f < right (strictly), with f = b * bin_hz.
            //
            // first_bin: smallest b such that b * bin_hz > left
            //   → b > left / bin_hz → b = floor(left / bin_hz) + 1
            //   (clamped to [0, n_bins) to avoid out-of-bounds)
            //
            // last_bin: largest b such that b * bin_hz < right
            //   → b < right / bin_hz → b = ceil(right / bin_hz) - 1
            //   (clamped to n_bins - 1)
            let first_bin = ((left / bin_hz).floor() as usize + 1).min(n_bins);
            let last_bin_raw = (right / bin_hz).ceil() as usize;
            let last_bin = if last_bin_raw == 0 {
                0
            } else {
                (last_bin_raw - 1).min(n_bins - 1)
            };

            if first_bin <= last_bin && first_bin < n_bins {
                let mut weights = Vec::with_capacity(last_bin - first_bin + 1);
                for b in first_bin..=last_bin {
                    let f = b as f32 * bin_hz;
                    // Guard: floating-point edge cases where the analytical
                    // bound slightly overshoots. The original code would
                    // produce 0.0 for f <= left or f >= right.
                    let w = if f <= left || f >= right {
                        0.0
                    } else if f <= centre {
                        norm * (f - left) / (centre - left).max(1e-10)
                    } else {
                        norm * (right - f) / (right - centre).max(1e-10)
                    };
                    weights.push(w);
                }
                // Trim leading/trailing zeros from float edge cases to
                // match the original scan-based CSR construction exactly.
                let first_nz = weights.iter().position(|&w| w != 0.0);
                let last_nz = weights.iter().rposition(|&w| w != 0.0);
                match (first_nz, last_nz) {
                    (Some(f), Some(l)) => {
                        sparse.push(MelBand {
                            start_bin: first_bin + f,
                            weights: weights[f..=l].to_vec(),
                        });
                    }
                    _ => {
                        sparse.push(MelBand {
                            start_bin: 0,
                            weights: Vec::new(),
                        });
                    }
                }
            } else {
                sparse.push(MelBand {
                    start_bin: 0,
                    weights: Vec::new(),
                });
            }
        }

        Ok(Self {
            n_mels,
            n_fft,
            sr,
            fmin,
            fmax,
            scale,
            sparse,
        })
    }

    /// Number of FFT bins each filter spans (`n_fft / 2 + 1`).
    #[must_use]
    pub const fn n_bins(&self) -> usize {
        self.n_fft / 2 + 1
    }

    /// Reconstruct the row-major weight matrix from the sparse
    /// representation. Allocates a new `Vec<f32>` on each call.
    ///
    /// This method is provided for inspection/debugging; the hot path
    /// uses the sparse CSR representation directly.
    #[must_use]
    pub fn matrix(&self) -> Vec<f32> {
        let n_bins = self.n_bins();
        let mut mat = vec![0.0_f32; self.n_mels * n_bins];
        for (k, band) in self.sparse.iter().enumerate() {
            for (j, &w) in band.weights.iter().enumerate() {
                mat[k * n_bins + band.start_bin + j] = w;
            }
        }
        mat
    }

    /// Compute one log-mel frame from a magnitude spectrum.
    ///
    /// Computes `log10(M · |X|² + 1e-10)` per librosa: the magnitude is
    /// squared to power before the matrix-vector product, and a small
    /// floor avoids `log10(0)`.
    ///
    /// # Panics
    ///
    /// Panics if `magnitude.len() != n_bins()` or `out.len() != n_mels`.
    pub fn log_mel(&self, magnitude: &[f32], out: &mut [f32]) {
        assert_eq!(
            magnitude.len(),
            self.n_bins(),
            "magnitude length must equal n_bins"
        );
        assert_eq!(out.len(), self.n_mels, "out length must equal n_mels");

        // Use the sparse representation: only iterate non-zero bins per band.
        // Square magnitude to power inline while accumulating via SIMD.
        for (k, slot) in out.iter_mut().enumerate() {
            let band = &self.sparse[k];
            let acc = super::dot_sq_wide(
                &band.weights,
                &magnitude[band.start_bin..band.start_bin + band.weights.len()],
            );
            *slot = core::f32::consts::LOG10_2 * log2f(acc + 1e-10);
        }
    }

    /// Compute one log-mel frame from a **power** spectrum
    /// (`re² + im²` per bin, e.g. one row of
    /// [`ShortTimeFFT::power_flat`]).
    ///
    /// Equivalent to [`log_mel`] but skips the per-bin square — feed the
    /// output of `power_flat` / `process_frame_power` directly to avoid
    /// doing the work twice.
    ///
    /// [`log_mel`]: MelFilterBank::log_mel
    /// [`ShortTimeFFT::power_flat`]: crate::dsp::stft::ShortTimeFFT::power_flat
    ///
    /// # Panics
    ///
    /// Panics if `power.len() != n_bins()` or `out.len() != n_mels`.
    pub fn log_mel_from_power(&self, power: &[f32], out: &mut [f32]) {
        assert_eq!(power.len(), self.n_bins(), "power length must equal n_bins");
        assert_eq!(out.len(), self.n_mels, "out length must equal n_mels");

        // Use the sparse representation: only iterate non-zero bins per band.
        // Each triangular filter spans ~20-40 bins instead of all n_bins.
        // The dot product is vectorized 8-wide via `wide::f32x8`.
        for (k, slot) in out.iter_mut().enumerate() {
            let band = &self.sparse[k];
            let acc = super::dot_wide(
                &band.weights,
                &power[band.start_bin..band.start_bin + band.weights.len()],
            );
            *slot = core::f32::consts::LOG10_2 * log2f(acc + 1e-10);
        }
    }
}

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

    #[test]
    fn htk_round_trip() {
        for &hz in &[0.0_f32, 100.0, 440.0, 1_000.0, 5_000.0, 11_025.0] {
            let m = MelScale::Htk.hz_to_mel(hz);
            assert_relative_eq!(MelScale::Htk.mel_to_hz(m), hz, max_relative = 1e-5);
        }
    }

    #[test]
    fn slaney_round_trip() {
        for &hz in &[
            0.0_f32, 100.0, 440.0, 999.0, 1_000.0, 1_001.0, 5_000.0, 11_025.0,
        ] {
            let m = MelScale::Slaney.hz_to_mel(hz);
            assert_relative_eq!(MelScale::Slaney.mel_to_hz(m), hz, max_relative = 1e-4);
        }
    }

    #[test]
    fn matrix_dimensions() {
        let fb = MelFilterBank::new(64, 1024, 16_000, 0.0, 8_000.0, MelScale::Htk);
        assert_eq!(fb.n_bins(), 513);
        assert_eq!(fb.matrix().len(), 64 * 513);
    }

    #[test]
    fn each_filter_has_a_peak_in_band() {
        let fb = MelFilterBank::new(40, 2048, 22_050, 0.0, 11_025.0, MelScale::Slaney);
        let n_bins = fb.n_bins();
        let mat = fb.matrix();
        for k in 0..fb.n_mels {
            let row = &mat[k * n_bins..(k + 1) * n_bins];
            let max = row.iter().cloned().fold(0.0_f32, f32::max);
            assert!(max > 0.0, "filter {k} is all-zero");
        }
    }

    #[test]
    fn log_mel_floor_at_silence() {
        let fb = MelFilterBank::new(16, 512, 16_000, 0.0, 8_000.0, MelScale::Htk);
        let zeros = vec![0.0_f32; fb.n_bins()];
        let mut out = vec![0.0_f32; fb.n_mels];
        fb.log_mel(&zeros, &mut out);
        // log10(1e-10) = -10.0 exactly.
        for v in out {
            assert_relative_eq!(v, -10.0, max_relative = 1e-5);
        }
    }

    #[test]
    fn htk_and_slaney_diverge_above_1khz() {
        // Below 1 kHz the two scales should agree to within ~5 mel.
        // Above 1 kHz Slaney is logarithmic with a different slope, so the
        // converted mel values diverge.
        let lo = 500.0_f32;
        let hi = 4_000.0_f32;
        let m_htk_lo = MelScale::Htk.hz_to_mel(lo);
        let m_sla_lo = MelScale::Slaney.hz_to_mel(lo);
        let m_htk_hi = MelScale::Htk.hz_to_mel(hi);
        let m_sla_hi = MelScale::Slaney.hz_to_mel(hi);

        let diff_lo = (m_htk_lo - m_sla_lo).abs();
        let diff_hi = (m_htk_hi - m_sla_hi).abs();
        assert!(
            diff_hi > diff_lo,
            "expected divergence to grow above 1 kHz: lo={diff_lo} hi={diff_hi}",
        );
    }

    #[test]
    fn matrix_rows_are_non_negative() {
        let fb = MelFilterBank::new(64, 2048, 22_050, 0.0, 11_025.0, MelScale::Slaney);
        for w in fb.matrix() {
            assert!(w >= 0.0, "negative weight in mel matrix: {w}");
        }
    }

    #[test]
    fn log_mel_from_power_matches_log_mel_on_squared_input() {
        let fb = MelFilterBank::new(32, 1024, 16_000, 0.0, 8_000.0, MelScale::Slaney);
        let n_bins = fb.n_bins();

        // Synthetic spiky magnitude spectrum.
        let mag: Vec<f32> = (0..n_bins)
            .map(|b| ((b as f32 * 0.073).sin().abs() + 0.001) * (1 + b % 7) as f32)
            .collect();
        let pow: Vec<f32> = mag.iter().map(|m| m * m).collect();

        let mut out_mag = vec![0.0_f32; fb.n_mels];
        let mut out_pow = vec![0.0_f32; fb.n_mels];
        fb.log_mel(&mag, &mut out_mag);
        fb.log_mel_from_power(&pow, &mut out_pow);

        for (a, b) in out_mag.iter().zip(out_pow.iter()) {
            assert_relative_eq!(*a, *b, max_relative = 1e-6);
        }
    }

    #[test]
    fn log_mel_picks_up_dirac_in_band() {
        let fb = MelFilterBank::new(40, 2048, 22_050, 0.0, 11_025.0, MelScale::Slaney);
        // Dirac at bin 200 ≈ 200 * 22050/2048 ≈ 2154 Hz.
        let mut mag = vec![0.0_f32; fb.n_bins()];
        mag[200] = 1.0;
        let mut out = vec![0.0_f32; fb.n_mels];
        fb.log_mel(&mag, &mut out);

        // Some band must respond above the silence floor.
        let max = out.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        assert!(max > -9.0, "no band responded: max={max}");
    }

    // -----------------------------------------------------------------
    // Constructor panic coverage.
    //
    // Pin the panic-asserts documented on `MelFilterBank::new`.
    // A refactor that loosens or removes any of these assertions would
    // silently change the filter bank's behaviour on bad input. Each
    // `should_panic` test below pins exactly one assertion.
    // -----------------------------------------------------------------

    #[test]
    #[should_panic(expected = "n_mels must be > 0")]
    fn mel_filter_bank_panics_on_zero_n_mels() {
        let _ = MelFilterBank::new(0, 1024, 16_000, 0.0, 8_000.0, MelScale::Slaney);
    }

    #[test]
    #[should_panic(expected = "n_fft must be even and >= 2")]
    fn mel_filter_bank_panics_on_odd_n_fft() {
        let _ = MelFilterBank::new(64, 1023, 16_000, 0.0, 8_000.0, MelScale::Slaney);
    }

    #[test]
    #[should_panic(expected = "n_fft must be even and >= 2")]
    fn mel_filter_bank_panics_on_n_fft_below_two() {
        let _ = MelFilterBank::new(64, 1, 16_000, 0.0, 8_000.0, MelScale::Slaney);
    }

    #[test]
    #[should_panic(expected = "fmin must be strictly less than fmax")]
    fn mel_filter_bank_panics_when_fmin_equals_fmax() {
        let _ = MelFilterBank::new(64, 1024, 16_000, 1_000.0, 1_000.0, MelScale::Slaney);
    }

    #[test]
    #[should_panic(expected = "fmin must be strictly less than fmax")]
    fn mel_filter_bank_panics_when_fmin_above_fmax() {
        let _ = MelFilterBank::new(64, 1024, 16_000, 4_000.0, 1_000.0, MelScale::Slaney);
    }

    #[test]
    #[should_panic(expected = "fmin must be >= 0")]
    fn mel_filter_bank_panics_on_negative_fmin() {
        let _ = MelFilterBank::new(64, 1024, 16_000, -10.0, 8_000.0, MelScale::Slaney);
    }
}