stratum-dsp 1.0.0

//! Autocorrelation-based BPM estimation
//!
//! Finds periodicity in onset signal using FFT-accelerated autocorrelation.
//!
//! # Algorithm
//!
//! This module implements the autocorrelation-based tempo estimation algorithm
//! described in Ellis & Pikrakis (2006). The process:
//!
//! 1. Convert onset list to binary beat signal (frame-based)
//! 2. Compute autocorrelation using FFT acceleration: `ACF = IFFT(|FFT(signal)|²)`
//! 3. Find peaks in autocorrelation function
//! 4. Convert lag values to BPM: `BPM = (60 * sample_rate) / (lag * hop_size)`
//! 5. Filter candidates within BPM range
//!
//! # Reference
//!
//! Ellis, D. P. W., & Pikrakis, A. (2006). Real-time Beat Induction.
//! *Proceedings of the International Conference on Music Information Retrieval*.
//!
//! # Example
//!
//! ```no_run
//! use stratum_dsp::features::period::autocorrelation::estimate_bpm_from_autocorrelation;
//! use stratum_dsp::features::period::BpmCandidate;
//!
//! let onsets = vec![0, 11025, 22050, 33075]; // 120 BPM at 44.1kHz, 512 hop
//! let candidates = estimate_bpm_from_autocorrelation(
//!     &onsets,
//!     44100,
//!     512,
//!     60.0,
//!     180.0,
//! )?;
//! # Ok::<(), stratum_dsp::AnalysisError>(())
//! ```

use super::BpmCandidate;
use rustfft::FftPlanner;
use rustfft::num_complex::Complex;

const EPSILON: f32 = 1e-10;

