rosa 0.1.0

//! VQT and CQT core implementation.

use crate::matrix::Matrix;
use realfft::RealFftPlanner;

use super::wavelet::{WaveletLengthsParams, WaveletParams, wavelet, wavelet_lengths};

/// VQT parameters.
pub struct VqtParams {
    /// Sample rate in Hz.
    pub sr: f64,
    /// Number of samples between successive frames.
    pub hop_length: usize,
    /// Minimum frequency in Hz.
    pub fmin: f64,
    /// Total number of frequency bins.
    pub n_bins: usize,
    /// Number of bins per octave.
    pub bins_per_octave: usize,
    /// Tuning offset in fractions of a bin.
    pub tuning: f64,
    /// Scaling factor applied to filter lengths.
    pub filter_scale: f64,
    /// None = no norm, Some(1.0) = L1 (default), Some(2.0) = L2.
    pub norm: Option<f64>,
    /// Window function used to shape each filter.
    pub window: fn(usize, bool) -> Vec<f64>,
    /// Whether to scale each bin by the square root of its filter length.
    pub scale: bool,
    /// None = auto (ERB for VQT), Some(0.0) = CQT.
    pub gamma: Option<f64>,
}

impl Default for VqtParams {
    fn default() -> Self {
        Self {
            sr: 22050.0,
            hop_length: 512,
            fmin: 32.703_195_662_574_83, // C1
            n_bins: 84,
            bins_per_octave: 12,
            tuning: 0.0,
            filter_scale: 1.0,
            norm: Some(1.0),
            window: crate::windows::hann,
            scale: true,
            gamma: None, // auto-ERB for VQT
        }
    }
}

/// CQT parameters (same as VQT but gamma is always 0).
pub struct CqtParams {
    /// Sample rate in Hz.
    pub sr: f64,
    /// Number of samples between successive frames.
    pub hop_length: usize,
    /// Minimum frequency in Hz.
    pub fmin: f64,
    /// Total number of frequency bins.
    pub n_bins: usize,
    /// Number of bins per octave.
    pub bins_per_octave: usize,
    /// Tuning offset in fractions of a bin.
    pub tuning: f64,
    /// Scaling factor applied to filter lengths.
    pub filter_scale: f64,
    /// None = no norm, Some(1.0) = L1 (default), Some(2.0) = L2.
    pub norm: Option<f64>,
    /// Window function used to shape each filter.
    pub window: fn(usize, bool) -> Vec<f64>,
    /// Whether to scale each bin by the square root of its filter length.
    pub scale: bool,
}

impl Default for CqtParams {
    fn default() -> Self {
        Self {
            sr: 22050.0,
            hop_length: 512,
            fmin: 32.703_195_662_574_83, // C1
            n_bins: 84,
            bins_per_octave: 12,
            tuning: 0.0,
            filter_scale: 1.0,
            norm: Some(1.0),
            window: crate::windows::hann,
            scale: true,
        }
    }
}

/// Complex matrix multiply: (A_re + j*A_im) @ (B_re + j*B_im)
/// Returns (C_re, C_im).
pub(crate) fn complex_matmul(
    a_re: &Matrix,
    a_im: &Matrix,
    b_re: &Matrix,
    b_im: &Matrix,
) -> (Matrix, Matrix) {
    // C_re = A_re @ B_re - A_im @ B_im
    // C_im = A_re @ B_im + A_im @ B_re
    let ar_br = a_re.matmul(b_re);
    let ai_bi = a_im.matmul(b_im);
    let ar_bi = a_re.matmul(b_im);
    let ai_br = a_im.matmul(b_re);

    let c_re = ar_br.zip_map(&ai_bi, |a, b| a - b);
    let c_im = ar_bi.zip_map(&ai_br, |a, b| a + b);
    (c_re, c_im)
}

