#![allow(incomplete_features)]
#![feature(generic_const_exprs)]
#![feature(const_trait_impl)]

use criterion::{black_box, criterion_group, criterion_main, Criterion};
use cryptix_bigint::{BigUInt, bigint, property::IsBigInt};
use cryptix_field::field::montgomery::MontgomeryOps;
use cryptix_field::{Element, Modular, field::montgomery::Montgomery, PrimeModular, OddModular};
use cryptix_field::field::primefield::FpElement;

type U256 = BigUInt<u64, 4>;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct BN254;

impl Modular<U256> for BN254 {
    const P: U256 = bigint!(U256, "2523648240000001BA344D80000000086121000000000013A700000000000013");
}

/// # Safety
/// 
/// the modular P is indeed a prime number, this comes from [the parameter of ec BN254](https://neuromancer.sk/std/bn/bn254#)
impl PrimeModular<U256> for BN254 { }

/// # Safety
/// 
/// P is odd
impl OddModular<U256> for BN254 { }

impl Montgomery<U256> for BN254 {
    const R_P: FpElement<U256, Self> = {
        FpElement(Element::new_unchecked(bigint!(U256, "212ba4f27ffffff5a2c62effffffffcdb939ffffffffff8a15ffffffffffff8e")))
    };

    const R_INV_P: FpElement<U256, Self> = {
        FpElement(Element::new_unchecked(bigint!(U256, "1a7344bac91f117ea513ec0ed5682406b6c15140174d61b28b762ae9cf6d3b46")))
    };
    
    /*
     * # Safety
     * 
     * this is the pre-computed value of R * R mod P, must less than P
     */
    const RR_P: FpElement<U256, BN254> = {
        FpElement(Element::new_unchecked(bigint!(U256, "1B0A32FDF6403A3D281E3A1B7F86954F55EFBF6E8C1CC3F1B3E886745370473D")))
    };
    const NEG_P_INV_B: <U256 as IsBigInt>::Dig = 0x8435e50d79435e5;
}


fn bench_fp(c: &mut Criterion) {
    // c.bench_function("bench ring multiplication", |b| {
    //     let x: Element<U256, BN254> = bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648").into();
    //     let y: Element<U256, BN254> = bigint!(U256, "0f44148ca6427affe2d16f7402ad1da381fb3002db6aa1c0b414a52e715abe88").into();
    //     b.iter(|| black_box(x * y))
    // });

    c.bench_function("bench addition", |b| {
        let x: Element<U256, BN254> = bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648").into();
        let y: Element<U256, BN254> = bigint!(U256, "0f44148ca6427affe2d16f7402ad1da381fb3002db6aa1c0b414a52e715abe88").into();

        let x = FpElement(x);
        let y = FpElement(y);
        b.iter(|| black_box(x + y))
    });

    c.bench_function("bench montgomery multiplication", |b| {
        let x: Element<U256, BN254> = bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648").into();
        let y: Element<U256, BN254> = bigint!(U256, "0f44148ca6427affe2d16f7402ad1da381fb3002db6aa1c0b414a52e715abe88").into();

        let x = FpElement(x);
        let y = FpElement(y);
        b.iter(|| black_box(x.mont_mul(y)))
    });

    c.bench_function("bench mont rdc multiplication", |b| {
        let x: Element<U256, BN254> = bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648").into();
        let y: Element<U256, BN254> = bigint!(U256, "0f44148ca6427affe2d16f7402ad1da381fb3002db6aa1c0b414a52e715abe88").into();

        let x = FpElement(x);
        let y = FpElement(y);
        b.iter(|| black_box(x.mont_form().mont_mul(y)))
    });

    c.bench_function("bench normal multiplication wrap", |b| {
        use cryptix_bigint::ops::WrapAround;
        let x: U256 = bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648");
        let y: U256 = bigint!(U256, "0f44148ca6427affe2d16f7402ad1da381fb3002db6aa1c0b414a52e715abe88");

        b.iter(|| black_box(BN254::P.wrap(x * y)))
    });

    c.bench_function("bench normal multiplication", |b| {
        let x: U256 = bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648");
        let y: U256 = bigint!(U256, "0f44148ca6427affe2d16f7402ad1da381fb3002db6aa1c0b414a52e715abe88");

        b.iter(|| black_box(x * y))
    });

    c.bench_function("bench normal addition", |b| {
        let x: U256 = bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648");
        let y: U256 = bigint!(U256, "0f44148ca6427affe2d16f7402ad1da381fb3002db6aa1c0b414a52e715abe88");

        b.iter(|| black_box(x + y))
    });

    c.bench_function("bench normal mod", |b| {
        use cryptix_bigint::ops::WrapAround;
        let x: U256 = bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648");
        let y: U256 = bigint!(U256, "0f44148ca6427affe2d16f7402ad1da381fb3002db6aa1c0b414a52e715abe88");
        let xy = x * y;

        b.iter(|| black_box(BN254::P.wrap(xy)))
    });

    c.bench_function("bench montgomery square", |b| {
        let x: Element<U256, BN254> = bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648").into();
        
        let x = FpElement(x);
        b.iter(|| black_box(x.mont_sqr()))
    });

    c.bench_function("bench montgomery inversion", |b| {
        let x: Element<U256, BN254> = bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648").into();
        
        let x = FpElement(x);
        b.iter(|| black_box(x.mont_inv()))
    });
}

criterion_group!(benches, bench_fp);
criterion_main!(benches);