rosa 0.1.0

//! Array utilities matching `librosa.util`.

/// Smallest positive normal f64 value.
///
/// Matches `librosa.util.tiny(np.float64)`.
pub fn tiny() -> f64 {
    f64::MIN_POSITIVE
}

/// Center-pad `data` with zeros to `size`.
///
/// Matches `librosa.util.pad_center`. Panics if `size < data.len()`.
pub fn pad_center(data: &[f64], size: usize) -> Vec<f64> {
    let n = data.len();
    assert!(
        size >= n,
        "pad_center: target size ({size}) must be >= data length ({n})"
    );
    if size == n {
        return data.to_vec();
    }
    let lpad = (size - n) / 2;
    let rpad = size - n - lpad;
    let mut result = vec![0.0; lpad];
    result.extend_from_slice(data);
    result.extend(std::iter::repeat_n(0.0, rpad));
    result
}

/// Truncate or zero-pad `data` to exactly `size`.
///
/// Matches `librosa.util.fix_length`.
pub fn fix_length(data: &[f64], size: usize) -> Vec<f64> {
    let n = data.len();
    if n == size {
        data.to_vec()
    } else if n > size {
        data[..size].to_vec()
    } else {
        let mut result = data.to_vec();
        result.resize(size, 0.0);
        result
    }
}

/// Normalization mode for vectors and matrix rows/columns.
///
/// Used by [`normalize`] and `Matrix::normalize_axis`.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Norm {
    /// L-infinity (max absolute value)
    Linf,
    /// L1 (sum of absolute values)
    L1,
    /// L2 (Euclidean norm)
    L2,
}

/// Normalize `data` in-place using the given norm.
///
/// Matches `librosa.util.normalize`. If the norm value is smaller than
/// `tiny()`, the data is left unchanged to avoid division by near-zero.
///
/// # Examples
///
/// ```
/// use rosa::dsp::util::{normalize, Norm};
///
/// let mut data = vec![1.0, 2.0, 3.0];
/// normalize(&mut data, Norm::Linf);
/// assert!((data[2] - 1.0).abs() < 1e-12);
/// ```
pub fn normalize(data: &mut [f64], norm: Norm) {
    let norm_val = match norm {
        Norm::Linf => data.iter().map(|x| x.abs()).fold(0.0_f64, f64::max),
        Norm::L1 => data.iter().map(|x| x.abs()).sum(),
        Norm::L2 => data.iter().map(|x| x * x).sum::<f64>().sqrt(),
    };

    if norm_val > tiny() {
        let inv = 1.0 / norm_val;
        for val in data.iter_mut() {
            *val *= inv;
        }
    }
}

/// Detect local maxima.
///
/// Matches `librosa.util.localmax`. Returns a boolean vector where `true`
/// indicates a local maximum: `x[i] > x[i-1]` (strict) and `x[i] >= x[i+1]`.
/// First element is always `false`. Last element compares only with left neighbor.
pub fn localmax(x: &[f64]) -> Vec<bool> {
    let n = x.len();
    let mut result = vec![false; n];
    for i in 1..n {
        if i < n - 1 {
            result[i] = x[i] > x[i - 1] && x[i] >= x[i + 1];
        } else {
            result[i] = x[i] > x[i - 1];
        }
    }
    result
}

/// Detect local minima.
///
/// Matches `librosa.util.localmin`. Returns a boolean vector where `true`
/// indicates a local minimum: `x[i] < x[i-1]` (strict) and `x[i] <= x[i+1]`.
/// First element is always `false`. Last element compares only with left neighbor.
pub fn localmin(x: &[f64]) -> Vec<bool> {
    let n = x.len();
    let mut result = vec![false; n];
    for i in 1..n {
        if i < n - 1 {
            result[i] = x[i] < x[i - 1] && x[i] <= x[i + 1];
        } else {
            result[i] = x[i] < x[i - 1];
        }
    }
    result
}

/// Wiener-style soft mask.
///
/// Matches `librosa.util.softmask`.
/// `mask[i] = x[i]^power / (x[i]^power + x_ref[i]^power)`
///
/// Uses rescaling by `max(x, x_ref)` for numerical stability.
/// When both values are near zero, returns 0.0.
/// When `power` is infinite, returns a hard binary mask (`x > x_ref`).
pub fn softmask(x: &[f64], x_ref: &[f64], power: f64) -> Vec<f64> {
    assert_eq!(x.len(), x_ref.len(), "softmask: input lengths must match");

    let n = x.len();
    let tiny_val = tiny();
    let mut mask = vec![0.0; n];

    if power.is_infinite() {
        // Hard mask
        for i in 0..n {
            mask[i] = if x[i] > x_ref[i] { 1.0 } else { 0.0 };
        }
        return mask;
    }

    for i in 0..n {
        let z = x[i].max(x_ref[i]);
        if z < tiny_val {
            mask[i] = 0.0;
            continue;
        }
        let xn = (x[i] / z).powf(power);
        let rn = (x_ref[i] / z).powf(power);
        let denom = xn + rn;
        mask[i] = if denom > 0.0 { xn / denom } else { 0.0 };
    }

    mask
}

