rosa 0.1.0

//! Wavelet filter construction for CQT/VQT.

use crate::matrix::Matrix;

/// Parameters for computing wavelet filter lengths.
pub struct WaveletLengthsParams {
    /// Sample rate in Hz.
    pub sr: f64,
    /// Window function used to shape each filter.
    pub window: fn(usize, bool) -> Vec<f64>,
    /// Scaling factor applied to filter lengths.
    pub filter_scale: f64,
    /// None = auto (ERB-based for VQT), Some(0.0) = CQT, Some(v) = custom VQT offset.
    pub gamma: Option<f64>,
}

impl Default for WaveletLengthsParams {
    fn default() -> Self {
        Self {
            sr: 22050.0,
            window: crate::windows::hann,
            filter_scale: 1.0,
            gamma: Some(0.0),
        }
    }
}

/// Parameters for building wavelet filters.
pub struct WaveletParams {
    /// Sample rate in Hz.
    pub sr: f64,
    /// Window function used to shape each filter.
    pub window: fn(usize, bool) -> Vec<f64>,
    /// Scaling factor applied to filter lengths.
    pub filter_scale: f64,
    /// Whether to zero-pad filters to the next power of two.
    pub pad_fft: bool,
    /// None = no normalization, Some(1.0) = L1, Some(2.0) = L2.
    pub norm: Option<f64>,
    /// None = auto (ERB), Some(0.0) = CQT.
    pub gamma: Option<f64>,
}

impl Default for WaveletParams {
    fn default() -> Self {
        Self {
            sr: 22050.0,
            window: crate::windows::hann,
            filter_scale: 1.0,
            pad_fft: true,
            norm: Some(1.0),
            gamma: Some(0.0),
        }
    }
}

/// Equivalent noise bandwidth (ENBW) of a window function.
///
/// Matches `librosa.filters.window_bandwidth(window, n)`.
/// Formula: `enbw = n * sum(w^2) / sum(w)^2`
pub fn window_bandwidth(window_fn: fn(usize, bool) -> Vec<f64>, n: usize) -> f64 {
    let w = window_fn(n, true);
    let sum_sq: f64 = w.iter().map(|x| x * x).sum();
    let sum_w: f64 = w.iter().sum();
    n as f64 * sum_sq / (sum_w * sum_w)
}

/// Compute relative bandwidth alpha from frequency array.
///
/// Matches `librosa.filters._relative_bandwidth(freqs=freqs)`.
fn relative_bandwidth(freqs: &[f64]) -> Vec<f64> {
    let n = freqs.len();
    assert!(n >= 2, "relative_bandwidth requires at least 2 frequencies");

    let logf: Vec<f64> = freqs.iter().map(|f| f.log2()).collect();
    let mut bpo = vec![0.0; n];

    // Boundary: first and last use one-sided differences
    bpo[0] = 1.0 / (logf[1] - logf[0]);
    bpo[n - 1] = 1.0 / (logf[n - 1] - logf[n - 2]);

    // Interior: centered differences
    for i in 1..n - 1 {
        bpo[i] = 2.0 / (logf[i + 1] - logf[i - 1]);
    }

    // Convert local bins-per-octave to relative bandwidth
    bpo.iter()
        .map(|&b| {
            let r2 = 2.0_f64.powf(2.0 / b);
            (r2 - 1.0) / (r2 + 1.0)
        })
        .collect()
}

/// Relative bandwidth for equal-tempered spacing.
fn _et_relative_bandwidth(bins_per_octave: usize) -> f64 {
    let r = 2.0_f64.powf(1.0 / bins_per_octave as f64);
    (r * r - 1.0) / (r * r + 1.0)
}

/// Compute filter lengths for CQT/VQT wavelet basis.
///
/// Matches `librosa.filters.wavelet_lengths(freqs=..., sr=..., ...)`.
/// Returns (lengths, cutoff_frequency).
pub fn wavelet_lengths(freqs: &[f64], params: &WaveletLengthsParams) -> (Vec<f64>, f64) {
    let alpha = relative_bandwidth(freqs);
    let n = freqs.len();

    // Resolve gamma for each bin
    let gamma_vals: Vec<f64> = match params.gamma {
        None => {
            // Auto ERB: gamma = alpha * 24.7 / 0.108
            alpha.iter().map(|&a| a * 24.7 / 0.108).collect()
        }
        Some(g) => vec![g; n],
    };

    // Q factor per bin
    let q: Vec<f64> = alpha.iter().map(|&a| params.filter_scale / a).collect();

    // Filter lengths
    let lengths: Vec<f64> = (0..n)
        .map(|i| q[i] * params.sr / (freqs[i] + gamma_vals[i] / alpha[i]))
        .collect();

    // Cutoff frequency
    let win_bw = window_bandwidth(params.window, 1000);
    let f_cutoff = (0..n)
        .map(|i| freqs[i] * (1.0 + 0.5 * win_bw / q[i]) + 0.5 * gamma_vals[i])
        .fold(0.0_f64, f64::max);

    (lengths, f_cutoff)
}

