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

mod bn254;
use bn254::{U256, BN254};
use cryptix_bigint::bigint;
use cryptix_field::{Element, field::{montgomery::{Montgomery, MontgomeryOps}, MulInverse, MulIdentity}};
use cryptix_field::field::primefield::FpElement;

#[test]
fn test_montgomery() {
    let a: FpElement<U256, BN254> = FpElement(
        bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648").into()
    );
    let b: FpElement<U256, BN254> = FpElement(
        bigint!(U256, "0f44148ca6427affe2d16f7402ad1da381fb3002db6aa1c0b414a52e715abe88").into()
    );
    let c = a + b;

    assert_eq!(a.repr(), bigint!(U256, "0ad7c538859337e08ebea0b5cfbe18db78c5f6eea014b7bbda29a88101a76622"));
    assert_eq!(b.repr(), bigint!(U256, "0f44148ca6427affe2d16f7402ad1da381fb3002db6aa1c0b414a52e715abe88"));
    assert_eq!(c, bigint!(U256, "1a1bd9c52bd5b2e071901029d26b367efac126f17b7f597c8e3e4daf730224aa").into());
    assert_eq!(a * b, bigint!(U256, "0f8dcf5d18658b1ff8eb23bd0ae4fd4801689488098f2f5c185f9802964e39bd").into());

    // arm = a * R
    let arm = a.mont_mul(BN254::RR_P);
    assert_eq!(arm.repr(), bigint!(U256, "1027875403430b1a3dc34ddf239bd4db28cf5019cda9c9b3529122b3de7d0f95"));
    
    // abrm = a * b * R^(-1)
    let abrm = a.mont_mul(b);
    assert_eq!(abrm.repr(), bigint!(U256, "0216be32303d9cf1f1256f6bbd63c823acc53b5c690ebec6bf9a774f76125f63"));

    assert_eq!(abrm.mont_mul(BN254::RR_P), a * b);
}

#[test]
fn test_mod_inv() {
    let a: FpElement<U256, BN254> = FpElement(
        bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648").into()
    );

    assert_eq!(a.mul_inv().repr(), bigint!(U256, "0c0ccee68605383a5f08ba36b05b7be41c6c9e21bed3b715f8e388047e7a7f3f"));
}

#[test]
fn test_mod_exp() {
    // a^1 = a
    let a: FpElement<U256, BN254> = FpElement(
        bigint!(U256, "551e8e3d059337e403273bb5cfbe18ec3b07f6eea014b7e32829a88101a76648").into()
    );
    assert_eq!(a.mont_exp(FpElement::<U256, BN254>::ONE.repr()), a);

    // a^2 = a * a
    let two: FpElement<U256, BN254> = FpElement::ONE + FpElement::ONE;
    assert_eq!(a.mont_exp(two.repr()), a.mont_sqr().mont_form());

    // a ^ 65536 = 0x19c560e6d983782e87ed3259d0de1dec7b2e209f404178badfdff0c6d7531141
    let tmp: FpElement<U256, BN254> = FpElement(Element::new(U256::from(65536)));
    assert_eq!(a.mont_exp(tmp.repr()).repr(), bigint!(U256, "19c560e6d983782e87ed3259d0de1dec7b2e209f404178badfdff0c6d7531141"));

    // a ^ a = 0x8449fd9613c0461cd38feeca4fa5f6ec12507accbde5c548b25ebd63264c4c3
    assert_eq!(a.mont_exp(a.repr()).repr(), bigint!(U256, "08449fd9613c0461cd38feeca4fa5f6ec12507accbde5c548b25ebd63264c4c3"));
}