rosa 0.1.0

//! Chroma feature extraction.

use crate::matrix::{Axis, Matrix};

/// Chroma filter bank parameters.
pub struct ChromaFilterParams {
    /// Sample rate in Hz.
    pub sr: f64,
    /// FFT window size.
    pub n_fft: usize,
    /// Number of chroma bins (pitch classes).
    pub n_chroma: usize,
    /// Tuning offset in fractions of a bin.
    pub tuning: f64,
    /// Center octave for the Gaussian octave weighting envelope.
    pub ctroct: f64,
    /// Width of the octave weighting envelope; None disables weighting.
    pub octwidth: Option<f64>,
    /// Column-wise normalization: Some(2.0) = L2 (default), None = no normalization.
    pub norm: Option<f64>,
    /// If true, roll the filter bank so that C is the first chroma bin.
    pub base_c: bool,
}

impl Default for ChromaFilterParams {
    fn default() -> Self {
        Self {
            sr: 22050.0,
            n_fft: 2048,
            n_chroma: 12,
            tuning: 0.0,
            ctroct: 5.0,
            octwidth: Some(2.0),
            norm: Some(2.0),
            base_c: true,
        }
    }
}

/// Chroma STFT parameters.
pub struct ChromaStftParams {
    /// Sample rate in Hz.
    pub sr: f64,
    /// FFT window size.
    pub n_fft: usize,
    /// Number of samples between successive frames.
    pub hop_length: usize,
    /// Length of the windowed segment used for each FFT.
    pub win_length: usize,
    /// Number of chroma bins (pitch classes).
    pub n_chroma: usize,
    /// Tuning offset; None = auto-estimate.
    pub tuning: Option<f64>,
    /// Column normalization: Some(f64::INFINITY) = Linf (default).
    pub norm: Option<f64>,
    /// Center octave for the Gaussian octave weighting envelope.
    pub ctroct: f64,
    /// Width of the octave weighting envelope; None disables weighting.
    pub octwidth: Option<f64>,
    /// If true, roll the filter bank so that C is the first chroma bin.
    pub base_c: bool,
}

impl Default for ChromaStftParams {
    fn default() -> Self {
        Self {
            sr: 22050.0,
            n_fft: 2048,
            hop_length: 512,
            win_length: 2048,
            n_chroma: 12,
            tuning: None,
            norm: Some(f64::INFINITY),
            ctroct: 5.0,
            octwidth: Some(2.0),
            base_c: true,
        }
    }
}

/// CQT mode for chroma computation.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum CqtMode {
    /// Use the standard full CQT.
    Full,
    /// Use the hybrid pseudo/full CQT.
    Hybrid,
}

/// Chroma CQT parameters.
pub struct ChromaCqtParams {
    /// Sample rate in Hz.
    pub sr: f64,
    /// Number of samples between successive frames.
    pub hop_length: usize,
    /// Minimum frequency in Hz; None defaults to C1.
    pub fmin: Option<f64>,
    /// Number of chroma bins (pitch classes).
    pub n_chroma: usize,
    /// Number of octaves to analyze.
    pub n_octaves: usize,
    /// Number of CQT bins per octave.
    pub bins_per_octave: usize,
    /// Tuning offset; None = auto-estimate.
    pub tuning: Option<f64>,
    /// Column normalization. Default: Some(f64::INFINITY).
    pub norm: Option<f64>,
    /// Minimum chroma value to keep; values below this are set to zero.
    pub threshold: f64,
    /// Which CQT variant to use.
    pub cqt_mode: CqtMode,
}

impl Default for ChromaCqtParams {
    fn default() -> Self {
        Self {
            sr: 22050.0,
            hop_length: 512,
            fmin: None,
            n_chroma: 12,
            n_octaves: 7,
            bins_per_octave: 36,
            tuning: None,
            norm: Some(f64::INFINITY),
            threshold: 0.0,
            cqt_mode: CqtMode::Full,
        }
    }
}

/// Chroma CENS parameters.
pub struct ChromaCensParams {
    /// Sample rate in Hz.
    pub sr: f64,
    /// Number of samples between successive frames.
    pub hop_length: usize,
    /// Minimum frequency in Hz; None defaults to C1.
    pub fmin: Option<f64>,
    /// Number of chroma bins (pitch classes).
    pub n_chroma: usize,
    /// Number of octaves to analyze.
    pub n_octaves: usize,
    /// Number of CQT bins per octave.
    pub bins_per_octave: usize,
    /// Which CQT variant to use.
    pub cqt_mode: CqtMode,
    /// Final normalization. Default: Some(2.0) = L2.
    pub norm: Option<f64>,
    /// Length of the temporal smoothing window in frames; None disables smoothing.
    pub win_len_smooth: Option<usize>,
    /// Window function used for temporal smoothing.
    pub smoothing_window: fn(usize, bool) -> Vec<f64>,
}

