use std::{
    hint::black_box,
    iter::repeat_with,
    num::{NonZeroU8, NonZeroUsize},
};

use campo::{BitField, Campo};
use criterion::{criterion_group, criterion_main, BatchSize, Criterion};

// NOTE: there's a pervasive use of balck_box for both inputs and function
// calls. The inputs are all randomized and the compiler should not be able to
// optimize based on this inputs. Some benchmarked functions can instead be
// compiled away without that (unused returned values, side effects on values
// that are not used afterwards, etc...). The extensive use of black_box is to
// ensure no possible optimizations can be applied (except for those relative to
// the function's logic per-se), even in cases where black_box might not be
// needed. Also, I'm not interested in tha absolute value of these benchmarks,
// but rather in the deltas after a refactor/optimization: black_box should
// serve as a normalizer.

fn bitwise_ops(c: &mut Criterion) {
    let scalar_field1: Campo<u8> = fastrand::u8(..).into();
    let mut scalar_field2: Campo<u8> = fastrand::u8(..).into();
    scalar_field2
        .resize_in_place(NonZeroU8::new(7).unwrap())
        .unwrap();

    let vec_field1: Campo<Vec<u8>> = repeat_with(|| fastrand::u8(..))
        .take(100)
        .collect::<Vec<u8>>()
        .try_into()
        .unwrap();
    let mut vec_field2: Campo<Vec<u8>> = repeat_with(|| fastrand::u8(..))
        .take(101)
        .collect::<Vec<u8>>()
        .try_into()
        .unwrap();
    vec_field2.resize(NonZeroUsize::new(805).unwrap());

    // The benchmarks on the u8 variant are just a reference point for the
    // performance of the vectorized implementation
    let mut group = c.benchmark_group("Bitwise");
    group.bench_function("u8_and", |b| {
        b.iter(|| black_box(scalar_field1 & scalar_field2))
    });
    group.bench_function("u8_or", |b| {
        b.iter(|| black_box(scalar_field1 | scalar_field2))
    });
    group.bench_function("u8_xor", |b| {
        b.iter(|| black_box(scalar_field1 ^ scalar_field2))
    });
    group.bench_function("vector_and", |b| {
        b.iter_batched(
            || (vec_field1.clone(), vec_field2.clone()),
            |(vec_field1, vec_field2)| black_box(vec_field1 & vec_field2),
            BatchSize::SmallInput,
        )
    });
    group.bench_function("vector_or", |b| {
        b.iter_batched(
            || (vec_field1.clone(), vec_field2.clone()),
            |(vec_field1, vec_field2)| black_box(vec_field1 | vec_field2),
            BatchSize::SmallInput,
        )
    });
    group.bench_function("vector_xor", |b| {
        b.iter_batched(
            || (vec_field1.clone(), vec_field2.clone()),
            |(vec_field1, vec_field2)| black_box(vec_field1 ^ vec_field2),
            BatchSize::SmallInput,
        )
    });

    group.finish();
}

fn shifts(c: &mut Criterion) {
    let scalar_shift_amt = fastrand::u8(..61);
    let vec_shift_amt = fastrand::usize(..500);
    let mut scalar_field: Campo<u64> = fastrand::u64(..).into();
    scalar_field
        .resize_in_place(NonZeroU8::new(62).unwrap())
        .unwrap();

    let mut vec_field: Campo<Vec<u8>> = repeat_with(|| fastrand::u8(..))
        .take(101)
        .collect::<Vec<u8>>()
        .try_into()
        .unwrap();
    vec_field.resize(NonZeroUsize::new(805).unwrap());

    // The benchmarks on the u64 variant are just a reference point for the
    // performance of the vectorized implementation
    let mut group = c.benchmark_group("Shifts");
    group.bench_function("u64_shl", |b| {
        b.iter(|| black_box(scalar_field << scalar_shift_amt))
    });
    group.bench_function("u64_shr", |b| {
        b.iter(|| black_box(scalar_field >> scalar_shift_amt))
    });
    group.bench_function("vector_shl", |b| {
        b.iter_batched(
            || vec_field.clone(),
            |vec_field| black_box(vec_field << vec_shift_amt),
            BatchSize::SmallInput,
        )
    });
    group.bench_function("vector_shr", |b| {
        b.iter_batched(
            || vec_field.clone(),
            |vec_field| black_box(vec_field >> vec_shift_amt),
            BatchSize::SmallInput,
        )
    });

    group.finish();
}

