use criterion::{criterion_group, criterion_main, Criterion, Fun};
/**
 * This module benchmarks the performance of different storage schemes
 */
use rand::distributions::{Distribution, Standard};
use rand::prelude::*;

static SEED: [u8; 32] = [3; 32];
static BUF_SIZE: usize = 1_000_000;

#[inline]
fn operation<T: num_traits::Float>(a: [T; 9], b: [T; 9]) -> [T; 9] {
    let mut c = [T::zero(); 9];
    for i in 0..9 {
        c[i] = a[i] + b[i];
    }
    c
}

#[inline]
fn make_random_vec<T>() -> Vec<[T; 9]>
where
    Standard: Distribution<T>,
{
    let n = BUF_SIZE;

    let mut v = Vec::new();
    let mut rng: StdRng = SeedableRng::from_seed(SEED);
    for _ in 0..n {
        v.push([
            rng.gen(),
            rng.gen(),
            rng.gen(),
            rng.gen(),
            rng.gen(),
            rng.gen(),
            rng.gen(),
            rng.gen(),
            rng.gen(),
        ]);
    }

    v
}

fn array_of_structs<T: std::ops::AddAssign<T> + num_traits::Float>(
    i: &Vec<([T; 9], [T; 9])>,
) -> [T; 9] {
    let mut sum = [T::zero(); 9];
    for &(a, b) in i.iter() {
        let res = operation(a, b);
        for i in 0..9 {
            sum[i] += res[i];
        }
    }
    sum
}

fn struct_of_arrays<T: std::ops::AddAssign<T> + num_traits::Float>(
    vec_a: &Vec<[T; 9]>,
    vec_b: &Vec<[T; 9]>,
) -> [T; 9] {
    let mut sum = [T::zero(); 9];
    for (&a, &b) in vec_a.iter().zip(vec_b.iter()) {
        let res = operation(a, b);
        for i in 0..9 {
            sum[i] += res[i];
        }
    }
    sum
}

fn storage(c: &mut Criterion) {
    let aos = Fun::new("Array of Structs", move |b, _| {
        let v = make_random_vec::<f64>()
            .into_iter()
            .zip(make_random_vec::<f64>().into_iter())
            .collect();
        b.iter(|| array_of_structs(&v))
    });

    let soa = Fun::new("Struct of Arrays", move |b, _| {
        let vec_a = make_random_vec::<f64>();
        let vec_b = make_random_vec::<f64>();
        b.iter(|| struct_of_arrays(&vec_a, &vec_b))
    });

    let fns = vec![aos, soa];
    c.bench_functions("storage", fns, ());
}

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