use ndarray::{Array1, Array2, Axis};
use crate::signal_processing::time_frequency::stft;
use thiserror::Error;
#[derive(Debug, Clone)]
pub struct TempoBuilder<'a> {
y: &'a [f32],
sr: u32,
hop_length: usize,
win_length: usize,
}
impl<'a> TempoBuilder<'a> {
pub fn hop_length(mut self, hop_length: usize) -> Self {
self.hop_length = hop_length;
self
}
pub fn win_length(mut self, win_length: usize) -> Self {
self.win_length = win_length;
self
}
pub fn compute(self) -> Result<f32, RhythmError> {
tempo_impl(Some(self.y), Some(self.sr), None, Some(self.hop_length))
}
}
pub fn tempo(y: &[f32], sr: u32) -> TempoBuilder<'_> {
TempoBuilder {
y,
sr,
hop_length: 512,
win_length: 2048,
}
}
#[derive(Error, Debug)]
pub enum RhythmError {
#[error("Invalid input: {0}")]
InvalidInput(String),
#[error("Computation failed: {0}")]
ComputationFailed(String),
}
pub fn tempo_impl(
y: Option<&[f32]>,
sr: Option<u32>,
onset_envelope: Option<&Array1<f32>>,
hop_length: Option<usize>,
) -> Result<f32, RhythmError> {
let sr = sr.unwrap_or(44100);
let hop = hop_length.unwrap_or(512);
let onset_owned = if onset_envelope.is_none() {
let y = y.ok_or(RhythmError::InvalidInput("Audio signal required when onset_envelope is None".to_string()))?;
let s = stft(y)
.hop_length(hop)
.compute()
.map_err(|e| RhythmError::ComputationFailed(format!("STFT computation failed: {}", e)))?
.mapv(|x| x.norm());
s.map_axis(Axis(0), |row| row.iter().map(|&x| x.max(0.0)).sum::<f32>())
} else {
onset_envelope.unwrap().to_owned()
};
let onset = &onset_owned;
let tempogram = tempogram_impl(None, Some(sr), Some(onset), hop_length, None)?;
Ok(tempogram.axis_iter(Axis(1)).map(|col| {
let max_val = col.iter().max_by(|a, b| a.partial_cmp(b).unwrap()).unwrap_or(&0.0);
let max_idx = col.iter().position(|&x| x == *max_val).unwrap_or(0);
crate::utils::frequency::tempo_frequencies(tempogram.shape()[0], Some(hop), Some(sr))[max_idx]
}).sum::<f32>() / tempogram.shape()[1] as f32)
}
pub fn tempogram_impl(
y: Option<&[f32]>,
sr: Option<u32>,
onset_envelope: Option<&Array1<f32>>,
hop_length: Option<usize>,
win_length: Option<usize>,
) -> Result<Array2<f32>, RhythmError> {
let _sr = sr.unwrap_or(44100);
let hop = hop_length.unwrap_or(512);
let win = win_length.unwrap_or(384);
let onset_owned = if onset_envelope.is_none() {
let y = y.ok_or(RhythmError::InvalidInput("Audio signal required when onset_envelope is None".to_string()))?;
let s = stft(y)
.hop_length(hop)
.compute()
.map_err(|e| RhythmError::ComputationFailed(format!("STFT computation failed: {}", e)))?
.mapv(|x| x.norm());
s.map_axis(Axis(0), |row| row.iter().map(|&x| x.max(0.0)).sum::<f32>())
} else {
onset_envelope.unwrap().to_owned()
};
let onset = &onset_owned;
let mut tempogram = Array2::zeros((win / 2 + 1, onset.len()));
for t in 0..onset.len() {
for lag in 0..(win / 2 + 1) {
let past = (t as isize - lag as isize).max(0) as usize;
tempogram[[lag, t]] = onset[t] * onset[past];
}
}
Ok(tempogram)
}
pub fn tempogram_ratio_impl(
y: Option<&[f32]>,
sr: Option<u32>,
onset_envelope: Option<&Array1<f32>>,
hop_length: Option<usize>,
ratios: Option<&[f32]>,
) -> Result<Array2<f32>, RhythmError> {
let tempogram = tempogram_impl(y, sr, onset_envelope, hop_length, None)?;
let ratios = ratios.unwrap_or(&[2.0, 3.0, 4.0]);
let mut ratio_map = Array2::zeros((ratios.len(), tempogram.shape()[1]));
for (r_idx, &r) in ratios.iter().enumerate() {
for t in 0..tempogram.shape()[1] {
let mut sum = 0.0;
for f in 0..tempogram.shape()[0] {
let target_f = f as f32 * r;
let bin = target_f.round() as usize;
if bin < tempogram.shape()[0] {
sum += tempogram[[bin, t]];
}
}
ratio_map[[r_idx, t]] = sum;
}
}
Ok(ratio_map)
}