/// Build complex wavelet filters for CQT/VQT.
///
/// Matches `librosa.filters.wavelet(freqs=..., sr=..., ...)`.
/// Returns (filters_re, filters_im, lengths).
/// filters_re/im are `Matrix(n_freqs, max_len)`, lengths is `Vec<f64>`.
pub fn wavelet(freqs: &[f64], params: &WaveletParams) -> (Matrix, Matrix, Vec<f64>) {
    let lengths_params = WaveletLengthsParams {
        sr: params.sr,
        window: params.window,
        filter_scale: params.filter_scale,
        gamma: params.gamma,
    };
    let (lengths, _f_cutoff) = wavelet_lengths(freqs, &lengths_params);

    let n_freqs = freqs.len();

    // Build each filter
    let mut filters_re: Vec<Vec<f64>> = Vec::with_capacity(n_freqs);
    let mut filters_im: Vec<Vec<f64>> = Vec::with_capacity(n_freqs);

    for (k, &freq) in freqs.iter().enumerate() {
        let ilen = lengths[k];
        // Match librosa: n_k = int(-ilen//2) to int(ilen//2) which gives ceil(ilen) samples
        let half = (ilen / 2.0).floor() as i64;
        let n_k = (half + half) as usize; // This gives 2*floor(ilen/2), matching arange(-ilen//2, ilen//2)
        let start = -half;

        // Generate time axis and complex exponential
        let mut re = Vec::with_capacity(n_k);
        let mut im = Vec::with_capacity(n_k);
        for j in 0..n_k {
            let t = (start + j as i64) as f64;
            let phase = 2.0 * std::f64::consts::PI * freq * t / params.sr;
            re.push(phase.cos());
            im.push(phase.sin());
        }

        // Apply window
        let win = (params.window)(n_k, true);
        for j in 0..n_k {
            re[j] *= win[j];
            im[j] *= win[j];
        }

        // Normalize complex signal
        if let Some(p) = params.norm {
            let norm_val = if p == 1.0 {
                // L1: sum of |z|
                re.iter()
                    .zip(im.iter())
                    .map(|(&r, &i)| (r * r + i * i).sqrt())
                    .sum::<f64>()
            } else if p == 2.0 {
                // L2: sqrt(sum(|z|^2))
                re.iter()
                    .zip(im.iter())
                    .map(|(&r, &i)| r * r + i * i)
                    .sum::<f64>()
                    .sqrt()
            } else if p.is_infinite() {
                // Linf: max(|z|)
                re.iter()
                    .zip(im.iter())
                    .map(|(&r, &i)| (r * r + i * i).sqrt())
                    .fold(0.0_f64, f64::max)
            } else {
                // Lp norm
                re.iter()
                    .zip(im.iter())
                    .map(|(&r, &i)| (r * r + i * i).sqrt().powf(p))
                    .sum::<f64>()
                    .powf(1.0 / p)
            };

            if norm_val > f64::MIN_POSITIVE {
                let inv = 1.0 / norm_val;
                for j in 0..n_k {
                    re[j] *= inv;
                    im[j] *= inv;
                }
            }
        }

        filters_re.push(re);
        filters_im.push(im);
    }

    // Determine max_len
    let raw_max = filters_re.iter().map(|f| f.len()).max().unwrap_or(0);
    let max_len = if params.pad_fft {
        raw_max.next_power_of_two()
    } else {
        raw_max
    };

    // Center-pad all filters to max_len
    let mut re_data = vec![0.0; n_freqs * max_len];
    let mut im_data = vec![0.0; n_freqs * max_len];

    for k in 0..n_freqs {
        let flen = filters_re[k].len();
        let lpad = (max_len - flen) / 2;
        for j in 0..flen {
            re_data[k * max_len + lpad + j] = filters_re[k][j];
            im_data[k * max_len + lpad + j] = filters_im[k][j];
        }
    }

    (
        Matrix::from_vec(re_data, n_freqs, max_len),
        Matrix::from_vec(im_data, n_freqs, max_len),
        lengths,
    )
}

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

    #[test]
    fn test_window_bandwidth_hann() {
        let bw = window_bandwidth(crate::windows::hann, 1000);
        // Hann window ENBW should be approximately 1.5
        assert!((bw - 1.5).abs() < 0.01, "Hann ENBW expected ~1.5, got {bw}");
    }

    #[test]
    fn test_cqt_frequencies_basic() {
        let freqs = crate::convert::cqt_frequencies(84, 32.703195662574829, 12, 0.0);
        assert_eq!(freqs.len(), 84);
        // First freq should be C1
        assert!((freqs[0] - 32.703195662574829).abs() < 1e-10);
        // 12 bins up should be exactly 2x (one octave)
        assert!((freqs[12] / freqs[0] - 2.0).abs() < 1e-10);
    }

    #[test]
    fn test_wavelet_lengths_cqt() {
        let freqs = crate::convert::cqt_frequencies(84, 32.703195662574829, 12, 0.0);
        let params = WaveletLengthsParams {
            sr: 22050.0,
            window: crate::windows::hann,
            filter_scale: 1.0,
            gamma: Some(0.0),
        };
        let (lengths, f_cutoff) = wavelet_lengths(&freqs, &params);
        assert_eq!(lengths.len(), 84);
        // Lowest freq should have longest filter
        assert!(lengths[0] > lengths[83]);
        // Cutoff should be positive
        assert!(f_cutoff > 0.0);
        // Cutoff should not exceed Nyquist (22050/2)
        assert!(f_cutoff < 22050.0);
    }

    #[test]
    fn test_wavelet_construction() {
        let freqs = crate::convert::cqt_frequencies(12, 32.703195662574829, 12, 0.0);
        let params = WaveletParams::default();
        let (re, im, lengths) = wavelet(&freqs, &params);
        assert_eq!(re.rows(), 12);
        assert_eq!(im.rows(), 12);
        assert_eq!(re.cols(), im.cols());
        assert_eq!(lengths.len(), 12);
        // All lengths should be positive
        for &l in &lengths {
            assert!(l > 0.0);
        }
    }
}