use criterion::Criterion;

#[cfg(all(target_arch = "aarch64", target_os = "macos"))]
mod benches {
    use criterion::{BatchSize, BenchmarkId, Criterion};
    use miden_crypto::Felt;
    use miden_gpu::{
        metal::{build_merkle_tree, utils::page_aligned_uninit_vector, ColumnHasher, RowHasher},
        HashFn,
    };
    use winter_math::FieldElement;

    const BENCHMARK_INPUT_SIZES: [usize; 5] = [2048, 4096, 32768, 262144, 8388608];
    const NUM_BASE: usize = 8;

    pub fn rpo_bench(c: &mut Criterion, name: &str) {
        let mut group = c.benchmark_group(name);
        group.sample_size(10);

        for n in BENCHMARK_INPUT_SIZES {
            let input_size = 8;

            group.bench_with_input(BenchmarkId::new("RpoColumns", n), &n, |b, _| {
                let mut col0 = unsafe { page_aligned_uninit_vector(n) };
                let mut col1 = unsafe { page_aligned_uninit_vector(n) };
                let mut col2 = unsafe { page_aligned_uninit_vector(n) };
                let mut col3 = unsafe { page_aligned_uninit_vector(n) };
                let mut col4 = unsafe { page_aligned_uninit_vector(n) };
                let mut col5 = unsafe { page_aligned_uninit_vector(n) };
                let mut col6 = unsafe { page_aligned_uninit_vector(n) };
                let mut col7 = unsafe { page_aligned_uninit_vector(n) };
                col0.fill(Felt::ONE);
                col1.fill(Felt::ONE);
                col2.fill(Felt::ONE);
                col3.fill(Felt::ONE);
                col4.fill(Felt::ONE);
                col5.fill(Felt::ONE);
                col6.fill(Felt::ONE);
                col7.fill(Felt::ONE);

                b.iter_batched(
                    || ColumnHasher::new(n, NUM_BASE, HashFn::Rpo256),
                    |mut rpo| {
                        for i in 0..input_size {
                            let col = match i % 8 {
                                0 => &col0,
                                1 => &col1,
                                2 => &col2,
                                3 => &col3,
                                4 => &col4,
                                5 => &col5,
                                6 => &col6,
                                7 => &col7,
                                _ => unreachable!(),
                            };

                            rpo.update(col);
                        }

                        pollster::block_on(rpo.finish());
                    },
                    BatchSize::SmallInput,
                )
            });

            group.bench_with_input(BenchmarkId::new("RpoRows", n), &n, |b, _| {
                let mut rows: Vec<[Felt; 8]> = unsafe { page_aligned_uninit_vector(n) };
                rows.fill([Felt::ONE; 8]);

                b.iter_batched(
                    || RowHasher::new(n, NUM_BASE, HashFn::Rpo256),
                    |mut rpo| {
                        rpo.update(&rows);
                        pollster::block_on(rpo.finish());
                    },
                    BatchSize::SmallInput,
                )
            });

            group.bench_with_input(BenchmarkId::new("RpoMerkleTreeRow", n), &n, |b, _| {
                let mut rows: Vec<[Felt; 8]> = unsafe { page_aligned_uninit_vector(n) };
                rows.fill([Felt::ONE; 8]);

                let mut rpo = RowHasher::new(n, NUM_BASE, HashFn::Rpo256);
                rpo.update(&rows);
                let hashes = pollster::block_on(rpo.finish());

                b.iter(|| {
                    pollster::block_on(build_merkle_tree(&hashes, HashFn::Rpo256));
                })
            });
        }

        group.finish();
    }
}

#[cfg(all(target_arch = "aarch64", target_os = "macos"))]
pub fn rpo_benches(c: &mut Criterion) {
    benches::rpo_bench(c, "Rescue-Prime Optimized")
}

#[cfg(not(all(target_arch = "aarch64", target_os = "macos")))]
pub fn rpo_benches(_c: &mut Criterion) {}