rosa 0.1.0

//! Window functions for spectral analysis.
//!
//! All functions match `scipy.signal.get_window` output.
//! The `periodic` parameter corresponds to `fftbins=True` in scipy/librosa.

use std::f64::consts::PI;

/// Generalized cosine window.
///
/// `w[k] = sum(coeffs[i] * (-1)^i * cos(2*pi*i*k / N))`
///
/// When `periodic=true`, computes a symmetric window of size `n+1` and drops
/// the last sample — matching `scipy.signal.get_window(name, n, fftbins=True)`.
pub fn general_cosine(n: usize, coeffs: &[f64], periodic: bool) -> Vec<f64> {
    if n == 0 {
        return Vec::new();
    }
    // scipy special-cases M<=1: returns ones(M) before any periodic logic
    if n == 1 {
        return vec![1.0];
    }

    let m = if periodic { n + 1 } else { n };

    let mut w: Vec<f64> = (0..m)
        .map(|k| {
            let frac = (k as f64) / ((m - 1) as f64);
            coeffs
                .iter()
                .enumerate()
                .map(|(i, &c)| {
                    let sign = if i % 2 == 0 { 1.0 } else { -1.0 };
                    sign * c * (2.0 * PI * (i as f64) * frac).cos()
                })
                .sum()
        })
        .collect();

    w.truncate(n);
    w
}

/// Hann (raised cosine) window.
///
/// Coefficients: [0.5, 0.5]
pub fn hann(n: usize, periodic: bool) -> Vec<f64> {
    general_cosine(n, &[0.5, 0.5], periodic)
}

/// Hamming window.
///
/// Coefficients: [0.54, 0.46]
pub fn hamming(n: usize, periodic: bool) -> Vec<f64> {
    general_cosine(n, &[0.54, 0.46], periodic)
}

/// Blackman window.
///
/// Coefficients: [0.42, 0.50, 0.08]
pub fn blackman(n: usize, periodic: bool) -> Vec<f64> {
    general_cosine(n, &[0.42, 0.50, 0.08], periodic)
}

/// Blackman-Harris window (4-term).
///
/// Coefficients: [0.35875, 0.48829, 0.14128, 0.01168]
pub fn blackmanharris(n: usize, periodic: bool) -> Vec<f64> {
    general_cosine(n, &[0.35875, 0.48829, 0.14128, 0.01168], periodic)
}

/// Bartlett (triangular, zero-ended) window.
///
/// The Bartlett window is a triangular window that always starts and ends at zero.
pub fn bartlett(n: usize, periodic: bool) -> Vec<f64> {
    if n == 0 {
        return Vec::new();
    }
    if n == 1 {
        return vec![1.0];
    }

    let m = if periodic { n + 1 } else { n };

    let half = (m - 1) as f64 / 2.0;
    let mut w: Vec<f64> = (0..m)
        .map(|k| 1.0 - ((k as f64 - half) / half).abs())
        .collect();

    w.truncate(n);
    w
}

/// Triangular (non-zero-ended) window.
///
/// Unlike Bartlett, the endpoints are not forced to zero.
/// Matches `scipy.signal.get_window('triang', n)`.
pub fn triangular(n: usize, periodic: bool) -> Vec<f64> {
    if n == 0 {
        return Vec::new();
    }
    if n == 1 {
        return vec![1.0];
    }

    let m = if periodic { n + 1 } else { n };

    let mut w: Vec<f64> = if m % 2 == 1 {
        let half = (m + 1) as f64 / 2.0;
        (0..m)
            .map(|k| 1.0 - ((k as f64 + 1.0 - half) / half).abs())
            .collect()
    } else {
        let half = m as f64 / 2.0;
        (0..m)
            .map(|k| 1.0 - ((k as f64 + 0.5 - half) / half).abs())
            .collect()
    };

    w.truncate(n);
    w
}

/// Rectangular (boxcar) window — all ones.
pub fn boxcar(n: usize) -> Vec<f64> {
    vec![1.0; n]
}

