use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
use rns::modulus::WordOps;
use rns::poly::Poly;
use rns::ring::Ring;

fn ntt(c: &mut Criterion) {
    fn runner<'a, const INPLACE: bool, const LAZY: bool>(
        ring: &'a Ring<u64>,
    ) -> Box<dyn FnMut() + 'a> {
        let mut a: Poly<u64> = ring.new_poly();
        for i in 0..a.n() {
            a.0[i] = i as u64;
        }
        if INPLACE {
            Box::new(move || ring.ntt_inplace::<LAZY>(&mut a))
        } else {
            let mut b: Poly<u64> = ring.new_poly();
            Box::new(move || ring.ntt::<LAZY>(&a, &mut b))
        }
    }

    let q: u64 = 0x1fffffffffe00001u64;

    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
        c.benchmark_group("ntt");

    for log_n in 10..17 {
        let ring = Ring::new(1 << log_n, q, 1);

        let runners: [(String, Box<dyn FnMut()>); 4] = [
            (format!("inplace=true/LAZY=true/q={}", q.log2()), {
                runner::<true, true>(&ring)
            }),
            (format!("inplace=true/LAZY=false/q={}", q.log2()), {
                runner::<true, false>(&ring)
            }),
            (format!("inplace=false/LAZY=true/q={}", q.log2()), {
                runner::<false, true>(&ring)
            }),
            (format!("inplace=false/LAZY=false/q={}", q.log2()), {
                runner::<false, false>(&ring)
            }),
        ];

        for (name, mut runner) in runners {
            let id: BenchmarkId = BenchmarkId::new(name, format!("n={}", 1 << log_n));
            b.bench_with_input(id, &(), |b: &mut criterion::Bencher<'_>, _| {
                b.iter(&mut runner)
            });
        }
    }
}

fn intt(c: &mut Criterion) {
    fn runner<'a, const INPLACE: bool, const LAZY: bool>(
        ring: &'a Ring<u64>,
    ) -> Box<dyn FnMut() + 'a> {
        let mut a: Poly<u64> = ring.new_poly();
        for i in 0..a.n() {
            a.0[i] = i as u64;
        }
        if INPLACE {
            Box::new(move || ring.intt_inplace::<LAZY>(&mut a))
        } else {
            let mut b: Poly<u64> = ring.new_poly();
            Box::new(move || ring.intt::<LAZY>(&a, &mut b))
        }
    }

    let q: u64 = 0x1fffffffffe00001u64;

    let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> =
        c.benchmark_group("intt");

    for log_n in 10..17 {
        let ring = Ring::new(1 << log_n, q, 1);

        let runners: [(String, Box<dyn FnMut()>); 4] = [
            (format!("inplace=true/LAZY=true/q={}", q.log2()), {
                runner::<true, true>(&ring)
            }),
            (format!("inplace=true/LAZY=false/q={}", q.log2()), {
                runner::<true, false>(&ring)
            }),
            (format!("inplace=false/LAZY=true/q={}", q.log2()), {
                runner::<false, true>(&ring)
            }),
            (format!("inplace=false/LAZY=false/q={}", q.log2()), {
                runner::<false, false>(&ring)
            }),
        ];

        for (name, mut runner) in runners {
            let id: BenchmarkId = BenchmarkId::new(name, format!("n={}", 1 << log_n));
            b.bench_with_input(id, &(), |b: &mut criterion::Bencher<'_>, _| {
                b.iter(&mut runner)
            });
        }
    }
}

criterion_group!(benches, ntt, intt,);
criterion_main!(benches);