melpe-rs 0.1.2

/// Pitch detection for MELPe
///
/// Normalized autocorrelation pitch estimator with:
/// - Candidate peak picking
/// - Sub-multiple rejection (octave error avoidance)
/// - Parabolic interpolation for fractional pitch

use crate::core_types::{PITCH_MIN, PITCH_MAX, FRAME_SAMPLES, SAMPLE_RATE};
use crate::math::{sqrtf, fabsf, clampf, fmaxf};

/// Result of pitch analysis for one frame
#[derive(Debug, Clone, Copy)]
pub struct PitchResult {
    /// Pitch period in samples (fractional)
    pub period: f32,
    /// Pitch frequency in Hz
    pub frequency: f32,
    /// Normalized correlation at the chosen lag (0.0..1.0)
    pub correlation: f32,
    /// Whether the frame is likely voiced
    pub voiced: bool,
}

/// Voicing threshold for normalized autocorrelation
const VOICING_THRESHOLD: f32 = 0.3;

/// Threshold for preferring a sub-multiple pitch (octave correction)
const SUBMULTIPLE_RATIO: f32 = 0.85;

/// Compute normalized autocorrelation for lag range [lag_min, lag_max].
/// Returns vector of (lag, normalized_correlation) pairs.
fn normalized_autocorrelation(
    signal: &[f32],
    lag_min: usize,
    lag_max: usize,
) -> ([f32; PITCH_MAX + 1], usize, usize) {
    let n = signal.len();
    let lag_max = lag_max.min(n / 2).min(PITCH_MAX);
    let lag_min = lag_min.max(1);

    let mut corr = [0.0f32; PITCH_MAX + 1];

    // Energy of the original window
    let mut energy0 = 0.0f32;
    for i in 0..n.saturating_sub(lag_max) {
        energy0 += signal[i] * signal[i];
    }

    if energy0 < 1e-10 {
        return (corr, lag_min, lag_max);
    }

    for lag in lag_min..=lag_max {
        let mut num = 0.0f32;
        let mut energy_lag = 0.0f32;
        let len = n - lag;

        for i in 0..len {
            num += signal[i] * signal[i + lag];
            energy_lag += signal[i + lag] * signal[i + lag];
        }

        let denom = sqrtf(energy0 * energy_lag);
        if denom > 1e-10 {
            corr[lag] = num / denom;
        }
    }

    (corr, lag_min, lag_max)
}

/// Pick the best pitch candidate from correlation values.
/// Handles sub-multiple rejection to avoid octave errors.
fn pick_best_candidate(
    corr: &[f32; PITCH_MAX + 1],
    lag_min: usize,
    lag_max: usize,
) -> (usize, f32) {
    // Find all local maxima
    let mut best_lag = lag_min;
    let mut best_val = corr[lag_min];

    for lag in (lag_min + 1)..lag_max {
        if corr[lag] > corr[lag - 1] && corr[lag] >= corr[lag + 1] {
            // Local peak
            if corr[lag] > best_val {
                best_val = corr[lag];
                best_lag = lag;
            }
        }
    }

    // Sub-multiple check: prefer shorter period if its correlation is close
    // This prevents doubling errors (picking 2T instead of T)
    if best_lag > lag_min * 2 {
        let half = best_lag / 2;
        if half >= lag_min && corr[half] > best_val * SUBMULTIPLE_RATIO {
            best_lag = half;
            best_val = corr[half];
        }

        let third = best_lag / 3;
        if third >= lag_min && corr[third] > best_val * SUBMULTIPLE_RATIO {
            best_lag = third;
            best_val = corr[third];
        }
    }

    (best_lag, best_val)
}

/// Parabolic interpolation around an integer lag for fractional pitch.
/// Given corr[lag-1], corr[lag], corr[lag+1], returns the refined lag.
fn parabolic_interp(
    corr: &[f32; PITCH_MAX + 1],
    lag: usize,
    lag_min: usize,
    lag_max: usize,
) -> f32 {
    if lag <= lag_min || lag >= lag_max {
        return lag as f32;
    }

    let left = corr[lag - 1];
    let center = corr[lag];
    let right = corr[lag + 1];

    let denom = 2.0 * (2.0 * center - left - right);
    if fabsf(denom) < 1e-10 {
        return lag as f32;
    }

    let offset = (right - left) / denom;
    lag as f32 + clampf(offset, -0.5, 0.5)
}

