rosa 0.1.0

use super::convert::{fft_frequencies, mel_frequencies};
use super::transform::{StftParams, spectrogram};
use crate::matrix::Matrix;

/// Mel filterbank normalization mode.
#[derive(Clone, Debug, PartialEq)]
pub enum MelNorm {
    /// No normalization.
    None,
    /// Slaney normalization: divide each filter by its bandwidth.
    Slaney,
}

/// Parameters for mel filterbank construction.
pub struct MelFilterParams {
    /// Sample rate of the audio in Hz.
    pub sr: f64,
    /// Number of FFT points.
    pub n_fft: usize,
    /// Number of mel filter bands.
    pub n_mels: usize,
    /// Lowest frequency in Hz.
    pub fmin: f64,
    /// Highest frequency in Hz. None defaults to `sr / 2`.
    pub fmax: Option<f64>,
    /// If true, use HTK mel scale; otherwise use Slaney/O'Shaughnessy scale.
    pub htk: bool,
    /// Normalization mode applied to each filter.
    pub norm: MelNorm,
}

impl MelFilterParams {
    /// Create `MelFilterParams` with the given sample rate and FFT size and sensible defaults.
    pub fn new(sr: f64, n_fft: usize) -> Self {
        Self {
            sr,
            n_fft,
            n_mels: 128,
            fmin: 0.0,
            fmax: Option::None,
            htk: false,
            norm: MelNorm::Slaney,
        }
    }
}

/// Parameters for mel spectrogram computation.
pub struct MelSpectrogramParams {
    /// STFT parameters.
    pub stft: StftParams,
    /// Mel filterbank parameters.
    pub mel: MelFilterParams,
    /// Exponent for the magnitude spectrogram (1 = magnitude, 2 = power).
    pub power: f64,
}

impl MelSpectrogramParams {
    /// Create `MelSpectrogramParams` for the given sample rate with sensible defaults.
    pub fn new(sr: f64) -> Self {
        Self {
            stft: StftParams::default(),
            mel: MelFilterParams::new(sr, 2048),
            power: 2.0,
        }
    }
}

/// Build a mel filterbank matrix.
///
/// Returns a Matrix of shape `(n_mels, 1 + n_fft/2)`.
///
/// # Examples
///
/// ```
/// use rosa::stft::{mel_filter, MelFilterParams};
///
/// let params = MelFilterParams::new(22050.0, 2048);
/// let fb = mel_filter(&params);
/// assert_eq!(fb.rows(), 128); // default n_mels
/// assert_eq!(fb.cols(), 1025); // 1 + 2048/2
/// ```
pub fn mel_filter(params: &MelFilterParams) -> Matrix {
    let fmax = params.fmax.unwrap_or(params.sr / 2.0);
    let n_bins = 1 + params.n_fft / 2;

    // FFT bin center frequencies
    let fftfreqs = fft_frequencies(params.sr, params.n_fft);

    // Mel center frequencies (n_mels + 2 points including edges)
    let mel_f = mel_frequencies(params.n_mels + 2, params.fmin, fmax, params.htk);

    // Differences between consecutive mel frequencies
    let fdiff: Vec<f64> = mel_f.windows(2).map(|w| w[1] - w[0]).collect();

    // Ramps: mel_f[i] - fftfreqs[j], shape (n_mels+2, n_bins)
    // Build triangular filters
    let n_mels = params.n_mels;
    let mut data = vec![0.0; n_mels * n_bins];

    for i in 0..n_mels {
        for j in 0..n_bins {
            let lower = (fftfreqs[j] - mel_f[i]) / fdiff[i];
            let upper = (mel_f[i + 2] - fftfreqs[j]) / fdiff[i + 1];
            let val = lower.min(upper).max(0.0);
            data[i * n_bins + j] = val;
        }
    }

    // Apply normalization
    if params.norm == MelNorm::Slaney {
        for i in 0..n_mels {
            let enorm = 2.0 / (mel_f[i + 2] - mel_f[i]);
            for j in 0..n_bins {
                data[i * n_bins + j] *= enorm;
            }
        }
    }

    Matrix::from_vec(data, n_mels, n_bins)
}

/// Compute a mel spectrogram.
///
/// Returns a Matrix of shape `(n_mels, n_frames)`.
///
/// # Examples
///
/// ```
/// use rosa::stft::{mel_spectrogram, MelSpectrogramParams};
///
/// let signal = vec![0.0; 22050]; // 1 second at 22050 Hz
/// let params = MelSpectrogramParams::new(22050.0);
/// let mel = mel_spectrogram(&signal, &params);
/// assert_eq!(mel.rows(), 128); // default n_mels
/// ```
pub fn mel_spectrogram(y: &[f64], params: &MelSpectrogramParams) -> Matrix {
    let s = spectrogram(y, &params.stft, params.power);
    let mel_basis = mel_filter(&params.mel);
    mel_basis.matmul(&s)
}

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

    #[derive(Deserialize)]
    struct MelFilterTestCase {
        name: String,
        sr: f64,
        n_fft: usize,
        n_mels: usize,
        fmin: f64,
        fmax: f64,
        htk: bool,
        norm: Option<String>,
        rows: usize,
        cols: usize,
        data: Vec<f64>,
    }

    fn load_mel_refs() -> Vec<MelFilterTestCase> {
        let data = include_str!("tests/references/mel_filterbank.json");
        serde_json::from_str(data).expect("Failed to parse mel_filterbank.json")
    }

    #[test]
    fn test_mel_filter_matches_librosa() {
        let refs = load_mel_refs();
        assert!(!refs.is_empty(), "No mel filterbank test cases found");

        for case in &refs {
            let norm = match case.norm.as_deref() {
                Some("slaney") => MelNorm::Slaney,
                _ => MelNorm::None,
            };
            let params = MelFilterParams {
                sr: case.sr,
                n_fft: case.n_fft,
                n_mels: case.n_mels,
                fmin: case.fmin,
                fmax: Some(case.fmax),
                htk: case.htk,
                norm,
            };

            let result = mel_filter(&params);
            assert_eq!(
                result.shape(),
                (case.rows, case.cols),
                "{}: shape mismatch",
                case.name
            );

            let actual = result.as_slice();
            assert_eq!(
                actual.len(),
                case.data.len(),
                "{}: data length mismatch",
                case.name
            );

            let tol = 1e-7;
            for (i, (a, e)) in actual.iter().zip(case.data.iter()).enumerate() {
                let diff = (a - e).abs();
                assert!(
                    diff < tol,
                    "{}[{}]: got {}, expected {}, diff {} > tol {}",
                    case.name,
                    i,
                    a,
                    e,
                    diff,
                    tol
                );
            }
        }
    }
}