microspectrogram 0.1.0

use core::num::NonZeroUsize;

pub enum Assert<const CHECK: bool> {}
pub trait IsTrue {}
impl IsTrue for Assert<true> {}

pub fn linspace<const N: usize>(start: f32, end: f32) -> [f32; N] {
    let step = (end - start) / (N - 1) as f32;
    let mut arr = [0.0; N];
    arr.iter_mut().enumerate().for_each(|(i, x)| {
        *x = start + step * i as f32;
    });
    arr
}

pub fn detrend_constant(arr: &mut [f32]) {
    let mut mean = 0.0;
    arr.iter().for_each(|x| {
        mean += x;
    });
    mean /= arr.len() as f32;

    arr.iter_mut().for_each(|x| {
        *x -= mean;
    });
}

pub struct OverlappingWindows<'a, T> {
    v: &'a [T],
    size: NonZeroUsize,
    step: NonZeroUsize,
}

pub fn overlapping_windows<T>(
    slice: &[T],
    size: NonZeroUsize,
    step: NonZeroUsize,
) -> OverlappingWindows<T> {
    OverlappingWindows {
        v: slice,
        size,
        step,
    }
}

impl<'a, T> Iterator for OverlappingWindows<'a, T> {
    type Item = &'a [T];

    #[inline]
    fn next(&mut self) -> Option<&'a [T]> {
        if self.size.get() > self.v.len() {
            None
        } else {
            let ret = Some(&self.v[..self.size.get()]);
            self.v = &self.v[self.step.get()..];
            ret
        }
    }
}

#[cfg(test)]
mod test {
    use core::{
        f32::consts::{FRAC_PI_3, PI},
        ops,
    };

    use num_traits::ToPrimitive;

    use super::*;

    fn all_close<T, const N: usize>(a: &[T; N], b: &[T; N]) -> bool
    where
        T: ops::Sub<Output = T> + PartialOrd + ToPrimitive + Copy,
    {
        const EPSILON: f32 = 1e-5;
        for i in 0..N {
            if (a[i] - b[i]).to_f32().unwrap() > EPSILON {
                return false;
            }
        }
        true
    }

    #[test]
    #[allow(clippy::excessive_precision)]
    fn test_linspace() {
        assert!(all_close(&linspace(1.0, 5.0), &[1.0, 2.0, 3.0, 4.0, 5.0]));
        assert!(all_close(&linspace(1.0, 5.0), &[1., 2.33333, 3.66666, 5.]));
        assert!(all_close(
            &linspace(-PI, PI),
            &[-PI, -FRAC_PI_3, FRAC_PI_3, PI]
        ));
    }

    #[test]
    fn test_detrend_constant() {
        let mut de = linspace(1.0, 5.0);
        detrend_constant(&mut de);
        assert_eq!(de, [-2., -1., 0., 1., 2.]);
    }

    #[test]
    fn test_overlapping_windows() {
        let data = [1, 2, 3, 4];
        let mut it = overlapping_windows(
            &data,
            NonZeroUsize::new(2).unwrap(),
            NonZeroUsize::new(1).unwrap(),
        );
        assert_eq!(it.next().unwrap(), &[1, 2]);
        assert_eq!(it.next().unwrap(), &[2, 3]);
        assert_eq!(it.next().unwrap(), &[3, 4]);
        assert!(it.next().is_none());

        let mut it = overlapping_windows(
            &data,
            NonZeroUsize::new(2).unwrap(),
            NonZeroUsize::new(2).unwrap(),
        );
        assert_eq!(it.next().unwrap(), &[1, 2]);
        assert_eq!(it.next().unwrap(), &[3, 4]);
        assert!(it.next().is_none());

        let data = [1, 2, 3, 4, 5, 6, 7, 8];
        let mut it = overlapping_windows(
            &data,
            NonZeroUsize::new(4).unwrap(),
            NonZeroUsize::new(2).unwrap(),
        );
        assert_eq!(it.next().unwrap(), &[1, 2, 3, 4]);
        assert_eq!(it.next().unwrap(), &[3, 4, 5, 6]);
        assert_eq!(it.next().unwrap(), &[5, 6, 7, 8]);
        assert!(it.next().is_none());

        let data = [1, 2, 3, 4, 5, 6, 7, 8, 9];
        let mut it = overlapping_windows(
            &data,
            NonZeroUsize::new(4).unwrap(),
            NonZeroUsize::new(2).unwrap(),
        );
        assert_eq!(it.next().unwrap(), &[1, 2, 3, 4]);
        assert_eq!(it.next().unwrap(), &[3, 4, 5, 6]);
        assert_eq!(it.next().unwrap(), &[5, 6, 7, 8]);
        assert!(it.next().is_none());
    }
}