/// Estimate BPM from autocorrelation
///
/// This function implements the autocorrelation-based tempo estimation algorithm
/// described in Ellis & Pikrakis (2006). The algorithm finds periodicity in the
/// onset signal by computing the autocorrelation function, which reveals repeating
/// patterns corresponding to the beat period.
///
/// # Reference
///
/// Ellis, D. P. W., & Pikrakis, A. (2006). Real-time Beat Induction.
/// *Proceedings of the International Conference on Music Information Retrieval*.
///
/// # Arguments
///
/// * `onsets` - Onset times in samples
/// * `sample_rate` - Sample rate in Hz
/// * `hop_size` - Hop size used for onset detection (samples per frame)
/// * `min_bpm` - Minimum BPM to consider (default: 60.0)
/// * `max_bpm` - Maximum BPM to consider (default: 180.0)
///
/// # Returns
///
/// Vector of BPM candidates ranked by confidence (highest first)
///
/// # Errors
///
/// Returns `AnalysisError` if:
/// - Onset list is empty or too short
/// - Invalid parameters (sample_rate=0, hop_size=0)
/// - Numerical errors during FFT computation
///
/// # Algorithm Details
///
/// 1. **Binary Signal**: Convert onsets to frame-based binary signal
///    - Frame index = `onset_sample / hop_size`
///    - Signal[frame] = 1 if onset present, else 0
///
/// 2. **Autocorrelation**: Compute using FFT acceleration
///    - `ACF[lag] = IFFT(|FFT(signal)|²)`
///    - Complexity: O(n log n) instead of O(n²)
///    - This FFT acceleration is a key contribution of Ellis & Pikrakis (2006)
///
/// 3. **Peak Detection**: Find local maxima in ACF
///    - Filter by BPM range (min_lag, max_lag)
///    - Compute prominence (height relative to neighbors)
///    - Peaks in ACF correspond to periodicities in the onset signal
///
/// 4. **BPM Conversion**: `BPM = (60 * sample_rate) / (lag * hop_size)`
///    - Lag values represent the period in frames
///    - Convert to BPM using the relationship between frames and time
///
/// # Performance
///
/// Typical performance: 5-15ms for 30s track (44100 Hz, 512 hop)
/// The FFT acceleration makes this method efficient for real-time applications.
pub fn estimate_bpm_from_autocorrelation(
    onsets: &[usize],
    sample_rate: u32,
    hop_size: usize,
    min_bpm: f32,
    max_bpm: f32,
) -> Result<Vec<BpmCandidate>, crate::error::AnalysisError> {
    log::debug!(
        "Estimating BPM from autocorrelation: {} onsets, {} Hz, hop={}, range=[{:.1}, {:.1}] BPM",
        onsets.len(),
        sample_rate,
        hop_size,
        min_bpm,
        max_bpm
    );

    // Validate inputs
    if onsets.is_empty() {
        return Err(crate::error::AnalysisError::InvalidInput(
            "Empty onset list".to_string(),
        ));
    }

    if sample_rate == 0 {
        return Err(crate::error::AnalysisError::InvalidInput(
            "Invalid sample rate: 0".to_string(),
        ));
    }

    if hop_size == 0 {
        return Err(crate::error::AnalysisError::InvalidInput(
            "Invalid hop size: 0".to_string(),
        ));
    }

    if min_bpm <= 0.0 || max_bpm <= 0.0 || min_bpm >= max_bpm {
        return Err(crate::error::AnalysisError::InvalidInput(format!(
            "Invalid BPM range: [{:.1}, {:.1}]",
            min_bpm, max_bpm
        )));
    }

    // Need at least 2 onsets for autocorrelation
    if onsets.len() < 2 {
        log::warn!("Too few onsets for autocorrelation: {}", onsets.len());
        return Ok(vec![]);
    }

    // Step 1: Convert onsets to frame-based binary signal
    let max_frame = onsets.iter().max().copied().unwrap_or(0) / hop_size;
    let signal_length = max_frame + 1;

    if signal_length < 2 {
        return Err(crate::error::AnalysisError::ProcessingError(
            "Signal too short for autocorrelation".to_string(),
        ));
    }

    let mut beat_signal = vec![0.0f32; signal_length];
    for &onset_sample in onsets {
        let frame_idx = onset_sample / hop_size;
        if frame_idx < signal_length {
            beat_signal[frame_idx] = 1.0;
        }
    }

    // Step 2: Compute autocorrelation using FFT acceleration
    let acf = compute_autocorrelation_fft(&beat_signal)?;

    // Step 3: Convert lag range to frame indices
    // BPM = (60 * sample_rate) / (lag * hop_size)
    // lag = (60 * sample_rate) / (BPM * hop_size)
    let lag_min = ((60.0 * sample_rate as f32) / (max_bpm * hop_size as f32)).ceil() as usize;
    let lag_max = ((60.0 * sample_rate as f32) / (min_bpm * hop_size as f32)).floor() as usize;

    if lag_min >= lag_max || lag_min >= acf.len() || lag_max >= acf.len() {
        log::warn!(
            "Invalid lag range: [{}, {}] for ACF length {}",
            lag_min,
            lag_max,
            acf.len()
        );
        return Ok(vec![]);
    }

    // Step 4: Find peaks in ACF within lag range
    let peaks = find_peaks_in_acf(&acf[lag_min..=lag_max], lag_min)?;

    // Step 5: Convert peaks to BPM candidates
    let mut candidates = Vec::new();
    for (lag, value) in peaks {
        let bpm = (60.0 * sample_rate as f32) / ((lag as f32) * (hop_size as f32));
        
        // Double-check BPM is in range (due to rounding)
        if bpm >= min_bpm && bpm <= max_bpm {
            // Normalize confidence by maximum ACF value
            let max_acf = acf.iter().copied().fold(0.0f32, f32::max);
            let confidence = if max_acf > EPSILON {
                (value / max_acf).min(1.0)
            } else {
                0.0
            };

            candidates.push(BpmCandidate { bpm, confidence });
        }
    }

    // Sort by confidence (highest first)
    candidates.sort_by(|a, b| b.confidence.partial_cmp(&a.confidence).unwrap_or(std::cmp::Ordering::Equal));

    log::debug!(
        "Autocorrelation found {} BPM candidates",
        candidates.len()
    );

    Ok(candidates)
}

