use std::sync::Arc;

use criterion::{black_box, criterion_group, criterion_main, Criterion};
use imsearch::multithreading::ThreadPool;

fn dot(a: &[f64], b: &[f64]) -> f64 {
    let mut sum = 0.0;

    for i in 0..a.len() {
        sum += a[i] * b[i];
    }

    sum
}

fn bench_single_threaded(a: &Vec<f64>, b: &Vec<f64>) {
    black_box(dot(a, b));
}

fn bench_threadpool(a: Arc<Vec<f64>>, b: Arc<Vec<f64>>) {
    let mut pool = ThreadPool::new();

    const CHUNKS: usize = 100;

    let steps = a.len() / CHUNKS;

    for i in 0..CHUNKS {
        let chunk = i * steps;
        let aa = a.clone();
        let bb = b.clone();
        pool.enqueue(move || {
            let a = &aa[chunk..(chunk + steps)];
            let b = &bb[chunk..(chunk + steps)];
            dot(a, b)
        });
    }

    black_box(
        pool.join_all()
            .into_iter()
            .map(|r| r.unwrap())
            .reduce(|a, b| a + b),
    );
}

#[inline]
fn hash(x: f64) -> f64 {
    ((x * 234.8743 + 3.8274).sin() * 87624.58376).fract()
}

fn create_vec(size: usize) -> Arc<Vec<f64>> {
    let mut vec = Vec::with_capacity(size);

    for i in 0..size {
        vec.push(hash(i as f64));
    }

    Arc::new(vec)
}

pub fn benchmark_threadpool(c: &mut Criterion) {
    let vec_a = create_vec(1_000_000);
    let vec_b = create_vec(1_000_000);

    c.bench_function("single threaded", |b| {
        b.iter(|| bench_single_threaded(&vec_a, &vec_b))
    });
    c.bench_function("multi threaded", |b| {
        b.iter(|| bench_threadpool(vec_a.clone(), vec_b.clone()))
    });
}

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