rosa 0.1.0

//! Dynamic programming beat tracker matching `librosa.beat.beat_track`.
//!
//! Implements Ellis (2007) "Beat tracking by dynamic programming."

use crate::dsp::util::tiny;

/// Parameters for beat tracking.
pub struct BeatTrackParams {
    /// Sample rate in Hz.
    pub sr: f64,
    /// Number of samples between frames.
    pub hop_length: usize,
    /// Initial tempo estimate in BPM for tempo estimation. Default: 120.0
    pub start_bpm: f64,
    /// Rigidity of beat spacing. Higher values enforce stricter tempo. Default: 100.0
    pub tightness: f64,
    /// If true, trim leading/trailing weak beats. Default: true
    pub trim: bool,
    /// If Some, use this BPM directly instead of estimating.
    pub bpm: Option<f64>,
}

impl BeatTrackParams {
    /// Create `BeatTrackParams` with the given sample rate and default values.
    pub fn new(sr: f64) -> Self {
        Self {
            sr,
            hop_length: 512,
            start_bpm: 120.0,
            tightness: 100.0,
            trim: true,
            bpm: None,
        }
    }
}

/// Dynamic programming beat tracker.
///
/// Returns `(estimated_tempo_bpm, beat_frame_indices)`.
/// Matches `librosa.beat.beat_track`.
pub fn beat_track(onset_envelope: &[f64], params: &BeatTrackParams) -> (f64, Vec<usize>) {
    let n = onset_envelope.len();
    if n == 0 {
        return (params.start_bpm, Vec::new());
    }

    // Stage 1: Estimate tempo if not provided
    let bpm = if let Some(bpm) = params.bpm {
        bpm
    } else {
        let tempo_params = super::tempo::TempoParams {
            sr: params.sr,
            hop_length: params.hop_length,
            start_bpm: params.start_bpm,
            std_bpm: 1.0,
            ac_size: 8.0,
            max_tempo: Some(320.0),
        };
        super::tempo::tempo(onset_envelope, &tempo_params)
    };

    // Stage 2: Compute frames per beat
    let frame_rate = params.sr / params.hop_length as f64;
    let fpb = (frame_rate * 60.0 / bpm).round() as usize;

    if fpb == 0 {
        return (bpm, Vec::new());
    }

    // Stage 3: Normalize onsets
    let normalized = normalize_onsets(onset_envelope);

    // Stage 4: Compute local score (Gaussian-windowed onset strength)
    let localscore = beat_local_score(&normalized, fpb);

    // Stage 5: Dynamic programming
    let (backlink, cumscore) = beat_track_dp(&localscore, fpb, params.tightness);

    // Stage 6: Backtrack
    let mut beats = dp_backtrack(&backlink, &cumscore);

    // Stage 7: Trim
    if params.trim {
        beats = trim_beats(&localscore, &beats);
    }

    (bpm, beats)
}

/// Normalize onset envelope by sample standard deviation.
fn normalize_onsets(onsets: &[f64]) -> Vec<f64> {
    let n = onsets.len();
    if n <= 1 {
        return onsets.to_vec();
    }

    let mean = onsets.iter().sum::<f64>() / n as f64;
    let var = onsets.iter().map(|&x| (x - mean).powi(2)).sum::<f64>() / (n - 1) as f64;
    let std = var.sqrt();

    let denom = std + tiny();
    onsets.iter().map(|&x| x / denom).collect()
}

/// Compute local score: convolve onset envelope with Gaussian window.
///
/// The Gaussian has sigma = fpb/32, spanning ±fpb samples.
fn beat_local_score(onset_envelope: &[f64], fpb: usize) -> Vec<f64> {
    let n = onset_envelope.len();
    if n == 0 {
        return Vec::new();
    }

    // Build Gaussian kernel: exp(-0.5 * (x * 32/fpb)^2) for x in [-fpb, fpb]
    let kernel_half = fpb;
    let kernel_len = 2 * kernel_half + 1;
    let mut kernel = Vec::with_capacity(kernel_len);
    let scale = 32.0 / fpb as f64;

    for i in 0..kernel_len {
        let x = i as f64 - kernel_half as f64;
        kernel.push((-0.5 * (x * scale).powi(2)).exp());
    }

    // Normalize kernel
    let ksum: f64 = kernel.iter().sum();
    if ksum > 0.0 {
        for k in kernel.iter_mut() {
            *k /= ksum;
        }
    }

    // Convolve (mode='same')
    let mut result = vec![0.0; n];
    for i in 0..n {
        let mut sum = 0.0;
        for (j, &k) in kernel.iter().enumerate() {
            let idx = i as isize + j as isize - kernel_half as isize;
            if idx >= 0 && (idx as usize) < n {
                sum += onset_envelope[idx as usize] * k;
            }
        }
        result[i] = sum;
    }

    result
}