/// Modified Bessel function of the first kind, order 0.
///
/// Uses the power series: `I0(x) = sum_{k=0}^{inf} ((x/2)^k / k!)^2`
/// Converges to machine precision for all values used in window functions.
fn i0(x: f64) -> f64 {
    let half_x = x.abs() * 0.5;
    let mut sum = 1.0;
    let mut term = 1.0;
    // 30 terms is sufficient for f64 precision up to x ~ 30
    for k in 1..=30 {
        term *= half_x / (k as f64);
        let t2 = term * term;
        sum += t2;
        if t2 < sum * f64::EPSILON {
            break;
        }
    }
    sum
}

/// Kaiser window.
///
/// Uses a Bessel I0 approximation. When `beta=0`, equivalent to boxcar.
pub fn kaiser(n: usize, beta: f64, periodic: bool) -> Vec<f64> {
    if n == 0 {
        return Vec::new();
    }
    if n == 1 {
        return vec![1.0];
    }

    let m = if periodic { n + 1 } else { n };

    let alpha = ((m - 1) as f64) / 2.0;
    let denom = i0(beta);

    let mut w: Vec<f64> = (0..m)
        .map(|k| {
            let ratio = (k as f64 - alpha) / alpha;
            let arg = beta * (1.0 - ratio * ratio).max(0.0).sqrt();
            i0(arg) / denom
        })
        .collect();

    w.truncate(n);
    w
}

