audiofp 0.3.7

//! 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,

    /// Row-major `(n_mels, n_fft/2 + 1)` weight matrix.
    matrix: Vec<f32>,

    /// 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 hot-path
    /// representation; `matrix` is kept for 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 {
        assert!(n_mels > 0, "n_mels must be > 0");
        assert!(
            n_fft >= 2 && n_fft.is_multiple_of(2),
            "n_fft must be even and >= 2"
        );
        assert!(fmin >= 0.0, "fmin must be >= 0");
        assert!(fmin < fmax, "fmin must be strictly less than fmax");

        let n_bins = n_fft / 2 + 1;
        let mut matrix = vec![0.0_f32; n_mels * n_bins];

        // 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;

        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);

            let row = &mut matrix[k * n_bins..(k + 1) * n_bins];
            for (b, w) in row.iter_mut().enumerate() {
                let f = b as f32 * bin_hz;
                *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)
                };
            }
        }

        // Build CSR (sparse) representation: for each band, find the
        // contiguous range of non-zero bins and store only those weights.
        // Each triangular filter is non-zero only in (left_hz, right_hz),
        // so the sparse representation skips all zero-tail bins.
        let mut sparse = Vec::with_capacity(n_mels);
        for k in 0..n_mels {
            let row = &matrix[k * n_bins..(k + 1) * n_bins];
            // Find first and last non-zero bin.
            let first = row.iter().position(|&w| w != 0.0).unwrap_or(n_bins);
            let last = row.iter().rposition(|&w| w != 0.0).unwrap_or(0);
            if first <= last {
                sparse.push(MelBand {
                    start_bin: first,
                    weights: row[first..=last].to_vec(),
                });
            } else {
                sparse.push(MelBand {
                    start_bin: 0,
                    weights: Vec::new(),
                });
            }
        }

        Self {
            n_mels,
            n_fft,
            sr,
            fmin,
            fmax,
            scale,
            matrix,
            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
    }

    /// Borrow the row-major weight matrix.
    #[must_use]
    pub fn matrix(&self) -> &[f32] {
        &self.matrix
    }

    /// 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.
        for (k, slot) in out.iter_mut().enumerate() {
            let band = &self.sparse[k];
            let mut acc = 0.0_f32;
            for (w, m) in band.weights.iter().zip(magnitude[band.start_bin..].iter()) {
                acc += w * (m * m);
            }
            *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.
        for (k, slot) in out.iter_mut().enumerate() {
            let band = &self.sparse[k];
            let mut acc = 0.0_f32;
            for (w, p) in band.weights.iter().zip(power[band.start_bin..].iter()) {
                acc += w * p;
            }
            *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();
        for k in 0..fb.n_mels {
            let row = &fb.matrix[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);
    }
}