/// Pad `data` with linear ramps on both sides.
///
/// Left ramp goes from 0 up to `data[0]`. Right ramp goes from `data[last]` down to 0.
/// Matches `np.pad(x, (pad_left, pad_right), mode='linear_ramp', end_values=0)`.
pub fn pad_linear_ramp(data: &[f64], pad_left: usize, pad_right: usize) -> Vec<f64> {
    let n = data.len();
    if n == 0 {
        return vec![0.0; pad_left + pad_right];
    }

    let total = pad_left + n + pad_right;
    let mut result = Vec::with_capacity(total);

    // Left ramp: linearly from 0 to data[0]
    if pad_left > 0 {
        let first = data[0];
        for i in 0..pad_left {
            result.push(first * (i as f64) / (pad_left as f64));
        }
    }

    // Original data
    result.extend_from_slice(data);

    // Right ramp: linearly from data[last] to 0
    if pad_right > 0 {
        let last = data[n - 1];
        for i in 1..=pad_right {
            result.push(last * (1.0 - i as f64 / pad_right as f64));
        }
    }

    result
}

/// 1D maximum filter matching `scipy.ndimage.maximum_filter1d`.
///
/// Replaces each element with the maximum in a centered window of `size`.
/// Uses scipy's default "reflect" boundary mode.
pub fn maximum_filter1d(x: &[f64], size: usize) -> Vec<f64> {
    let n = x.len();
    if n == 0 {
        return Vec::new();
    }
    if size <= 1 {
        return x.to_vec();
    }
    let half = size / 2;
    let mut result = vec![0.0; n];
    for (i, result_val) in result.iter_mut().enumerate().take(n) {
        // scipy uses origin=0 → window is [i - half, i - half + size)
        // with "reflect" boundary: d c b a | a b c d | d c b a
        let mut max_val = f64::NEG_INFINITY;
        for j in 0..size {
            let idx = i as isize - half as isize + j as isize;
            let reflected = if idx < 0 {
                // reflect: -1→0, -2→1, etc.
                (-idx - 1) as usize
            } else if idx >= n as isize {
                // reflect: n→n-1, n+1→n-2, etc.
                (2 * n as isize - 1 - idx) as usize
            } else {
                idx as usize
            };
            let reflected = reflected.min(n - 1);
            max_val = max_val.max(x[reflected]);
        }
        *result_val = max_val;
    }
    result
}

