use std::sync::Arc;
use std::sync::Barrier;
use std::sync::atomic::AtomicBool;

use rayon::prelude::*;

const OPS: &[usize] = &[10];
const THREADS: &[usize] = &[2, 4, 8, 16];
const KIB: usize = 1024;
const MIB: usize = KIB * KIB;
const SIZES: &[(&str, usize)] = &[
    ("16KiB", 16 * KIB),
    ("32KiB", 32 * KIB),
    ("64KiB", 64 * KIB),
    ("128KiB", 128 * KIB),
    ("256KiB", 256 * KIB),
    ("512KiB", 512 * KIB),
    ("1MiB", MIB),
    ("2MiB", 2 * MIB),
    ("4MiB", 4 * MIB),
    ("8MiB", 8 * MIB),
];

#[global_allocator]
static ALLOC: std::alloc::System = std::alloc::System;

struct Size(usize);

fn alloc(size: Size) {
    let v = Vec::<u8>::with_capacity(size.0);
    std::hint::black_box(v);
}

fn alloc_set(size: Size) {
    let v = vec![42_u8; size.0];
    std::hint::black_box(v);
}

fn reserve(size: Size) {
    let mut v = Vec::<u8>::with_capacity(1);
    while v.capacity() < size.0 {
        let capacity = v.capacity();
        let additional = capacity.min(size.0 - capacity);
        v.reserve(capacity + additional);
    }
    std::hint::black_box(v);
}

fn reserve_set(size: Size) {
    let mut v = vec![42_u8; 1];
    while v.len() < size.0 {
        let len = v.len();
        let additional = len.min(size.0 - len);
        v.extend((0..additional).map(|n| n as u8));
    }
    std::hint::black_box(v);
}

fn bench_iter(
    criterion: &mut criterion::Criterion,
    group_name: &str,
    size_str: &str,
    size: usize,
    ops: usize,
    routine: fn(Size),
) {
    let mut group = criterion.benchmark_group(group_name);

    let input = vec![size; ops];

    group.throughput(criterion::Throughput::Elements(ops as u64));

    group.bench_with_input(
        format!("iter/size:{size_str}/ops:{ops}"),
        &input,
        |bencher, input| {
            bencher.iter(|| {
                input.iter().for_each(|&size| {
                    routine(Size(size));
                })
            })
        },
    );

    group.finish();
}

fn bench_par_iter(
    criterion: &mut criterion::Criterion,
    group_name: &str,
    size_str: &str,
    size: usize,
    ops: usize,
    routine: fn(Size),
) {
    let mut group = criterion.benchmark_group(group_name);

    let input = vec![size; ops];

    group.throughput(criterion::Throughput::Elements(ops as u64));

    for &threads in THREADS {
        let pool = rayon::ThreadPoolBuilder::new()
            .num_threads(threads)
            .build()
            .unwrap();

        group.bench_with_input(
            format!("par_iter/size:{size_str}/ops:{ops}/threads:{threads}"),
            &input,
            |bencher, input| {
                pool.install(|| {
                    bencher.iter(|| {
                        input.par_iter().for_each(|&size| {
                            routine(Size(size));
                        })
                    })
                })
            },
        );
    }

    group.finish();
}

fn bench_spawn(
    criterion: &mut criterion::Criterion,
    group_name: &str,
    size_str: &str,
    size: usize,
    ops: usize,
    routine: fn(Size),
) {
    let mut group = criterion.benchmark_group(group_name);

    let input = Arc::new(vec![size; ops]);

    group.throughput(criterion::Throughput::Elements(ops as u64));

    for &threads in THREADS {
        // Pre-spawn threads.
        let stop_flag = Arc::new(AtomicBool::new(false));
        let begin_barrier = Arc::new(Barrier::new(threads + 1));
        let end_barrier = Arc::new(Barrier::new(threads + 1));
        let mut join_handles = Vec::with_capacity(threads);
        for _ in 0..threads {
            let input = Arc::clone(&input);
            let stop_flag = Arc::clone(&stop_flag);
            let begin_barrier = Arc::clone(&begin_barrier);
            let end_barrier = Arc::clone(&end_barrier);
            join_handles.push(std::thread::spawn(move || {
                loop {
                    begin_barrier.wait();
                    input.iter().for_each(|&size| {
                        routine(Size(size));
                    });
                    end_barrier.wait();
                    if stop_flag.load(std::sync::atomic::Ordering::Acquire) {
                        break;
                    }
                }
            }));
        }

        // Run the benchmark.
        group.bench_with_input(
            format!("spawn/size:{size_str}/ops:{ops}/threads:{threads}"),
            &input,
            |bencher, _input| {
                bencher.iter(|| {
                    begin_barrier.wait();
                    // Wait for threads to execute the routine in parallel.
                    end_barrier.wait();
                })
            },
        );

        // Stop and join threads.
        stop_flag.store(true, std::sync::atomic::Ordering::Release);
        begin_barrier.wait();
        end_barrier.wait();
        join_handles
            .into_iter()
            .for_each(|handle| handle.join().unwrap());
    }

    group.finish();
}

fn all_benches(criterion: &mut criterion::Criterion, group_name: &str, routine: fn(Size)) {
    for &(size_str, size) in SIZES {
        for &ops in OPS {
            bench_iter(criterion, group_name, size_str, size, ops, routine);
            bench_par_iter(criterion, group_name, size_str, size, ops, routine);
            bench_spawn(criterion, group_name, size_str, size, ops, routine);
        }
    }
}

fn alloc_bench(c: &mut criterion::Criterion) {
    all_benches(c, "alloc", alloc);
    all_benches(c, "alloc_set", alloc_set);
}

fn reserve_bench(c: &mut criterion::Criterion) {
    all_benches(c, "reserve", reserve);
    all_benches(c, "reserve_set", reserve_set);
}

criterion::criterion_group! {
    name = benches;
    config = criterion::Criterion::default()
        .sample_size(10)
        .warm_up_time(std::time::Duration::from_secs(1))
        .measurement_time(std::time::Duration::from_secs(1))
        .sample_size(10);
    targets = alloc_bench, reserve_bench
}

criterion::criterion_main!(benches);