/// Core DP beat tracking.
///
/// Returns (backlink, cumscore).
fn beat_track_dp(localscore: &[f64], fpb: usize, tightness: f64) -> (Vec<i64>, Vec<f64>) {
    let n = localscore.len();
    let mut cumscore = vec![0.0; n];
    let mut backlink = vec![-1i64; n];

    if n == 0 {
        return (backlink, cumscore);
    }

    let max_localscore = localscore.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
    let score_thresh = 0.01 * max_localscore;
    let log_fpb = (fpb as f64).ln();

    cumscore[0] = localscore[0];
    let mut first_beat = true;

    for i in 1..n {
        // Search backward for best predecessor
        // Range: [i - 2*fpb, i - fpb/2)
        let search_start = i.saturating_sub(2 * fpb);
        let search_end = i.saturating_sub(fpb / 2);

        let mut best_score = f64::NEG_INFINITY;
        let mut best_loc: i64 = -1;

        for loc in search_start..search_end {
            let interval = (i - loc) as f64;
            let penalty = tightness * (interval.ln() - log_fpb).powi(2);
            let score = cumscore[loc] - penalty;
            if score > best_score {
                best_score = score;
                best_loc = loc as i64;
            }
        }

        if best_loc >= 0 {
            cumscore[i] = best_score + localscore[i];
        } else {
            cumscore[i] = localscore[i];
        }

        if first_beat && localscore[i] < score_thresh {
            backlink[i] = -1;
        } else {
            backlink[i] = best_loc;
            if best_loc >= 0 {
                first_beat = false;
            }
        }
    }

    (backlink, cumscore)
}

/// Find the last beat position and backtrack to recover beat sequence.
fn dp_backtrack(backlink: &[i64], cumscore: &[f64]) -> Vec<usize> {
    let n = cumscore.len();
    if n == 0 {
        return Vec::new();
    }

    // Find last beat: last local max of cumscore above threshold
    // Find local maxima of cumscore
    let maxima: Vec<usize> = (1..n)
        .filter(|&i| {
            let prev = cumscore[i] > cumscore[i - 1];
            let next = if i + 1 < n {
                cumscore[i] >= cumscore[i + 1]
            } else {
                true
            };
            prev && next
        })
        .collect();

    if maxima.is_empty() {
        return Vec::new();
    }

    // Find median of local max values
    let mut max_values: Vec<f64> = maxima.iter().map(|&i| cumscore[i]).collect();
    max_values.sort_by(|a, b| a.total_cmp(b));
    let median_val = if max_values.len().is_multiple_of(2) {
        (max_values[max_values.len() / 2 - 1] + max_values[max_values.len() / 2]) / 2.0
    } else {
        max_values[max_values.len() / 2]
    };

    // Last beat is last local max above 0.5 * median
    let threshold = 0.5 * median_val;
    let tail = match maxima.iter().rev().find(|&&i| cumscore[i] >= threshold) {
        Some(&i) => i,
        None => return Vec::new(),
    };

    // Backtrack
    let mut beats = Vec::new();
    let mut pos = tail as i64;
    while pos >= 0 {
        beats.push(pos as usize);
        pos = backlink[pos as usize];
    }
    beats.reverse();

    beats
}

/// Remove weak leading/trailing beats.
fn trim_beats(localscore: &[f64], beats: &[usize]) -> Vec<usize> {
    if beats.is_empty() || localscore.is_empty() {
        return beats.to_vec();
    }

    // Threshold: 0.5 * RMS of localscore
    let rms = (localscore.iter().map(|x| x * x).sum::<f64>() / localscore.len() as f64).sqrt();
    let threshold = 0.5 * rms;

    // Find first beat above threshold
    let start = beats
        .iter()
        .position(|&b| localscore[b] >= threshold)
        .unwrap_or(beats.len());

    // Find last beat above threshold
    let end = beats
        .iter()
        .rposition(|&b| localscore[b] >= threshold)
        .map(|i| i + 1)
        .unwrap_or(0);

    if start < end {
        beats[start..end].to_vec()
    } else {
        Vec::new()
    }
}