rosa 0.1.0

//! Onset strength envelope matching `librosa.onset.onset_strength`.

use crate::dsp::db::power_to_db;
use crate::matrix::Matrix;
use crate::stft::{MelSpectrogramParams, mel_spectrogram};

pub use crate::matrix::Aggregate;

/// Parameters for onset strength computation.
pub struct OnsetStrengthParams {
    /// Sample rate in Hz.
    pub sr: f64,
    /// FFT window size.
    pub n_fft: usize,
    /// Number of samples between frames.
    pub hop_length: usize,
    /// Number of frames to lag when computing spectral flux.
    pub lag: usize,
    /// Size of the maximum filter applied along frequency bins.
    pub max_size: usize,
    /// If true, apply IIR high-pass filter to remove DC offset.
    pub detrend: bool,
    /// If true, pad the onset envelope to align with centered STFT frames.
    pub center: bool,
    /// Aggregation function applied across frequency bins.
    pub aggregate: Aggregate,
}

impl OnsetStrengthParams {
    /// Create parameters with the given sample rate and library defaults.
    pub fn new(sr: f64) -> Self {
        Self {
            sr,
            n_fft: 2048,
            hop_length: 512,
            lag: 1,
            max_size: 1,
            detrend: false,
            center: true,
            aggregate: Aggregate::Mean,
        }
    }
}

fn median(values: &mut [f64]) -> f64 {
    let n = values.len();
    if n == 0 {
        return 0.0;
    }
    values.sort_by(|a, b| a.total_cmp(b));
    if n.is_multiple_of(2) {
        (values[n / 2 - 1] + values[n / 2]) / 2.0
    } else {
        values[n / 2]
    }
}

/// Compute spectral flux onset strength envelope.
///
/// Matches `librosa.onset.onset_strength(y=y, sr=sr)` with defaults.
///
/// Algorithm:
/// 1. Compute mel spectrogram, convert to dB
/// 2. Compute spectral flux: `S[:, lag:] - ref[:, :-lag]`
/// 3. Half-wave rectify (discard negatives)
/// 4. Mean across frequency bins
/// 5. Pad to compensate for lag and centering
/// 6. Optional detrend (DC removal)
pub fn onset_strength(y: &[f64], params: &OnsetStrengthParams) -> Vec<f64> {
    let sr = params.sr;
    let n_fft = params.n_fft;
    let hop_length = params.hop_length;
    let lag = params.lag;

    // Step 1: Compute mel spectrogram
    let mut mel_params = MelSpectrogramParams::new(sr);
    mel_params.stft.n_fft = n_fft;
    mel_params.stft.hop_length = hop_length;
    // librosa onset_strength uses fmax=sr/2 (same as default)
    mel_params.mel.n_fft = n_fft;
    mel_params.mel.fmax = Some(sr / 2.0);

    let mel_s = mel_spectrogram(y, &mel_params);
    let n_mels = mel_s.rows();
    let n_frames = mel_s.cols();

    // Step 2: Convert to dB
    let db_flat = power_to_db(mel_s.as_slice(), 1.0, 1e-10, Some(80.0));
    let s_db = Matrix::from_vec(db_flat, n_mels, n_frames);

    // Step 3: Compute reference spectrum
    // With max_size == 1, ref = S (no filtering)
    let ref_db = if params.max_size > 1 {
        // Apply maximum_filter1d along frequency axis (rows) for each frame
        let mut ref_data = vec![0.0; n_mels * n_frames];
        for t in 0..n_frames {
            // Extract column
            let col: Vec<f64> = (0..n_mels).map(|r| s_db.get(r, t)).collect();
            let filtered = crate::dsp::util::maximum_filter1d(&col, params.max_size);
            for (r, &val) in filtered.iter().enumerate() {
                ref_data[r * n_frames + t] = val;
            }
        }
        Matrix::from_vec(ref_data, n_mels, n_frames)
    } else {
        s_db.clone()
    };

    if n_frames <= lag {
        // Not enough frames to compute flux
        return vec![0.0; n_frames];
    }

    // Step 4: Compute spectral flux: S[:, lag:] - ref[:, :-lag]
    // Half-wave rectify (keep only positive values)
    let flux_frames = n_frames - lag;
    let mut onset_env = vec![0.0; flux_frames];

    for t in 0..flux_frames {
        let mut values: Vec<f64> = Vec::with_capacity(n_mels);
        for r in 0..n_mels {
            let diff = s_db.get(r, t + lag) - ref_db.get(r, t);
            values.push(diff.max(0.0));
        }
        onset_env[t] = match params.aggregate {
            Aggregate::Mean => values.iter().sum::<f64>() / n_mels as f64,
            Aggregate::Median => median(&mut values),
        };
    }

    // Step 5: Pad to compensate for lag and centering
    let mut pad_width = lag;
    if params.center {
        pad_width += n_fft / (2 * hop_length);
    }

    let mut padded = vec![0.0; pad_width];
    padded.extend_from_slice(&onset_env);

    // Step 6: Optional detrend (DC removal via IIR filter)
    if params.detrend {
        // y[n] = x[n] - x[n-1] + 0.99 * y[n-1]
        // lfilter([1.0, -1.0], [1.0, -0.99], x)
        let mut filtered = vec![0.0; padded.len()];
        if !padded.is_empty() {
            filtered[0] = padded[0];
            for i in 1..padded.len() {
                filtered[i] = padded[i] - padded[i - 1] + 0.99 * filtered[i - 1];
            }
        }
        padded = filtered;
    }

    // Step 7: Trim to match input frame count if center
    if params.center {
        padded.truncate(n_frames);
    }

    padded
}