use std::hint::black_box;

use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use iterstats::Iterstats;
use rand::seq::IteratorRandom;

fn naive<'a>(iter: impl Iterator<Item = &'a f64> + Clone) -> Vec<f64> {
    let sum = iter.clone().sum::<f64>();
    let count = iter.clone().count() as f64;
    let mean = sum / count;
    let stddev = (iter.clone().map(|i| (i - mean).powi(2)).sum::<f64>() / count).sqrt();
    iter.map(|i| (i - mean) / stddev).collect()
}

fn bench_zscore(c: &mut Criterion) {
    let mut group = c.benchmark_group("zscore");
    let mut rng = &mut rand::thread_rng();
    for len in [10, 100, 1000, 10000, 100_000, 1_000_000] {
        let data = (0..2 * len)
            .map(|i| i as f64)
            .choose_multiple(&mut rng, len);
        assert_eq!(data.len(), len);
        group.bench_with_input(BenchmarkId::new("naive", len), &data, |b, i| {
            b.iter(|| black_box(naive(black_box(i.iter()))))
        });
        group.bench_with_input(BenchmarkId::new("iterstats", len), &data, |b, i| {
            b.iter(|| black_box(black_box(black_box(i.iter()).zscore()).collect::<Vec<_>>()))
        });
    }
    group.finish();
}

criterion_group!(benches, bench_zscore);
criterion_main!(benches);