/// General-purpose peak detection matching `librosa.util.peak_pick`.
///
/// A sample `n` is selected as a peak if ALL three conditions are met:
/// 1. `x[n] == max(x[n - pre_max : n + post_max])` — local maximum
/// 2. `x[n] >= mean(x[n - pre_avg : n + post_avg]) + delta` — above local mean
/// 3. `n - previous_n > wait` — sufficient spacing from last peak
///
/// Returns indices of detected peaks.
pub fn peak_pick(
    x: &[f64],
    pre_max: usize,
    post_max: usize,
    pre_avg: usize,
    post_avg: usize,
    delta: f64,
    wait: usize,
) -> Vec<usize> {
    let n = x.len();
    if n == 0 {
        return Vec::new();
    }

    let mut peaks = Vec::new();
    let mut last_peak: Option<usize> = None;

    for i in 0..n {
        // Condition 1: local maximum in [i-pre_max, i+post_max)
        // librosa uses exclusive end: x[max(0, n-pre_max):min(n+post_max, len)]
        let max_start = i.saturating_sub(pre_max);
        let max_end = (i + post_max).min(n);
        let local_max = x[max_start..max_end]
            .iter()
            .copied()
            .fold(f64::NEG_INFINITY, f64::max);
        if x[i] < local_max {
            continue;
        }

        // Condition 2: above mean + delta
        // librosa uses exclusive end: x[max(0, n-pre_avg):min(n+post_avg, len)]
        let avg_start = i.saturating_sub(pre_avg);
        let avg_end = (i + post_avg).min(n);
        let window = &x[avg_start..avg_end];
        let local_mean = window.iter().sum::<f64>() / window.len() as f64;
        if x[i] < local_mean + delta {
            continue;
        }

        // Condition 3: sufficient spacing
        if let Some(last) = last_peak
            && i <= last + wait
        {
            continue;
        }

        peaks.push(i);
        last_peak = Some(i);
    }

    peaks
}

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

    fn load_util_references() -> Vec<serde_json::Value> {
        let data = include_str!("tests/references/util_functions.json");
        serde_json::from_str(data).expect("Failed to parse util_functions.json")
    }

    fn assert_f64_eq(actual: &[f64], expected: &[f64], tol: f64, name: &str) {
        assert_eq!(actual.len(), expected.len(), "{name}: length mismatch");
        for (i, (a, e)) in actual.iter().zip(expected.iter()).enumerate() {
            let diff = (a - e).abs();
            assert!(
                diff < tol,
                "{name}[{i}]: got {a}, expected {e}, diff {diff}"
            );
        }
    }

    #[test]
    fn test_pad_center_matches_scipy() {
        let refs = load_util_references();
        for case in refs.iter().filter(|c| c["function"] == "pad_center") {
            let name = case["name"].as_str().unwrap();
            let input: Vec<f64> = serde_json::from_value(case["input"].clone()).unwrap();
            let size = case["size"].as_u64().unwrap() as usize;
            let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();
            let result = pad_center(&input, size);
            assert_f64_eq(&result, &expected, 1e-12, name);
        }
    }

    #[test]
    fn test_fix_length_matches_scipy() {
        let refs = load_util_references();
        for case in refs.iter().filter(|c| c["function"] == "fix_length") {
            let name = case["name"].as_str().unwrap();
            let input: Vec<f64> = serde_json::from_value(case["input"].clone()).unwrap();
            let size = case["size"].as_u64().unwrap() as usize;
            let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();
            let result = fix_length(&input, size);
            assert_f64_eq(&result, &expected, 1e-12, name);
        }
    }

    #[test]
    fn test_normalize_matches_scipy() {
        let refs = load_util_references();
        for case in refs.iter().filter(|c| c["function"] == "normalize") {
            let name = case["name"].as_str().unwrap();
            let mut input: Vec<f64> = serde_json::from_value(case["input"].clone()).unwrap();
            let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();
            let norm_str = case["norm"].as_str().unwrap();
            let norm = match norm_str {
                "inf" => Norm::Linf,
                "1" => Norm::L1,
                "2" => Norm::L2,
                _ => panic!("Unknown norm: {norm_str}"),
            };
            normalize(&mut input, norm);
            assert_f64_eq(&input, &expected, 1e-12, name);
        }
    }

    #[test]
    fn test_localmax_matches_scipy() {
        let refs = load_util_references();
        for case in refs.iter().filter(|c| c["function"] == "localmax") {
            let name = case["name"].as_str().unwrap();
            let input: Vec<f64> = serde_json::from_value(case["input"].clone()).unwrap();
            let expected: Vec<bool> = serde_json::from_value(case["expected"].clone()).unwrap();
            let result = localmax(&input);
            assert_eq!(result, expected, "{name}");
        }
    }

    #[test]
    fn test_localmin_matches_scipy() {
        let refs = load_util_references();
        for case in refs.iter().filter(|c| c["function"] == "localmin") {
            let name = case["name"].as_str().unwrap();
            let input: Vec<f64> = serde_json::from_value(case["input"].clone()).unwrap();
            let expected: Vec<bool> = serde_json::from_value(case["expected"].clone()).unwrap();
            let result = localmin(&input);
            assert_eq!(result, expected, "{name}");
        }
    }

    #[test]
    fn test_softmask_matches_scipy() {
        let refs = load_util_references();
        for case in refs.iter().filter(|c| c["function"] == "softmask") {
            let name = case["name"].as_str().unwrap();
            let x: Vec<f64> = serde_json::from_value(case["x"].clone()).unwrap();
            let x_ref: Vec<f64> = serde_json::from_value(case["x_ref"].clone()).unwrap();
            let power = case["power"].as_f64().unwrap();
            let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();
            let result = softmask(&x, &x_ref, power);
            assert_f64_eq(&result, &expected, 1e-10, name);
        }
    }

    #[test]
    fn test_tiny() {
        assert_eq!(tiny(), f64::MIN_POSITIVE);
        assert!(tiny() > 0.0);
    }

    #[test]
    fn test_pad_center_panics_if_too_small() {
        let result = std::panic::catch_unwind(|| pad_center(&[1.0, 2.0, 3.0], 2));
        assert!(result.is_err());
    }

    #[test]
    fn test_maximum_filter1d_matches_scipy() {
        let refs = load_util_references();
        for case in refs.iter().filter(|c| c["function"] == "maximum_filter1d") {
            let name = case["name"].as_str().unwrap();
            let input: Vec<f64> = serde_json::from_value(case["input"].clone()).unwrap();
            let size = case["size"].as_u64().unwrap() as usize;
            let expected: Vec<f64> = serde_json::from_value(case["expected"].clone()).unwrap();
            let result = maximum_filter1d(&input, size);
            assert_f64_eq(&result, &expected, 1e-12, name);
        }
    }

    #[test]
    fn test_peak_pick_matches_librosa() {
        let refs = load_util_references();
        for case in refs.iter().filter(|c| c["function"] == "peak_pick") {
            let name = case["name"].as_str().unwrap();
            let input: Vec<f64> = serde_json::from_value(case["input"].clone()).unwrap();
            let pre_max = case["pre_max"].as_u64().unwrap() as usize;
            let post_max = case["post_max"].as_u64().unwrap() as usize;
            let pre_avg = case["pre_avg"].as_u64().unwrap() as usize;
            let post_avg = case["post_avg"].as_u64().unwrap() as usize;
            let delta = case["delta"].as_f64().unwrap();
            let wait = case["wait"].as_u64().unwrap() as usize;
            let expected: Vec<usize> = serde_json::from_value(case["expected"].clone()).unwrap();
            let result = peak_pick(&input, pre_max, post_max, pre_avg, post_avg, delta, wait);
            assert_eq!(result, expected, "{name}");
        }
    }
}