/// Compute autocorrelation using FFT acceleration
///
/// Uses the identity: ACF = IFFT(|FFT(signal)|²)
///
/// # Arguments
///
/// * `signal` - Input signal (binary beat signal)
///
/// # Returns
///
/// Autocorrelation function (same length as input)
fn compute_autocorrelation_fft(signal: &[f32]) -> Result<Vec<f32>, crate::error::AnalysisError> {
    let n = signal.len();
    
    // FFT size: next power of 2 >= 2*n (for zero-padding)
    let fft_size = (2 * n).next_power_of_two();
    
    // Convert to complex and zero-pad
    let mut fft_input: Vec<Complex<f32>> = signal
        .iter()
        .map(|&x| Complex::new(x, 0.0))
        .collect();
    fft_input.resize(fft_size, Complex::new(0.0, 0.0));

    // Forward FFT
    let mut planner = FftPlanner::new();
    let fft = planner.plan_fft_forward(fft_size);
    fft.process(&mut fft_input);

    // Compute |FFT|²
    for x in &mut fft_input {
        *x = *x * x.conj();
    }

    // Inverse FFT
    let ifft = planner.plan_fft_inverse(fft_size);
    ifft.process(&mut fft_input);

    // Extract real part and normalize by FFT size
    let scale = 1.0 / (fft_size as f32);
    let acf: Vec<f32> = fft_input[..n]
        .iter()
        .map(|x| (x.re * scale).max(0.0)) // Autocorrelation is non-negative
        .collect();

    // Note: Ellis & Pikrakis (2006) mentions optional normalization by signal length
    // (ACF[lag] = ACF[lag] / (n - lag)) for better consistency across signal lengths.
    // However, this normalization can affect peak detection by favoring shorter lags,
    // potentially causing octave errors. Since normalization is optional and the current
    // unnormalized approach works well, we keep it unnormalized for peak detection accuracy.
    // If normalization is needed for specific use cases, it can be added as an optional parameter.

    Ok(acf)
}

