rosa 0.1.0

//! Harmonic analysis: interpolation, salience, and tonnetz.

use crate::matrix::Matrix;

/// Parameters for harmonic salience.
pub struct SalienceParams {
    /// Harmonic numbers to include (e.g. `[1.0, 2.0, 3.0]`).
    pub harmonics: Vec<f64>,
    /// Per-harmonic weights; None = uniform weights.
    pub weights: Option<Vec<f64>>,
    /// If true, only retain values at local spectral peaks.
    pub filter_peaks: bool,
    /// Value used to fill non-peak positions when `filter_peaks` is true.
    pub fill_value: f64,
}

impl Default for SalienceParams {
    fn default() -> Self {
        Self {
            harmonics: vec![1.0, 2.0, 3.0, 4.0, 5.0],
            weights: None,
            filter_peaks: true,
            fill_value: f64::NAN,
        }
    }
}

/// 1D linear interpolation over sorted grid.
///
/// For each x_new value, finds bracketing interval in x_grid
/// and linearly interpolates y_grid. Values outside range get fill_value.
fn interp1d_linear(x_grid: &[f64], y_grid: &[f64], x_new: &[f64], fill_value: f64) -> Vec<f64> {
    let n = x_grid.len();
    x_new
        .iter()
        .map(|&x| {
            if n == 0 || x < x_grid[0] || x > x_grid[n - 1] {
                return fill_value;
            }
            // Binary search for bracketing interval
            let pos = x_grid.partition_point(|&v| v < x);
            if pos == 0 {
                return y_grid[0];
            }
            if pos >= n {
                return y_grid[n - 1];
            }
            // Linear interpolation between pos-1 and pos
            let x0 = x_grid[pos - 1];
            let x1 = x_grid[pos];
            let dx = x1 - x0;
            if dx.abs() < f64::MIN_POSITIVE {
                return y_grid[pos - 1];
            }
            let t = (x - x0) / dx;
            y_grid[pos - 1] * (1.0 - t) + y_grid[pos] * t
        })
        .collect()
}

/// Interpolate spectrogram at harmonic frequencies.
///
/// Matches `librosa.interp_harmonics(x, freqs=freqs, harmonics=harmonics)`.
/// S has shape (n_freqs, n_frames), freqs has length n_freqs.
/// Returns `Vec<Matrix>`, one per harmonic, each of shape `(n_freqs, n_frames)`.
pub fn interp_harmonics(
    s: &Matrix,
    freqs: &[f64],
    harmonics: &[f64],
    fill_value: f64,
) -> Vec<Matrix> {
    let n_freqs = s.rows();
    let n_frames = s.cols();

    harmonics
        .iter()
        .map(|&h| {
            let target_freqs: Vec<f64> = freqs.iter().map(|&f| h * f).collect();
            let mut data = vec![0.0f64; n_freqs * n_frames];

            for t in 0..n_frames {
                // Extract column
                let col: Vec<f64> = (0..n_freqs).map(|r| s.get(r, t)).collect();
                let interpolated = interp1d_linear(freqs, &col, &target_freqs, fill_value);
                for r in 0..n_freqs {
                    data[r * n_frames + t] = interpolated[r];
                }
            }

            Matrix::from_vec(data, n_freqs, n_frames)
        })
        .collect()
}

/// Extract spectrogram energy at harmonics of a time-varying f0.
///
/// Matches `librosa.f0_harmonics(x, f0=f0, freqs=freqs, harmonics=harmonics)`.
/// Returns Matrix of shape (n_harmonics, n_frames).
pub fn f0_harmonics(
    s: &Matrix,
    f0: &[f64],
    freqs: &[f64],
    harmonics: &[f64],
    fill_value: f64,
) -> Matrix {
    let n_freqs = s.rows();
    let n_frames = s.cols();
    let n_harmonics = harmonics.len();

    assert_eq!(f0.len(), n_frames, "f0 length must match number of frames");

    let mut data = vec![fill_value; n_harmonics * n_frames];

    for t in 0..n_frames {
        if f0[t].is_nan() || f0[t] <= 0.0 {
            continue;
        }
        let col: Vec<f64> = (0..n_freqs).map(|r| s.get(r, t)).collect();

        for (h_idx, &h) in harmonics.iter().enumerate() {
            let target = f0[t] * h;
            let result = interp1d_linear(freqs, &col, &[target], fill_value);
            data[h_idx * n_frames + t] = result[0];
        }
    }

    Matrix::from_vec(data, n_harmonics, n_frames)
}

/// Harmonic salience function.
///
/// Matches `librosa.salience(S, freqs=freqs, harmonics=..., ...)`.
/// Returns Matrix of shape (n_freqs, n_frames).
pub fn salience(s: &Matrix, freqs: &[f64], params: &SalienceParams) -> Matrix {
    let n_freqs = s.rows();
    let n_frames = s.cols();

    // Interpolate at harmonic frequencies
    let harm_data = interp_harmonics(s, freqs, &params.harmonics, 0.0);

    // Compute weights
    let n_harmonics = params.harmonics.len();
    let weights = match &params.weights {
        Some(w) => w.clone(),
        None => vec![1.0 / n_harmonics as f64; n_harmonics],
    };
    let weight_sum: f64 = weights.iter().sum();

    // Weighted average across harmonics
    let mut sal_data = vec![0.0f64; n_freqs * n_frames];
    for f in 0..n_freqs {
        for t in 0..n_frames {
            let mut sum = 0.0;
            for (h, w) in weights.iter().enumerate() {
                sum += w * harm_data[h].get(f, t);
            }
            sal_data[f * n_frames + t] = sum / weight_sum;
        }
    }
    let mut result = Matrix::from_vec(sal_data, n_freqs, n_frames);

    // Filter peaks: only keep values at local maxima along frequency axis
    if params.filter_peaks {
        for t in 0..n_frames {
            let col: Vec<f64> = (0..n_freqs).map(|f| result.get(f, t)).collect();
            let peaks = crate::dsp::util::localmax(&col);
            for f in 0..n_freqs {
                if !peaks[f] {
                    *result.get_mut(f, t) = params.fill_value;
                }
            }
        }
    }

    result
}