/// Build FFT-domain filter basis for one octave.
/// Returns (basis_re, basis_im, n_fft, lengths).
fn vqt_filter_fft(
    sr: f64,
    freqs: &[f64],
    filter_scale: f64,
    norm: Option<f64>,
    hop_length: usize,
    window: fn(usize, bool) -> Vec<f64>,
    gamma: Option<f64>,
) -> (Matrix, Matrix, usize, Vec<f64>) {
    // Build time-domain wavelet filters
    let wp = WaveletParams {
        sr,
        window,
        filter_scale,
        pad_fft: true,
        norm,
        gamma,
    };
    let (filters_re, filters_im, lengths) = wavelet(freqs, &wp);

    // Determine n_fft: at least 2 * hop_length, rounded to power of 2
    let min_fft = (2 * hop_length).next_power_of_two();
    let n_fft = filters_re.cols().max(min_fft);
    // Ensure n_fft is power of 2
    let n_fft = n_fft.next_power_of_two();

    let n_freqs = freqs.len();
    let n_bins_fft = n_fft / 2 + 1;

    // Zero-pad filter rows to n_fft if needed, then FFT each row
    let mut planner = RealFftPlanner::<f64>::new();
    let fft = planner.plan_fft_forward(n_fft);

    let mut basis_re_data = vec![0.0; n_freqs * n_bins_fft];
    let mut basis_im_data = vec![0.0; n_freqs * n_bins_fft];

    let mut input_buf = fft.make_input_vec();
    let mut output_buf = fft.make_output_vec();
    let mut scratch = fft.make_scratch_vec();

    for k in 0..n_freqs {
        // Renormalize: basis[k, :] *= lengths[k] / n_fft
        let scale = lengths[k] / n_fft as f64;

        // Copy filter row into input buffer (zero-padded)
        input_buf.fill(0.0);
        let filter_cols = filters_re.cols();

        // Center the filter in n_fft (pad_center)
        let pad_start = if n_fft > filter_cols {
            (n_fft - filter_cols) / 2
        } else {
            0
        };
        let copy_len = filter_cols.min(n_fft);

        // We need to combine re and im parts and FFT them separately
        // Actually, librosa does real FFT on each filter (which is complex).
        // The wavelet filters are complex: sig = cos + j*sin
        // librosa does: fft(complex_filters), but uses scipy.fft.fft, which handles complex input.
        //
        // For complex input FFT: F{re + j*im} = F{re} + j*F{im}
        // where F{} is the real-valued signal's DFT.
        // But realfft only handles real input. So we FFT re and im separately.

        // FFT the real part
        input_buf.fill(0.0);
        for j in 0..copy_len {
            input_buf[pad_start + j] = filters_re.get(k, j) * scale;
        }
        fft.process_with_scratch(&mut input_buf, &mut output_buf, &mut scratch)
            .expect("FFT failed");

        let mut fft_re_re = vec![0.0; n_bins_fft]; // FFT of re part, real component
        let mut fft_re_im = vec![0.0; n_bins_fft]; // FFT of re part, imag component
        for j in 0..n_bins_fft {
            fft_re_re[j] = output_buf[j].re;
            fft_re_im[j] = output_buf[j].im;
        }

        // FFT the imaginary part
        input_buf.fill(0.0);
        for j in 0..copy_len {
            input_buf[pad_start + j] = filters_im.get(k, src_j_val(k, j, &filters_im)) * scale;
        }
        fft.process_with_scratch(&mut input_buf, &mut output_buf, &mut scratch)
            .expect("FFT failed");

        // F{complex_filter} = F{re} + j * F{im}
        // = (fft_re_re + j*fft_re_im) + j*(fft_im_re + j*fft_im_im)
        // = (fft_re_re - fft_im_im) + j*(fft_re_im + fft_im_re)
        for j in 0..n_bins_fft {
            let fft_im_re = output_buf[j].re;
            let fft_im_im = output_buf[j].im;
            basis_re_data[k * n_bins_fft + j] = fft_re_re[j] - fft_im_im;
            basis_im_data[k * n_bins_fft + j] = fft_re_im[j] + fft_im_re;
        }
    }

    (
        Matrix::from_vec(basis_re_data, n_freqs, n_bins_fft),
        Matrix::from_vec(basis_im_data, n_freqs, n_bins_fft),
        n_fft,
        lengths,
    )
}

// Helper to get value from matrix safely
fn src_j_val(_k: usize, j: usize, _m: &Matrix) -> usize {
    j
}

/// Compute CQT response for one octave.
/// Uses STFT with boxcar window, then multiplies by filter basis.
fn cqt_response(
    y: &[f64],
    n_fft: usize,
    hop_length: usize,
    basis_re: &Matrix,
    basis_im: &Matrix,
) -> (Matrix, Matrix) {
    // Compute STFT with boxcar (all ones) window
    let stft_params = crate::stft::StftParams {
        n_fft,
        hop_length,
        win_length: n_fft,
        center: true,
        window: Some(vec![1.0; n_fft]),
    };
    let (stft_re, stft_im) = crate::stft::stft_as_real_imag(y, &stft_params);

    // Complex matrix multiply: response = basis @ STFT
    complex_matmul(basis_re, basis_im, &stft_re, &stft_im)
}

