use criterion::Criterion;
use criterion::{criterion_group, criterion_main};
use gridiron::fp_256;
use gridiron::fp_480;
use num_traits::{Inv, Pow};
use rand::RngCore;
use rand::rngs::ThreadRng;
use std::hint::black_box;
use std::ops::Neg;

fn criterion_benchmark(c: &mut Criterion) {
    // Determine limb size label for benchmark names
    #[cfg(feature = "limb31")]
    let limb_label = "31bit";

    #[cfg(feature = "limb62")]
    let limb_label = "62bit";

    #[cfg(feature = "limb31")]
    fn gen_rand_limbs(rng: &mut ThreadRng) -> [u32; fp_256::NUMLIMBS] {
        let mut limbs = [0u32; fp_256::NUMLIMBS];
        for limb in limbs.iter_mut() {
            *limb = rng.next_u32();
        }
        limbs[fp_256::NUMLIMBS - 1] &= 0xFF; //Ensure the last limb isn't too big
        limbs
    }

    #[cfg(feature = "limb62")]
    fn gen_rand_limbs(rng: &mut ThreadRng) -> [u64; fp_256::NUMLIMBS] {
        let mut limbs = [0u64; fp_256::NUMLIMBS];
        for limb in limbs.iter_mut() {
            *limb = rng.next_u64();
        }
        limbs[fp_256::NUMLIMBS - 1] &= 0xFF; //Ensure the last limb isn't too big
        limbs
    }

    fn gen_rand_sixty_four_bytes(rng: &mut ThreadRng) -> [u8; 64] {
        let mut limbs = [0u8; 64];
        rng.fill_bytes(&mut limbs[..]);
        limbs
    }

    fn gen_rand_fp256_raw(rng: &mut ThreadRng) -> fp_256::Fp256 {
        fp_256::Fp256::new(gen_rand_limbs(rng))
    }

    fn gen_rand_fp256(rng: &mut ThreadRng) -> fp_256::Fp256 {
        gen_rand_fp256_raw(rng).normalize_little()
    }

    #[cfg(feature = "limb31")]
    fn gen_rand_480_limbs(rng: &mut ThreadRng) -> [u32; fp_480::NUMLIMBS] {
        let mut limbs = [0u32; fp_480::NUMLIMBS];
        for limb in limbs.iter_mut() {
            *limb = rng.next_u32();
        }
        limbs[fp_480::NUMLIMBS - 1] &= 0xFEFF;
        limbs
    }

    #[cfg(feature = "limb62")]
    fn gen_rand_480_limbs(rng: &mut ThreadRng) -> [u64; fp_480::NUMLIMBS] {
        let mut limbs = [0u64; fp_480::NUMLIMBS];
        for limb in limbs.iter_mut() {
            *limb = rng.next_u64();
        }
        limbs[fp_480::NUMLIMBS - 1] &= 0xFEFF;
        limbs
    }

    fn gen_rand_fp480_raw(rng: &mut ThreadRng) -> fp_480::Fp480 {
        fp_480::Fp480::new(gen_rand_480_limbs(rng))
    }

    fn gen_rand_fp480(rng: &mut ThreadRng) -> fp_480::Fp480 {
        (gen_rand_fp480_raw(rng)).normalize_little()
    }

    c.bench_function(
        &format!(
            "Fp256-{} - normalize_little (256 bits to Fp256 100 times)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || gen_rand_fp256_raw(&mut rng),
                |val_to_norm| {
                    for _ in 0..100 {
                        black_box(val_to_norm.normalize_little());
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!("Fp256-{} - add (add two Fp256s 100 times)", limb_label),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp256(&mut rng), gen_rand_fp256(&mut rng)),
                |(mut a, b)| {
                    for _ in 0..100 {
                        a = black_box(a + b);
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!(
            "Fp256-{} - add_assign (add an Fp256 into another Fp256 100 times)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp256(&mut rng), gen_rand_fp256(&mut rng)),
                |(mut a, b)| {
                    for _ in 0..100 {
                        a += b;
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!("Fp256-{} - sub (subtract two Fp256s 100 times)", limb_label),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp256(&mut rng), gen_rand_fp256(&mut rng)),
                |(mut a, b)| {
                    for _ in 0..100 {
                        a = black_box(a - b);
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!(
            "Fp256-{} - sub_assign (subtract an Fp256 from another Fp256 100 times)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp256(&mut rng), gen_rand_fp256(&mut rng)),
                |(mut a, b)| {
                    for _ in 0..100 {
                        a -= b;
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!("Fp256-{} - mul (two Fp256s 100 times)", limb_label),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp256(&mut rng), gen_rand_fp256(&mut rng)),
                |(mut a, b)| {
                    for _ in 0..100 {
                        a = black_box(a * b);
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!(
            "Fp256-{} - mul_assign (an Fp256 into another Fp256 100 times)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp256(&mut rng), gen_rand_fp256(&mut rng)),
                |(mut a, b)| {
                    for _ in 0..100 {
                        a *= b;
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!("Fp256-{} - div (divide two Fp256s)", limb_label),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp256(&mut rng), gen_rand_fp256(&mut rng)),
                |(a, b)| {
                    black_box(a / b);
                },
            );
        },
    );

    c.bench_function(
        &format!("Fp256-{} - neg (negate an Fp256 100 times)", limb_label),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || gen_rand_fp256(&mut rng),
                |mut a| {
                    for _ in 0..100 {
                        a = a.neg();
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!("Fp256-{} - inv (invert an Fp256)", limb_label),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(|| gen_rand_fp256(&mut rng), |a| a.inv());
        },
    );

    c.bench_function(
        &format!(
            "Fp256-{} - pow (exponentiate an Fp256 by an Fp256)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp256(&mut rng), gen_rand_fp256(&mut rng)),
                |(a, exp)| a.pow(exp),
            );
        },
    );

    c.bench_function(
        &format!(
            "Fp256-{} - pow (exponentiate an Fp256 by a u32)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp256(&mut rng), rng.next_u32()),
                |(a, exp)| a.pow(exp),
            );
        },
    );

    c.bench_function(
        &format!(
            "Fp256-{} - from, to_bytes_array (roundtrip to and from byte array 100 times)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || gen_rand_fp256(&mut rng),
                |a| {
                    for _ in 0..100 {
                        let byte_array = a.to_bytes_array();
                        let _ = fp_256::Fp256::from(byte_array);
                    }
                },
            );
        },
    );

    c.bench_function(&format!("Fp256-{} - from 64 bytes", limb_label), |bench| {
        let mut rng = rand::rng();
        bench.iter_with_setup(
            || gen_rand_sixty_four_bytes(&mut rng),
            |bytes| fp_256::Fp256::from(bytes),
        );
    });

    c.bench_function(
        &format!("Fp256-{} - Monty - Mul 100 times", limb_label),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || {
                    (
                        gen_rand_fp256(&mut rng).to_monty(),
                        gen_rand_fp256(&mut rng).to_monty(),
                    )
                },
                |(mut a, b)| {
                    for _ in 0..100 {
                        a = black_box(a * b);
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!("Fp256-{} - Monty - to_monty 100 times", limb_label),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || gen_rand_fp256(&mut rng),
                |a| {
                    for _ in 0..100 {
                        black_box(a.to_monty());
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!("Fp256-{} - Monty - to_norm 100 times", limb_label),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || gen_rand_fp256(&mut rng).to_monty(),
                |a| {
                    for _ in 0..100 {
                        black_box(a.to_norm());
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!(
            "Fp480-{} - normalize_little (480 bits to Fp480 100 times)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || gen_rand_fp480_raw(&mut rng),
                |val_to_norm| {
                    for _ in 0..100 {
                        black_box(val_to_norm.normalize_little());
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!(
            "Fp480-{} - add_assign (add an Fp480 into another Fp480 100 times)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp480(&mut rng), gen_rand_fp480(&mut rng)),
                |(mut a, b)| {
                    for _ in 0..100 {
                        a += b;
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!(
            "Fp480-{} - sub_assign (subtract an Fp480 from another Fp480 100 times)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp480(&mut rng), gen_rand_fp480(&mut rng)),
                |(mut a, b)| {
                    for _ in 0..100 {
                        a -= b;
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!(
            "Fp480-{} - mul_assign (an Fp480 into another Fp480 100 times)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || (gen_rand_fp480(&mut rng), gen_rand_fp480(&mut rng)),
                |(mut a, b)| {
                    for _ in 0..100 {
                        a *= b;
                    }
                },
            );
        },
    );

    c.bench_function(
        &format!("Fp480-{} - inv (invert an Fp480)", limb_label),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(|| gen_rand_fp480(&mut rng), |a| a.inv());
        },
    );

    c.bench_function(
        &format!(
            "Fp480-{} - from, to_bytes_array (roundtrip to and from byte array 100 times)",
            limb_label
        ),
        |bench| {
            let mut rng = rand::rng();
            bench.iter_with_setup(
                || gen_rand_fp480(&mut rng),
                |a| {
                    for _ in 0..100 {
                        let byte_array = a.to_bytes_array();
                        let _ = fp_480::Fp480::from(byte_array);
                    }
                },
            );
        },
    );

    c.bench_function(&format!("Fp480-{} - from 64 bytes", limb_label), |bench| {
        let mut rng = rand::rng();
        bench.iter_with_setup(
            || gen_rand_sixty_four_bytes(&mut rng),
            |bytes| fp_480::Fp480::from(bytes),
        );
    });
}

criterion_group! {
    name = benches;
    config = Criterion::default().sample_size(20);
    targets = criterion_benchmark
}
criterion_main!(benches);