/// Tonal centroid features (6D pitch space).
///
/// Matches `librosa.feature.tonnetz(y=y, sr=sr, chroma=chroma)`.
/// Returns Matrix of shape (6, n_frames).
pub fn tonnetz(y: Option<&[f64]>, sr: f64, chroma: Option<&Matrix>) -> Matrix {
    // If no chroma provided, compute it
    let owned_chroma;
    let chroma = match chroma {
        Some(c) => c,
        None => {
            let y = y.expect("tonnetz: either y or chroma must be provided");
            let params = super::chroma::ChromaCqtParams {
                sr,
                ..Default::default()
            };
            owned_chroma = super::chroma::chroma_cqt(y, &params);
            &owned_chroma
        }
    };

    let n_chroma = chroma.rows();

    // L1-normalize each column
    use crate::matrix::Axis;
    let chroma_norm = chroma.normalize_axis(Axis::Row, 1.0);

    // Tonnetz dimensions:
    // Row 0: sin(7π/6 * k')  where k' = 12*k/n_chroma in [0, 12)
    // Row 1: cos(7π/6 * k')
    // Row 2: sin(3π/2 * k')
    // Row 3: cos(3π/2 * k')
    // Row 4: 0.5 * sin(2π/3 * k')
    // Row 5: 0.5 * cos(2π/3 * k')
    //
    // But actually librosa uses a different formulation:
    // angles_per_chroma = [7/6, 7/6, 3/2, 3/2, 2/3, 2/3]
    // radii = [1, 1, 1, 1, 0.5, 0.5]
    // V[i][j] = angles[i] * dim_map[j]
    // Even rows: V -= 0.5  (which turns cos into sin via cos(pi*(x-0.5)) = sin(pi*x))
    // phi[i][j] = radii[i] * cos(pi * V[i][j])

    let angles = [
        7.0 / 6.0,
        7.0 / 6.0,
        3.0 / 2.0,
        3.0 / 2.0,
        2.0 / 3.0,
        2.0 / 3.0,
    ];
    let radii = [1.0, 1.0, 1.0, 1.0, 0.5, 0.5];

    let mut phi_data = vec![0.0f64; 6 * n_chroma];
    for i in 0..6 {
        for j in 0..n_chroma {
            let mut v = angles[i] * (j as f64) * (12.0 / n_chroma as f64);
            // Even rows (0, 2, 4) get the -0.5 offset
            if i % 2 == 0 {
                v -= 0.5;
            }
            phi_data[i * n_chroma + j] = radii[i] * (std::f64::consts::PI * v).cos();
        }
    }
    let phi = Matrix::from_vec(phi_data, 6, n_chroma);

    // tonnetz = phi @ chroma_normalized
    phi.matmul(&chroma_norm)
}

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

    #[test]
    fn test_interp1d_linear_basic() {
        let x = vec![0.0, 1.0, 2.0, 3.0];
        let y = vec![0.0, 10.0, 20.0, 30.0];
        let x_new = vec![0.5, 1.5, 2.5];
        let result = interp1d_linear(&x, &y, &x_new, -1.0);
        assert!((result[0] - 5.0).abs() < 1e-12);
        assert!((result[1] - 15.0).abs() < 1e-12);
        assert!((result[2] - 25.0).abs() < 1e-12);
    }

    #[test]
    fn test_interp1d_linear_out_of_bounds() {
        let x = vec![1.0, 2.0, 3.0];
        let y = vec![10.0, 20.0, 30.0];
        let x_new = vec![0.0, 4.0];
        let result = interp1d_linear(&x, &y, &x_new, -999.0);
        assert_eq!(result[0], -999.0);
        assert_eq!(result[1], -999.0);
    }

    #[test]
    fn test_interp_harmonics_shape() {
        let s = Matrix::from_vec(vec![1.0; 10 * 5], 10, 5);
        let freqs: Vec<f64> = (0..10).map(|i| 100.0 + i as f64 * 100.0).collect();
        let harmonics = vec![1.0, 2.0, 3.0];
        let result = interp_harmonics(&s, &freqs, &harmonics, 0.0);
        assert_eq!(result.len(), 3);
        for m in &result {
            assert_eq!(m.rows(), 10);
            assert_eq!(m.cols(), 5);
        }
    }

    #[test]
    fn test_tonnetz_shape() {
        // Create a simple 12x5 chroma matrix
        let mut data = vec![0.0; 12 * 5];
        // Put some energy in the first chroma bin
        for f in 0..5 {
            data[0 * 5 + f] = 1.0;
        }
        let chroma = Matrix::from_vec(data, 12, 5);
        let result = tonnetz(None, 22050.0, Some(&chroma));
        assert_eq!(result.rows(), 6);
        assert_eq!(result.cols(), 5);
    }
}