/// Stack octave responses, trimming to minimum frame count.
/// responses are ordered [highest_octave, ..., lowest_octave].
/// Output stacks from lowest to highest frequency.
fn trim_stack(responses: &[(Matrix, Matrix, usize)], n_bins: usize) -> (Matrix, Matrix) {
    if responses.is_empty() {
        return (Matrix::zeros(n_bins, 0), Matrix::zeros(n_bins, 0));
    }

    // Find minimum frame count
    let min_frames = responses
        .iter()
        .map(|(re, _im, _n)| re.cols())
        .min()
        .unwrap_or(0);

    let mut out_re = vec![0.0; n_bins * min_frames];
    let mut out_im = vec![0.0; n_bins * min_frames];

    // Stack from lowest octave (last in list) to highest (first)
    let mut row_offset = 0;
    for (resp_re, resp_im, n_bins_oct) in responses.iter().rev() {
        // If this octave has more bins than needed, take the highest-freq bins
        let skip = if resp_re.rows() > *n_bins_oct {
            resp_re.rows() - n_bins_oct
        } else {
            0
        };
        let bins_to_copy = (*n_bins_oct).min(n_bins - row_offset);

        for b in 0..bins_to_copy {
            let src_row = skip + b;
            let dst_row = row_offset + b;
            for f in 0..min_frames {
                out_re[dst_row * min_frames + f] = resp_re.get(src_row, f);
                out_im[dst_row * min_frames + f] = resp_im.get(src_row, f);
            }
        }
        row_offset += bins_to_copy;
    }

    (
        Matrix::from_vec(out_re, n_bins, min_frames),
        Matrix::from_vec(out_im, n_bins, min_frames),
    )
}

/// Variable-Q Transform using octave decimation.
///
/// Matches `librosa.vqt(y, sr=..., gamma=...)`.
/// Returns (re, im) matrices of shape (n_bins, n_frames).
pub fn vqt(y: &[f64], params: &VqtParams) -> (Matrix, Matrix) {
    let n_octaves = (params.n_bins as f64 / params.bins_per_octave as f64).ceil() as usize;

    // Generate all target frequencies with tuning correction
    let freqs = crate::convert::cqt_frequencies(
        params.n_bins,
        params.fmin,
        params.bins_per_octave,
        params.tuning,
    );

    // Compute all filter lengths to determine f_cutoff and early downsampling
    let wl_params = WaveletLengthsParams {
        sr: params.sr,
        window: params.window,
        filter_scale: params.filter_scale,
        gamma: params.gamma,
    };
    let (all_lengths, f_cutoff) = wavelet_lengths(&freqs, &wl_params);

    // Early downsampling
    let nyquist = params.sr / 2.0;
    let num_twos = params.hop_length.trailing_zeros() as usize;

    let downsample_count = if f_cutoff > 0.0 && nyquist > f_cutoff {
        let max_ds = ((nyquist / f_cutoff).log2().floor() as isize - 1 - (n_octaves as isize - 1))
            .max(0) as usize;
        max_ds.min(num_twos)
    } else {
        0
    };

    let mut my_y: Vec<f64>;
    let mut my_sr: f64;

    if downsample_count > 0 {
        let factor = 1u32 << downsample_count;
        let new_sr = (params.sr as u32) / factor;
        my_y = crate::audio::resample(y, params.sr as u32, new_sr)
            .expect("CQT internal resample failed");
        my_sr = new_sr as f64;
    } else {
        my_y = y.to_vec();
        my_sr = params.sr;
    }

    let original_sr = params.sr;
    let mut my_hop = params.hop_length >> downsample_count;

    // Process octaves from highest frequency to lowest
    let mut responses: Vec<(Matrix, Matrix, usize)> = Vec::with_capacity(n_octaves);

    for oct in 0..n_octaves {
        // Determine bins for this octave
        let bin_start = params
            .n_bins
            .saturating_sub((oct + 1) * params.bins_per_octave);
        let bin_end = params.n_bins.saturating_sub(oct * params.bins_per_octave);
        let n_bins_oct = bin_end - bin_start;

        if n_bins_oct == 0 {
            continue;
        }

        let freqs_oct = &freqs[bin_start..bin_end];

        // Build filter basis for this octave
        let (basis_re, basis_im, n_fft, lengths_oct) = vqt_filter_fft(
            my_sr,
            freqs_oct,
            params.filter_scale,
            params.norm,
            my_hop,
            params.window,
            params.gamma,
        );

        // Scale basis by sqrt(original_sr / my_sr) for resampling energy compensation
        let sr_scale = (original_sr / my_sr).sqrt();
        let basis_re = basis_re.map(|v| v * sr_scale);
        let basis_im = basis_im.map(|v| v * sr_scale);

        // Compute response for this octave
        let (resp_re, resp_im) = cqt_response(&my_y, n_fft, my_hop, &basis_re, &basis_im);

        responses.push((resp_re, resp_im, n_bins_oct));

        // Downsample for next octave (if hop_length is even)
        if oct < n_octaves - 1 && my_hop.is_multiple_of(2) {
            let new_sr = (my_sr / 2.0) as u32;
            if new_sr > 0 {
                my_y = crate::audio::resample(&my_y, my_sr as u32, new_sr)
                    .expect("CQT internal resample failed");
                my_sr /= 2.0;
                my_hop /= 2;
            }
        }
        // Suppress unused variable warning
        let _ = &lengths_oct;
    }

    // Stack all octave responses
    let (mut out_re, mut out_im) = trim_stack(&responses, params.n_bins);

    // Apply scaling: divide each bin by sqrt(filter_length)
    if params.scale {
        for k in 0..params.n_bins {
            let scale = 1.0 / all_lengths[k].sqrt();
            for f in 0..out_re.cols() {
                *out_re.get_mut(k, f) *= scale;
                *out_im.get_mut(k, f) *= scale;
            }
        }
    }

    (out_re, out_im)
}