impl Default for ChromaCensParams {
    fn default() -> Self {
        Self {
            sr: 22050.0,
            hop_length: 512,
            fmin: None,
            n_chroma: 12,
            n_octaves: 7,
            bins_per_octave: 36,
            cqt_mode: CqtMode::Full,
            norm: Some(2.0),
            win_len_smooth: Some(41),
            smoothing_window: crate::windows::hann,
        }
    }
}

/// Build STFT-to-chroma projection filter bank.
///
/// Matches `librosa.filters.chroma(sr=..., n_fft=..., ...)`.
/// Returns Matrix of shape (n_chroma, 1 + n_fft/2).
pub fn chroma_filter(params: &ChromaFilterParams) -> Matrix {
    let n_chroma = params.n_chroma;
    let n_fft = params.n_fft;
    let n_cols = 1 + n_fft / 2;

    // Compute frequencies: linspace(0, sr, n_fft, endpoint=false)
    // We only need the first n_fft values, then take [0..n_cols]
    let freq_step = params.sr / n_fft as f64;

    // Map each FFT bin to fractional octave position
    let mut frqbins = vec![0.0f64; n_fft];
    for k in 1..n_fft {
        let freq = k as f64 * freq_step;
        let octs = crate::convert::hz_to_octs_tuned(freq, params.tuning, n_chroma);
        frqbins[k] = n_chroma as f64 * octs;
    }
    // Fake DC bin
    frqbins[0] = frqbins[1] - 1.5 * n_chroma as f64;

    // Compute bin widths
    let mut binwidths = vec![1.0f64; n_fft];
    for k in 0..n_fft - 1 {
        binwidths[k] = (frqbins[k + 1] - frqbins[k]).max(1.0);
    }

    // Build Gaussian bumps: wts[c, k] = exp(-0.5 * (2 * D / binwidths[k])^2)
    // where D = frqbins[k] - c, wrapped to [-n_chroma/2, n_chroma/2]
    let n_chroma_f = n_chroma as f64;
    let half = n_chroma_f / 2.0;

    let mut wts = vec![0.0f64; n_chroma * n_fft];
    for c in 0..n_chroma {
        for k in 0..n_fft {
            let d = frqbins[k] - c as f64;
            // Wrap to [-half, half]
            let d_wrapped = ((d + half + 10.0 * n_chroma_f) % n_chroma_f) - half;
            let val = 2.0 * d_wrapped / binwidths[k];
            wts[c * n_fft + k] = (-0.5 * val * val).exp();
        }
    }

    // Column-wise normalization (normalize each column = each FFT bin)
    if let Some(p) = params.norm {
        for k in 0..n_fft {
            let norm_val = if p == 1.0 {
                (0..n_chroma).map(|c| wts[c * n_fft + k].abs()).sum::<f64>()
            } else if p == 2.0 {
                (0..n_chroma)
                    .map(|c| wts[c * n_fft + k] * wts[c * n_fft + k])
                    .sum::<f64>()
                    .sqrt()
            } else if p.is_infinite() {
                (0..n_chroma)
                    .map(|c| wts[c * n_fft + k].abs())
                    .fold(0.0_f64, f64::max)
            } else {
                1.0
            };
            if norm_val > f64::MIN_POSITIVE {
                for c in 0..n_chroma {
                    wts[c * n_fft + k] /= norm_val;
                }
            }
        }
    }

    // Apply octave weighting
    if let Some(octwidth) = params.octwidth {
        for k in 0..n_fft {
            let oct_pos = frqbins[k] / n_chroma_f;
            let diff = oct_pos - params.ctroct;
            let weight = (-0.5 * (diff / octwidth).powi(2)).exp();
            for c in 0..n_chroma {
                wts[c * n_fft + k] *= weight;
            }
        }
    }

    // Roll to start at C if base_c
    if params.base_c {
        let roll = (n_chroma as isize - 3 * (n_chroma as isize / 12)).rem_euclid(n_chroma as isize)
            as usize;
        if roll != 0 {
            // Roll rows
            let mut new_wts = vec![0.0; n_chroma * n_fft];
            for c in 0..n_chroma {
                let new_c = (c + roll) % n_chroma;
                for k in 0..n_fft {
                    new_wts[new_c * n_fft + k] = wts[c * n_fft + k];
                }
            }
            wts = new_wts;
        }
    }

    // Truncate to n_cols = 1 + n_fft/2
    let mut result_data = vec![0.0; n_chroma * n_cols];
    for c in 0..n_chroma {
        for k in 0..n_cols {
            result_data[c * n_cols + k] = wts[c * n_fft + k];
        }
    }

    Matrix::from_vec(result_data, n_chroma, n_cols)
}