/// Estimate pitch for a frame of audio samples.
///
/// `signal` should be one frame (~180 samples at 8kHz).
/// Optionally pass a center-clipping threshold to improve pitch detection
/// in noisy conditions (set to 0.0 to disable).
pub fn estimate_pitch(signal: &[f32], clip_threshold: f32) -> PitchResult {
    let n = signal.len();

    // Center clipping for noise robustness
    let mut clipped = [0.0f32; FRAME_SAMPLES];
    let len = n.min(FRAME_SAMPLES);
    for i in 0..len {
        if fabsf(signal[i]) > clip_threshold {
            clipped[i] = signal[i];
        }
    }

    let buf = if clip_threshold > 0.0 {
        &clipped[..len]
    } else {
        &signal[..len]
    };

    let (corr, lag_min, lag_max) = normalized_autocorrelation(buf, PITCH_MIN, PITCH_MAX);
    let (best_lag, best_corr) = pick_best_candidate(&corr, lag_min, lag_max);

    let voiced = best_corr >= VOICING_THRESHOLD;
    let refined_period = parabolic_interp(&corr, best_lag, lag_min, lag_max);
    let frequency = if refined_period > 0.0 {
        SAMPLE_RATE as f32 / refined_period
    } else {
        0.0
    };

    PitchResult {
        period: refined_period,
        frequency,
        correlation: best_corr,
        voiced,
    }
}

/// Compute adaptive center-clipping threshold from signal energy.
/// Returns a threshold at `fraction` of the peak amplitude.
pub fn adaptive_clip_threshold(signal: &[f32], fraction: f32) -> f32 {
    let peak = signal.iter().fold(0.0f32, |m, &s| fmaxf(m, fabsf(s)));
    peak * fraction
}

#[cfg(all(test, feature = "std"))]
mod tests {
    use super::*;

    /// Generate a sine wave at a given frequency
    fn sine_wave(freq: f32, n: usize) -> Vec<f32> {
        (0..n)
            .map(|i| {
                (2.0 * core::f32::consts::PI * freq * i as f32 / SAMPLE_RATE as f32).sin()
            })
            .collect()
    }

    #[test]
    fn test_pitch_100hz() {
        let signal = sine_wave(100.0, FRAME_SAMPLES);
        let result = estimate_pitch(&signal, 0.0);

        assert!(result.voiced, "100 Hz sine should be voiced");
        // Period should be ~80 samples (8000/100)
        assert!(
            (result.period - 80.0).abs() < 2.0,
            "Expected period ~80, got {}",
            result.period
        );
        assert!(
            (result.frequency - 100.0).abs() < 3.0,
            "Expected ~100 Hz, got {}",
            result.frequency
        );
    }

    #[test]
    fn test_pitch_200hz() {
        let signal = sine_wave(200.0, FRAME_SAMPLES);
        let result = estimate_pitch(&signal, 0.0);

        assert!(result.voiced);
        // Period ~40 samples
        assert!(
            (result.period - 40.0).abs() < 2.0,
            "Expected period ~40, got {}",
            result.period
        );
    }

    #[test]
    fn test_pitch_250hz() {
        let signal = sine_wave(250.0, FRAME_SAMPLES);
        let result = estimate_pitch(&signal, 0.0);

        assert!(result.voiced);
        // Period ~32 samples
        assert!(
            (result.period - 32.0).abs() < 2.0,
            "Expected period ~32, got {}",
            result.period
        );
    }

    #[test]
    fn test_silence_unvoiced() {
        let signal = [0.0f32; FRAME_SAMPLES];
        let result = estimate_pitch(&signal, 0.0);
        assert!(!result.voiced, "Silence should be unvoiced");
    }

    #[test]
    fn test_noise_unvoiced() {
        // Pseudo-random noise using simple LCG
        let mut noise = [0.0f32; FRAME_SAMPLES];
        let mut seed: u32 = 12345;
        for s in noise.iter_mut() {
            seed = seed.wrapping_mul(1103515245).wrapping_add(12345);
            *s = (seed as f32 / u32::MAX as f32) * 2.0 - 1.0;
        }
        let result = estimate_pitch(&noise, 0.0);
        // Noise correlation should be low
        assert!(
            result.correlation < 0.5,
            "Noise correlation should be low, got {}",
            result.correlation
        );
    }

    #[test]
    fn test_center_clipping() {
        let signal = sine_wave(150.0, FRAME_SAMPLES);
        let thresh = adaptive_clip_threshold(&signal, 0.3);
        assert!(thresh > 0.0);
        assert!(thresh < 1.0);

        let result = estimate_pitch(&signal, thresh);
        // Should still detect pitch even with clipping
        assert!(result.voiced);
        assert!(
            (result.frequency - 150.0).abs() < 10.0,
            "Expected ~150 Hz with clipping, got {}",
            result.frequency
        );
    }

    #[test]
    fn test_parabolic_refinement() {
        // Verify that fractional pitch is more accurate than integer
        let freq = 123.456;
        let expected_period = SAMPLE_RATE as f32 / freq;
        let signal = sine_wave(freq, FRAME_SAMPLES);
        let result = estimate_pitch(&signal, 0.0);

        let integer_error = (result.period.round() - expected_period).abs();
        let refined_error = (result.period - expected_period).abs();
        assert!(
            refined_error <= integer_error + 0.1,
            "Parabolic refinement should help: int_err={}, refined_err={}",
            integer_error,
            refined_error
        );
    }
}