pluto-sdr 0.1.1

// Author: Roman Hayn
// MIT License 2023

pub fn mean(data: &Vec<i16>) -> Option<i16> {
    let sum = data.iter().map(|&b| b as i64).sum::<i64>() as i16;
    let count = data.len() as i16;

    match count {
        positive if positive > 0 => Some(sum / count),
        _ => None,
    }
}

pub fn mean_float(data: &Vec<f32>) -> Option<f32> {
    let sum = data.iter().map(|&b| b as f64).sum::<f64>() as f32;
    let count = data.len() as f32;

    match count {
        positive if positive > 0.0 => Some(sum / count),
        _ => None,
    }
}

/// find the variance of a signal
pub fn variance(data: &Vec<i16>) -> Option<f32> {
    match (mean(data), data.len()) {
        (Some(d_mean), count) if count > 0 => Some(
            data.iter()
                .map(|val| {
                    let diff = d_mean - (*val);
                    diff as f32 * diff as f32
                })
                .sum::<f32>()
                / count as f32,
        ),
        _ => None,
    }
}

/// Just like [crate::signal::variance()] but assumes the signal mean is 0.
pub fn zero_mean_variance_float(data: &Vec<f32>) -> Option<f32> {
    if data.len() > 0 {
        return Some(
            data.iter()
                .map(|val| {
                    let diff = *val;
                    diff * diff
                })
                .sum::<f32>()
                / data.len() as f32,
        );
    } else {
        None
    }
}

/// Computes highest 2^n, while 2^n < size.
/// e.g. 513 -> 512 or 87 -> 64.
pub fn next_lower_power_two(size: usize) -> usize {
    let mut two_to_n: usize = 2;

    loop {
        let temp: usize = two_to_n * 2;
        if temp > size {
            break;
        }
        two_to_n = temp;
    }
    two_to_n
}

/// Computes next lowest 2^n while 2^n > size.
/// e.g. 501 -> 512 or 42 -> 64.
pub fn next_upper_power_two(size: usize) -> usize {
    let mut two_to_n: usize = 2;

    loop {
        let temp: usize = two_to_n * 2;
        two_to_n = temp;
        if temp >= size {
            break;
        }
    }
    two_to_n
}

/// Just like [crate::signal::variance()] but assumes the signal mean is 0.
pub fn zero_mean_variance(data: &Vec<i16>) -> Option<f32> {
    if data.len() > 0 {
        return Some(
            data.iter()
                .map(|val| {
                    let diff = *val;
                    diff as f32 * diff as f32
                })
                .sum::<f32>()
                / data.len() as f32,
        );
    } else {
        None
    }
}

pub fn std_deviation_float(data: &Vec<f32>) -> Option<f32> {
    match (mean_float(data), data.len()) {
        (Some(_mean), count) if count > 0 => {
            let var = data
                .iter()
                .map(|val| {
                    let diff = _mean - (*val);
                    diff * diff
                })
                .sum::<f32>()
                / count as f32;

            Some(var.sqrt())
        }
        _ => None,
    }
}

pub fn std_deviation(data: &Vec<i16>) -> Option<f32> {
    match (mean(data), data.len()) {
        (Some(_mean), count) if count > 0 => {
            let var = data
                .iter()
                .map(|val| {
                    let diff = _mean - (*val);
                    diff as f32 * diff as f32
                })
                .sum::<f32>()
                / count as f32;

            Some(var.sqrt())
        }
        _ => None,
    }
}

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

    #[test]
    fn test_mean() {
        let vec = vec![0, 1, 2, 4, 5, 6];
        assert_eq!(mean(&vec).unwrap(), 3);
    }

    // Should not error on larger i16 values
    #[test]
    fn test_mean_big() {
        let vec = vec![i16::MIN, i16::MIN, i16::MAX, i16::MAX];
        assert_eq!(mean(&vec).unwrap(), 0);
    }

    #[test]
    fn test_variance() {
        let vec = vec![66, 30, 40, 64];
        let res = variance(&vec).unwrap();
        assert_eq!(res, 238.0);

        let vec2 = vec![51, 21, 79, 49];
        let res2 = variance(&vec2).unwrap();
        assert_eq!(res2, 421.0);
    }

    #[test]
    fn test_zero_mean_variance() {
        let vec = vec![66, 30, 40, 64];
        let res = zero_mean_variance(&vec).unwrap();
        assert_eq!(res, 2738.0);

        let vec2 = vec![51, 21, 79, 49];
        let res2 = zero_mean_variance(&vec2).unwrap();
        assert_eq!(res2, 2921.0);
    }

    #[test]
    fn test_std_deviation() {
        let vec = vec![66, 30, 40, 64];
        println!("{:?}", std_deviation(&vec).unwrap());
        let res = format!("{:.2}", std_deviation(&vec).unwrap());
        assert_eq!(res, "15.43"); // 15.427

        let vec2 = vec![51, 21, 79, 49];
        println!("{:?}", std_deviation(&vec2).unwrap());
        let res2 = format!("{:.2}", std_deviation(&vec2).unwrap());
        assert_eq!(res2, "20.52"); // 20.518
    }
}