fn resize(c: &mut Criterion) {
    let scalar_shrink_amt = NonZeroU8::new(fastrand::u8(..32)).unwrap();
    let scalar_grow_in_place_amt =
        NonZeroU8::new(fastrand::u8(41..64)).unwrap();
    let scalar_grow_amt = NonZeroU8::new(fastrand::u8(65..128)).unwrap();
    let vec_shrink_amt = NonZeroUsize::new(fastrand::usize(..500)).unwrap();
    let vec_grow_amt = NonZeroUsize::new(fastrand::usize(806..1000)).unwrap();

    let mut scalar_field: Campo<u64> = fastrand::u64(..).into();
    scalar_field
        .resize_in_place(NonZeroU8::new(40).unwrap())
        .unwrap();
    let mut vec_field: Campo<Vec<u8>> = repeat_with(|| fastrand::u8(..))
        .take(101)
        .collect::<Vec<u8>>()
        .try_into()
        .unwrap();
    vec_field.resize(NonZeroUsize::new(805).unwrap());

    // The benchmarks on the u64 variant are just a reference point for the
    // performance of the vectorized implementation
    let mut group = c.benchmark_group("Resize");
    group.bench_function("u64_shrink_in_place", |b| {
        b.iter(|| {
            black_box(
                scalar_field
                    .resize_in_place(black_box(scalar_shrink_amt))
                    .unwrap(),
            )
        })
    });
    group.bench_function("u64_shrink_to_u32", |b| {
        b.iter(|| {
            black_box(
                scalar_field
                    .resize::<u32>(black_box(scalar_shrink_amt))
                    .unwrap(),
            )
        })
    });
    group.bench_function("u64_grow_in_place", |b| {
        b.iter(|| {
            black_box(
                scalar_field
                    .resize_in_place(black_box(scalar_grow_in_place_amt))
                    .unwrap(),
            )
        })
    });
    group.bench_function("u64_grow_to_u128", |b| {
        b.iter(|| {
            black_box(
                scalar_field
                    .resize::<u128>(black_box(scalar_grow_amt))
                    .unwrap(),
            )
        })
    });
    group.bench_function("vector_shrink", |b| {
        b.iter_batched(
            || vec_field.clone(),
            |mut vec_field| {
                black_box(vec_field.resize(black_box(vec_shrink_amt)))
            },
            BatchSize::SmallInput,
        )
    });
    group.bench_function("vector_grow", |b| {
        b.iter_batched(
            || vec_field.clone(),
            |mut vec_field| {
                black_box(vec_field.resize(black_box(vec_grow_amt)))
            },
            BatchSize::SmallInput,
        )
    });

    group.finish();
}

fn bitfield(c: &mut Criterion) {
    let scalar_idx = fastrand::u8(..90);
    let vec_idx = fastrand::usize(..805);

    let mut scalar_field: Campo<u128> = fastrand::u128(..).into();
    scalar_field
        .resize_in_place(NonZeroU8::new(90).unwrap())
        .unwrap();
    let mut vec_field: Campo<Vec<u8>> = repeat_with(|| fastrand::u8(..))
        .take(101)
        .collect::<Vec<u8>>()
        .try_into()
        .unwrap();
    vec_field.resize(NonZeroUsize::new(805).unwrap());

    // The benchmarks on the u128 variant are just a reference point for the
    // performance of the vectorized implementation
    let mut group = c.benchmark_group("Bitfield");

    group.bench_function("u128_get_unchecked", |b| {
        b.iter(|| black_box(scalar_field.get_unchecked(black_box(scalar_idx))))
    });
    group.bench_function("u128_get", |b| {
        b.iter(|| black_box(scalar_field.get(black_box(scalar_idx)).unwrap()))
    });
    group.bench_function("vector_get_unchecked", |b| {
        b.iter(|| black_box(vec_field.get_unchecked(black_box(vec_idx))));
    });
    group.bench_function("vector_get", |b| {
        b.iter(|| black_box(vec_field.get(black_box(vec_idx)).unwrap()));
    });

    group.bench_function("u128_set_unchecked", |b| {
        b.iter(|| black_box(scalar_field.set_unchecked(black_box(scalar_idx))))
    });
    group.bench_function("u128_set", |b| {
        b.iter(|| black_box(scalar_field.set(black_box(scalar_idx)).unwrap()))
    });
    group.bench_function("vector_set_unchecked", |b| {
        b.iter(|| black_box(vec_field.set_unchecked(black_box(vec_idx))));
    });
    group.bench_function("vector_set", |b| {
        b.iter(|| black_box(vec_field.set(black_box(vec_idx)).unwrap()));
    });

    group.bench_function("u128_clear_unchecked", |b| {
        b.iter(|| {
            black_box(scalar_field.clear_unchecked(black_box(scalar_idx)))
        })
    });
    group.bench_function("u128_clear", |b| {
        b.iter(|| black_box(scalar_field.clear(black_box(scalar_idx)).unwrap()))
    });
    group.bench_function("vector_clear_unchecked", |b| {
        b.iter(|| black_box(vec_field.clear_unchecked(black_box(vec_idx))));
    });
    group.bench_function("vector_clear", |b| {
        b.iter(|| black_box(vec_field.clear(black_box(vec_idx)).unwrap()));
    });

    group.bench_function("u128_count_sets", |b| {
        b.iter(|| black_box(scalar_field.count_ones()))
    });
    group.bench_function("vector_count_sets", |b| {
        b.iter(|| black_box(vec_field.count_ones()));
    });

    group.finish();
}

criterion_group!(benches, bitwise_ops, shifts, resize, bitfield);
criterion_main!(benches);