/// Look up a window function by name, matching `librosa.filters.get_window`.
///
/// Supported names (case-insensitive):
/// `"hann"`, `"hanning"`, `"hamming"`, `"blackman"`, `"blackmanharris"`,
/// `"bartlett"`, `"triangular"`, `"triang"`, `"boxcar"`, `"rectangular"`, `"ones"`.
///
/// For Kaiser windows, use `"kaiser:beta"` (e.g. `"kaiser:4.0"`).
///
/// # Errors
/// Returns `Err` for unknown window names or invalid kaiser beta values.
pub fn get_window(name: &str, n: usize, periodic: bool) -> Result<Vec<f64>, String> {
    let name_lower = name.to_lowercase();
    if let Some(beta_str) = name_lower.strip_prefix("kaiser:") {
        let beta: f64 = beta_str
            .parse()
            .map_err(|_| format!("Invalid kaiser beta: {}", beta_str))?;
        return Ok(kaiser(n, beta, periodic));
    }
    match name_lower.as_str() {
        "hann" | "hanning" => Ok(hann(n, periodic)),
        "hamming" => Ok(hamming(n, periodic)),
        "blackman" => Ok(blackman(n, periodic)),
        "blackmanharris" => Ok(blackmanharris(n, periodic)),
        "bartlett" => Ok(bartlett(n, periodic)),
        "triangular" | "triang" => Ok(triangular(n, periodic)),
        "boxcar" | "rectangular" | "ones" => Ok(boxcar(n)),
        _ => Err(format!("Unknown window: {}", name)),
    }
}

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

    #[derive(Deserialize)]
    struct WindowTestCase {
        name: String,
        function: String,
        n: usize,
        periodic: bool,
        #[serde(default)]
        beta: Option<f64>,
        expected: Vec<f64>,
    }

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

    fn assert_approx_eq(actual: &[f64], expected: &[f64], tol: f64, name: &str) {
        assert_eq!(
            actual.len(),
            expected.len(),
            "{name}: length mismatch: got {} expected {}",
            actual.len(),
            expected.len()
        );
        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} > tol {tol}"
            );
        }
    }

    #[test]
    fn test_hann_matches_scipy() {
        let refs = load_window_references();
        for case in refs.iter().filter(|c| c.function == "hann") {
            let result = hann(case.n, case.periodic);
            assert_approx_eq(&result, &case.expected, 1e-12, &case.name);
        }
    }

    #[test]
    fn test_hamming_matches_scipy() {
        let refs = load_window_references();
        for case in refs.iter().filter(|c| c.function == "hamming") {
            let result = hamming(case.n, case.periodic);
            assert_approx_eq(&result, &case.expected, 1e-12, &case.name);
        }
    }

    #[test]
    fn test_blackman_matches_scipy() {
        let refs = load_window_references();
        for case in refs.iter().filter(|c| c.function == "blackman") {
            let result = blackman(case.n, case.periodic);
            assert_approx_eq(&result, &case.expected, 1e-12, &case.name);
        }
    }

    #[test]
    fn test_blackmanharris_matches_scipy() {
        let refs = load_window_references();
        for case in refs.iter().filter(|c| c.function == "blackmanharris") {
            let result = blackmanharris(case.n, case.periodic);
            assert_approx_eq(&result, &case.expected, 1e-12, &case.name);
        }
    }

    #[test]
    fn test_bartlett_matches_scipy() {
        let refs = load_window_references();
        for case in refs.iter().filter(|c| c.function == "bartlett") {
            let result = bartlett(case.n, case.periodic);
            assert_approx_eq(&result, &case.expected, 1e-12, &case.name);
        }
    }

    #[test]
    fn test_triangular_matches_scipy() {
        let refs = load_window_references();
        for case in refs.iter().filter(|c| c.function == "triangular") {
            let result = triangular(case.n, case.periodic);
            assert_approx_eq(&result, &case.expected, 1e-12, &case.name);
        }
    }

    #[test]
    fn test_boxcar_matches_scipy() {
        let refs = load_window_references();
        for case in refs.iter().filter(|c| c.function == "boxcar") {
            let result = boxcar(case.n);
            assert_approx_eq(&result, &case.expected, 1e-12, &case.name);
        }
    }

    #[test]
    fn test_kaiser_matches_scipy() {
        let refs = load_window_references();
        for case in refs.iter().filter(|c| c.function == "kaiser") {
            let beta = case.beta.unwrap_or(5.0);
            let result = kaiser(case.n, beta, case.periodic);
            assert_approx_eq(&result, &case.expected, 1e-10, &case.name);
        }
    }

    #[test]
    fn test_n1_returns_one() {
        assert_eq!(hann(1, true), vec![1.0]);
        assert_eq!(hamming(1, true), vec![1.0]);
        assert_eq!(blackman(1, true), vec![1.0]);
        assert_eq!(kaiser(1, 5.0, true), vec![1.0]);
        assert_eq!(bartlett(1, true), vec![1.0]);
        assert_eq!(triangular(1, true), vec![1.0]);
    }

    #[test]
    fn test_get_window_dispatches_correctly() {
        // Each supported name should match the direct function call
        assert_eq!(get_window("hann", 64, true).unwrap(), hann(64, true));
        assert_eq!(get_window("Hanning", 64, true).unwrap(), hann(64, true));
        assert_eq!(get_window("HAMMING", 64, true).unwrap(), hamming(64, true));
        assert_eq!(
            get_window("blackman", 64, false).unwrap(),
            blackman(64, false)
        );
        assert_eq!(
            get_window("blackmanharris", 64, true).unwrap(),
            blackmanharris(64, true)
        );
        assert_eq!(
            get_window("bartlett", 64, true).unwrap(),
            bartlett(64, true)
        );
        assert_eq!(
            get_window("triang", 64, true).unwrap(),
            triangular(64, true)
        );
        assert_eq!(
            get_window("triangular", 64, true).unwrap(),
            triangular(64, true)
        );
        assert_eq!(get_window("boxcar", 64, true).unwrap(), boxcar(64));
        assert_eq!(get_window("rectangular", 64, true).unwrap(), boxcar(64));
        assert_eq!(get_window("ones", 64, true).unwrap(), boxcar(64));
        assert_eq!(
            get_window("kaiser:5.0", 64, true).unwrap(),
            kaiser(64, 5.0, true)
        );
    }

    #[test]
    fn test_get_window_unknown() {
        assert!(get_window("unknown", 64, true).is_err());
    }

    #[test]
    fn test_get_window_kaiser_invalid_beta() {
        assert!(get_window("kaiser:abc", 64, true).is_err());
    }

    #[test]
    fn test_n0_returns_empty() {
        assert!(hann(0, true).is_empty());
        assert!(hamming(0, true).is_empty());
        assert!(kaiser(0, 5.0, true).is_empty());
        assert!(bartlett(0, true).is_empty());
        assert!(boxcar(0).is_empty());
    }
}