/// Build CQT-bin to chroma-bin mapping matrix.
///
/// Matches `librosa.filters.cq_to_chroma(n_input, ...)`.
/// Returns Matrix of shape (n_chroma, n_input).
pub fn cq_to_chroma(
    n_input: usize,
    bins_per_octave: usize,
    n_chroma: usize,
    fmin: Option<f64>,
    base_c: bool,
) -> Matrix {
    assert!(
        bins_per_octave.is_multiple_of(n_chroma),
        "bins_per_octave ({bins_per_octave}) must be divisible by n_chroma ({n_chroma})"
    );

    let n_merge = bins_per_octave / n_chroma;

    // Build one-octave mapping: (n_chroma, bins_per_octave)
    let mut octave_map = vec![0.0f64; n_chroma * bins_per_octave];
    for c in 0..n_chroma {
        // Each chroma bin covers n_merge consecutive CQT bins, centered
        let start = c * n_merge;
        for m in 0..n_merge {
            let col = (start + m) % bins_per_octave;
            octave_map[c * bins_per_octave + col] = 1.0;
        }
    }

    // Roll left by n_merge/2 to center (matches np.roll(arr, -n_merge//2, axis=1))
    let roll_amount = n_merge as isize / 2;
    let mut rolled = vec![0.0f64; n_chroma * bins_per_octave];
    for c in 0..n_chroma {
        for b in 0..bins_per_octave {
            let src = ((b as isize + roll_amount) % bins_per_octave as isize
                + bins_per_octave as isize) as usize
                % bins_per_octave;
            rolled[c * bins_per_octave + b] = octave_map[c * bins_per_octave + src];
        }
    }

    // Tile to n_input columns
    let n_octaves_tile = n_input.div_ceil(bins_per_octave);
    let mut tiled = vec![0.0f64; n_chroma * n_input];
    for c in 0..n_chroma {
        for oct in 0..n_octaves_tile {
            for b in 0..bins_per_octave {
                let col = oct * bins_per_octave + b;
                if col < n_input {
                    tiled[c * n_input + col] = rolled[c * bins_per_octave + b];
                }
            }
        }
    }

    // Roll to start at C based on fmin MIDI note
    if base_c {
        let fmin_val = fmin.unwrap_or(crate::convert::midi_to_hz(24.0)); // C1
        let midi = crate::convert::hz_to_midi(fmin_val);
        // Compute the roll amount to align fmin to C
        let chroma_offset =
            (midi.round() as isize % n_chroma as isize + n_chroma as isize) as usize % n_chroma;
        if chroma_offset != 0 {
            let roll = n_chroma - chroma_offset;
            let mut rolled_final = vec![0.0f64; n_chroma * n_input];
            for c in 0..n_chroma {
                let new_c = (c + roll) % n_chroma;
                for k in 0..n_input {
                    rolled_final[new_c * n_input + k] = tiled[c * n_input + k];
                }
            }
            tiled = rolled_final;
        }
    }

    Matrix::from_vec(tiled, n_chroma, n_input)
}

/// Chromagram from STFT.
///
/// Matches `librosa.feature.chroma_stft(y=y, sr=sr, ...)`.
/// Returns Matrix of shape (n_chroma, n_frames).
pub fn chroma_stft(y: &[f64], params: &ChromaStftParams) -> Matrix {
    // Compute power spectrogram
    let stft_params = crate::stft::StftParams {
        n_fft: params.n_fft,
        hop_length: params.hop_length,
        win_length: params.win_length,
        center: true,
        window: None,
    };
    let s = crate::stft::spectrogram(y, &stft_params, 2.0);

    // Auto-estimate tuning if needed
    let tuning = params.tuning.unwrap_or_else(|| {
        crate::pitch::estimate_tuning(
            Some(y),
            params.sr,
            None,
            params.n_fft,
            0.05,
            params.n_chroma,
        )
    });

    // Build chroma filter bank
    let filter_params = ChromaFilterParams {
        sr: params.sr,
        n_fft: params.n_fft,
        n_chroma: params.n_chroma,
        tuning,
        ctroct: params.ctroct,
        octwidth: params.octwidth,
        norm: Some(2.0), // Filter bank always uses L2
        base_c: params.base_c,
    };
    let chromafb = chroma_filter(&filter_params);

    // Project: raw_chroma = chromafb @ S
    let raw_chroma = chromafb.matmul(&s);

    // Normalize columns
    match params.norm {
        Some(p) => raw_chroma.normalize_axis(Axis::Row, p),
        None => raw_chroma,
    }
}