/// Constant-Q Transform. Thin wrapper over VQT with gamma=0.
///
/// Matches `librosa.cqt(y, sr=..., ...)`.
/// Returns (re, im) matrices of shape (n_bins, n_frames).
pub fn cqt(y: &[f64], params: &CqtParams) -> (Matrix, Matrix) {
    let vqt_params = VqtParams {
        sr: params.sr,
        hop_length: params.hop_length,
        fmin: params.fmin,
        n_bins: params.n_bins,
        bins_per_octave: params.bins_per_octave,
        tuning: params.tuning,
        filter_scale: params.filter_scale,
        norm: params.norm,
        window: params.window,
        scale: params.scale,
        gamma: Some(0.0), // CQT: gamma = 0
    };
    vqt(y, &vqt_params)
}

/// Compute element-wise magnitude from complex (re, im) pair.
///
/// `mag[i,j] = sqrt(re[i,j]^2 + im[i,j]^2)`
pub fn complex_magnitude(re: &Matrix, im: &Matrix) -> Matrix {
    re.zip_map(im, |r, i| (r * r + i * i).sqrt())
}

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

    #[test]
    fn test_complex_matmul_simple() {
        // (1+2j) * (3+4j) = 3+4j+6j+8j^2 = (3-8)+(4+6)j = -5+10j
        let a_re = Matrix::from_vec(vec![1.0], 1, 1);
        let a_im = Matrix::from_vec(vec![2.0], 1, 1);
        let b_re = Matrix::from_vec(vec![3.0], 1, 1);
        let b_im = Matrix::from_vec(vec![4.0], 1, 1);
        let (c_re, c_im) = complex_matmul(&a_re, &a_im, &b_re, &b_im);
        assert!((c_re.get(0, 0) - (-5.0)).abs() < 1e-12);
        assert!((c_im.get(0, 0) - 10.0).abs() < 1e-12);
    }

    #[test]
    fn test_complex_magnitude() {
        let re = Matrix::from_vec(vec![3.0, 0.0], 1, 2);
        let im = Matrix::from_vec(vec![4.0, 1.0], 1, 2);
        let mag = complex_magnitude(&re, &im);
        assert!((mag.get(0, 0) - 5.0).abs() < 1e-12);
        assert!((mag.get(0, 1) - 1.0).abs() < 1e-12);
    }

    #[test]
    fn test_cqt_basic_shape() {
        // Generate a simple sine wave
        let sr = 22050.0;
        let duration = 1.0;
        let n_samples = (sr * duration) as usize;
        let y: Vec<f64> = (0..n_samples)
            .map(|i| (2.0 * std::f64::consts::PI * 440.0 * i as f64 / sr).sin())
            .collect();

        let params = CqtParams {
            sr,
            n_bins: 84,
            bins_per_octave: 12,
            hop_length: 512,
            ..CqtParams::default()
        };
        let (re, im) = cqt(&y, &params);
        assert_eq!(re.rows(), 84);
        assert_eq!(im.rows(), 84);
        assert_eq!(re.cols(), im.cols());
        assert!(re.cols() > 0);
    }
}