use crate::dsp::util::peak_pick;
pub fn onset_detect(
onset_envelope: &[f64],
sr: f64,
hop_length: usize,
normalize: bool,
backtrack: bool,
) -> Vec<usize> {
if onset_envelope.is_empty() {
return Vec::new();
}
let env = if normalize {
let min_val = onset_envelope.iter().copied().fold(f64::INFINITY, f64::min);
let shifted: Vec<f64> = onset_envelope.iter().map(|&v| v - min_val).collect();
let max_val = shifted.iter().copied().fold(0.0_f64, f64::max);
let tiny = (f64::EPSILON * max_val.abs()).sqrt().max(f64::MIN_POSITIVE);
shifted
.iter()
.map(|&v| v / (max_val + tiny))
.collect::<Vec<_>>()
} else {
onset_envelope.to_vec()
};
if !env.iter().any(|&v| v != 0.0) || !env.iter().all(|v| v.is_finite()) {
return Vec::new();
}
let pre_max = ((0.03 * sr / hop_length as f64).floor() as usize).max(1);
let post_max = ((0.00 * sr / hop_length as f64).floor() as usize) + 1;
let pre_avg = ((0.10 * sr / hop_length as f64).floor() as usize).max(1);
let post_avg = ((0.10 * sr / hop_length as f64).floor() as usize) + 1;
let wait = ((0.03 * sr / hop_length as f64).floor() as usize).max(1);
let delta = 0.07;
let mut onsets = peak_pick(&env, pre_max, post_max, pre_avg, post_avg, delta, wait);
if backtrack {
onsets = onset_backtrack(&onsets, &env);
}
onsets
}
pub fn onset_backtrack(events: &[usize], energy: &[f64]) -> Vec<usize> {
if events.is_empty() || energy.is_empty() {
return events.to_vec();
}
let n = energy.len();
let mut minima_mask = vec![false; n];
for i in 1..n.saturating_sub(1) {
minima_mask[i] = energy[i] <= energy[i - 1] && energy[i] < energy[i + 1];
}
let mut minima: Vec<usize> = minima_mask
.iter()
.enumerate()
.filter(|&(_, &b)| b)
.map(|(i, _)| i)
.collect();
if minima.is_empty() || minima[0] != 0 {
minima.insert(0, 0);
}
events
.iter()
.map(|&event| {
match minima.binary_search(&event) {
Ok(idx) => minima[idx],
Err(idx) => {
if idx == 0 {
minima[0]
} else {
minima[idx - 1]
}
}
}
})
.collect()
}