/// Chromagram from CQT.
///
/// Matches `librosa.feature.chroma_cqt(y=y, sr=sr, ...)`.
/// Returns Matrix of shape (n_chroma, n_frames).
pub fn chroma_cqt(y: &[f64], params: &ChromaCqtParams) -> Matrix {
    let fmin = params.fmin.unwrap_or(crate::convert::midi_to_hz(24.0)); // C1
    let n_bins = params.n_octaves * params.bins_per_octave;

    // Auto-estimate tuning (must use bins_per_octave, not n_chroma,
    // to match librosa which passes bins_per_octave to estimate_tuning)
    let tuning = params.tuning.unwrap_or_else(|| {
        crate::pitch::estimate_tuning(Some(y), params.sr, None, 2048, 0.05, params.bins_per_octave)
    });

    // Compute CQT
    let cqt_params = crate::cqt::CqtParams {
        sr: params.sr,
        hop_length: params.hop_length,
        fmin,
        n_bins,
        bins_per_octave: params.bins_per_octave,
        tuning,
        filter_scale: 1.0,
        norm: Some(1.0),
        window: crate::windows::hann,
        scale: true,
    };

    let (cqt_re, cqt_im) = match params.cqt_mode {
        CqtMode::Full => crate::cqt::cqt(y, &cqt_params),
        CqtMode::Hybrid => super::hybrid::hybrid_cqt(y, &cqt_params),
    };
    let c = crate::cqt::complex_magnitude(&cqt_re, &cqt_im);

    // Map CQT to chroma
    let cq2chr = cq_to_chroma(
        n_bins,
        params.bins_per_octave,
        params.n_chroma,
        Some(fmin),
        true,
    );
    let mut chroma = cq2chr.matmul(&c);

    // Apply threshold
    if params.threshold > 0.0 {
        chroma.map_inplace(|v| if v < params.threshold { 0.0 } else { v });
    }

    // Normalize columns
    match params.norm {
        Some(p) => chroma.normalize_axis(Axis::Row, p),
        None => chroma,
    }
}

/// CENS chroma — quantized and smoothed.
///
/// Matches `librosa.feature.chroma_cens(y=y, sr=sr, ...)`.
/// Returns Matrix of shape (n_chroma, n_frames).
pub fn chroma_cens(y: &[f64], params: &ChromaCensParams) -> Matrix {
    // Step 1: Compute chroma_cqt with norm=None
    let cqt_params = ChromaCqtParams {
        sr: params.sr,
        hop_length: params.hop_length,
        fmin: params.fmin,
        n_chroma: params.n_chroma,
        n_octaves: params.n_octaves,
        bins_per_octave: params.bins_per_octave,
        tuning: None,
        norm: None,
        threshold: 0.0,
        cqt_mode: params.cqt_mode,
    };
    let chroma = chroma_cqt(y, &cqt_params);

    let n_chroma = chroma.rows();
    let n_frames = chroma.cols();

    // Step 2: L1-normalize each column
    let chroma = chroma.normalize_axis(Axis::Row, 1.0);

    // Step 3: Quantize
    const STEPS: [f64; 4] = [0.4, 0.2, 0.1, 0.05];
    const WEIGHTS: [f64; 4] = [0.25, 0.25, 0.25, 0.25];

    let mut quant_data = vec![0.0f64; n_chroma * n_frames];
    for c in 0..n_chroma {
        for f in 0..n_frames {
            let val = chroma.get(c, f);
            let mut q = 0.0;
            for (&step, &weight) in STEPS.iter().zip(WEIGHTS.iter()) {
                if val > step {
                    q += weight;
                }
            }
            quant_data[c * n_frames + f] = q;
        }
    }
    let mut quant = Matrix::from_vec(quant_data, n_chroma, n_frames);

    // Step 4: Temporal smoothing
    if let Some(win_len) = params.win_len_smooth
        && win_len > 1
        && n_frames > 1
    {
        // Build smoothing window (symmetric!)
        let win = (params.smoothing_window)(win_len, false);
        let win_sum: f64 = win.iter().sum();
        let norm_win: Vec<f64> = win.iter().map(|&w| w / win_sum).collect();

        // Convolve each row (mode='same')
        let half = win_len / 2;
        let mut smoothed = vec![0.0f64; n_chroma * n_frames];
        for c in 0..n_chroma {
            for f in 0..n_frames {
                let mut sum = 0.0;
                for w in 0..win_len {
                    let idx = f as isize + w as isize - half as isize;
                    if idx >= 0 && (idx as usize) < n_frames {
                        sum += quant.get(c, idx as usize) * norm_win[w];
                    }
                }
                smoothed[c * n_frames + f] = sum;
            }
        }
        quant = Matrix::from_vec(smoothed, n_chroma, n_frames);
    }

    // Step 5: Final normalization (default L2)
    match params.norm {
        Some(p) => quant.normalize_axis(Axis::Row, p),
        None => quant,
    }
}