/// Find peaks in autocorrelation function
///
/// Finds local maxima with minimum prominence.
///
/// # Arguments
///
/// * `acf_slice` - Slice of ACF to search (already filtered to lag range)
/// * `offset` - Offset to add to indices (lag_min)
///
/// # Returns
///
/// Vector of (lag, value) pairs for detected peaks
fn find_peaks_in_acf(
    acf_slice: &[f32],
    offset: usize,
) -> Result<Vec<(usize, f32)>, crate::error::AnalysisError> {
    if acf_slice.is_empty() {
        return Ok(vec![]);
    }

    let max_value = acf_slice.iter().copied().fold(0.0f32, f32::max);
    if max_value < EPSILON {
        return Ok(vec![]);
    }

    // Minimum prominence: 10% of maximum
    let min_prominence = max_value * 0.1;
    
    // Minimum distance: corresponds to ~5 BPM separation
    // For 120 BPM at 44.1kHz, 512 hop: lag ~= 43
    // 5 BPM difference ≈ 2 lag difference
    let min_distance = 2;

    let mut peaks: Vec<(usize, f32)> = Vec::new();

    // Find local maxima
    for i in 1..(acf_slice.len() - 1) {
        let value = acf_slice[i];
        
        // Check if local maximum
        if value > acf_slice[i - 1] && value > acf_slice[i + 1] {
            // Check prominence (height relative to neighbors)
            let left_val = acf_slice[i - 1];
            let right_val = acf_slice[i + 1];
            let prominence = value - (left_val.max(right_val));
            
            if prominence >= min_prominence {
                // Check minimum distance from previous peak
                let lag = i + offset;
                if peaks.is_empty() || (lag as i32 - peaks.last().unwrap().0 as i32).abs() >= min_distance as i32 {
                    peaks.push((lag, value));
                } else {
                    // Keep the higher peak if too close
                    let last_idx = peaks.len() - 1;
                    if value > peaks[last_idx].1 {
                        peaks[last_idx] = (lag, value);
                    }
                }
            }
        }
    }

    // Sort by value (highest first)
    peaks.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

    Ok(peaks)
}

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

    #[test]
    fn test_autocorrelation_basic_120bpm() {
        // Create onsets for 120 BPM at 44.1kHz, 512 hop
        // 120 BPM = 0.5s per beat = 22050 samples per beat
        // At 512 hop: 22050 / 512 ≈ 43 frames per beat
        let sample_rate = 44100;
        let hop_size = 512;
        let bpm = 120.0;
        let period_samples = (60.0 * sample_rate as f32) / bpm;
        let period_frames = (period_samples / hop_size as f32).round() as usize;

        // Generate 4 beats worth of onsets
        let mut onsets = Vec::new();
        for beat in 0..4 {
            let frame = beat * period_frames;
            let sample = frame * hop_size;
            onsets.push(sample);
        }

        let candidates = estimate_bpm_from_autocorrelation(
            &onsets,
            sample_rate,
            hop_size,
            60.0,
            180.0,
        )
        .unwrap();

        assert!(!candidates.is_empty(), "Should find at least one candidate");
        
        // Best candidate should be close to 120 BPM
        let best = &candidates[0];
        assert!(
            (best.bpm - 120.0).abs() < 5.0,
            "Best BPM should be close to 120, got {:.2}",
            best.bpm
        );
        assert!(best.confidence > 0.0, "Confidence should be positive");
    }

    #[test]
    fn test_autocorrelation_empty_onsets() {
        let result = estimate_bpm_from_autocorrelation(&[], 44100, 512, 60.0, 180.0);
        assert!(result.is_err());
    }

    #[test]
    fn test_autocorrelation_single_onset() {
        let onsets = vec![1000];
        let result = estimate_bpm_from_autocorrelation(&onsets, 44100, 512, 60.0, 180.0);
        // Should return empty (need at least 2 onsets)
        assert!(result.is_ok());
        assert!(result.unwrap().is_empty());
    }

    #[test]
    fn test_autocorrelation_invalid_params() {
        let onsets = vec![1000, 2000];
        
        // Invalid sample rate
        let result = estimate_bpm_from_autocorrelation(&onsets, 0, 512, 60.0, 180.0);
        assert!(result.is_err());
        
        // Invalid hop size
        let result = estimate_bpm_from_autocorrelation(&onsets, 44100, 0, 60.0, 180.0);
        assert!(result.is_err());
        
        // Invalid BPM range
        let result = estimate_bpm_from_autocorrelation(&onsets, 44100, 512, 180.0, 60.0);
        assert!(result.is_err());
    }

    #[test]
    fn test_autocorrelation_128bpm() {
        // Test 128 BPM
        let sample_rate = 44100;
        let hop_size = 512;
        let bpm = 128.0;
        let period_samples = (60.0 * sample_rate as f32) / bpm;
        let period_frames = (period_samples / hop_size as f32).round() as usize;

        let mut onsets = Vec::new();
        for beat in 0..4 {
            let frame = beat * period_frames;
            let sample = frame * hop_size;
            onsets.push(sample);
        }

        let candidates = estimate_bpm_from_autocorrelation(
            &onsets,
            sample_rate,
            hop_size,
            60.0,
            180.0,
        )
        .unwrap();

        assert!(!candidates.is_empty());
        let best = &candidates[0];
        assert!(
            (best.bpm - 128.0).abs() < 5.0,
            "Best BPM should be close to 128, got {:.2}",
            best.bpm
        );
    }

    #[test]
    fn test_compute_autocorrelation_fft() {
        // Simple periodic signal: [1, 0, 1, 0, 1, 0]
        let signal = vec![1.0, 0.0, 1.0, 0.0, 1.0, 0.0];
        let acf = compute_autocorrelation_fft(&signal).unwrap();

        assert_eq!(acf.len(), signal.len());
        
        // ACF[0] should be maximum (self-correlation)
        assert!(acf[0] > 0.0);
        
        // ACF should be symmetric (approximately)
        // ACF[2] should be high (period of 2)
        if acf.len() > 2 {
            assert!(acf[2] > 0.0);
        }
    }

    #[test]
    fn test_find_peaks_in_acf() {
        // Create ACF with clear peaks
        let acf = vec![0.1, 0.2, 0.5, 0.3, 0.4, 0.6, 0.2, 0.1];
        let peaks = find_peaks_in_acf(&acf, 0).unwrap();

        assert!(!peaks.is_empty());
        // Should find peak at index 2 (value 0.5) and index 5 (value 0.6)
        assert!(peaks.iter().any(|(idx, _)| *idx == 2 || *idx == 5));
    }
}