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

use cryptix_bn254::galoisfield::fp4::Fp4Element;
use cryptix_bn254::{fp, fp4};
use cryptix_bn254::galoisfield::{U256, BN254};
use cryptix_field::field::MulInverse;
use cryptix_field::field::primefield::FpElement;
use cryptix_field::field::montgomery::MontgomeryOps;

#[allow(non_upper_case_globals)]
const A: Fp4Element = fp4!(
    "141d50d1c6b86ab5ffb4b3f342c49ef2227026ef1a5982fbc7963cd5e302e5c8",
    "07ad43f30e437b75cfb8a7ddf522caf0594dc5131f417fba880b7a68548afea7",
    "1a295e6b77513f83fa06e1a71fd80a0a680d0cd8bae4066f8526f358f8d93e1e",
    "2267cd85ff0afe17d551c5c7aa6e565ec25bafa662c234d0b83fb614f9ff5892"
);

#[allow(non_upper_case_globals)]
const B: Fp4Element = fp4!(
    "1d47b6b2b86121c093a063c713900388e9712336d44e38add2a5a4edd8a436e2",
    "18b87ee9362339303bcb5729b9afa4f4552f8d9dfc41f7e74d553ce2b12bd2d0",
    "1366aae804d9afda1dc1cbc83a9919b20a57003e211f85d1a1f2ebe3364a6373",
    "118ae946a31a20385bd9b35444dea16acea7468c95c209022d42e71d16e397bb"
);

#[test]
fn test_fp4_add() {
    assert_eq!(A + B, fp4!(
        "0c41a3023f198c74d920ca3a5654a272aac04a25eea7bb95f33be1c3bba71c97",
        "2065c2dc4466b4a60b83ff07aed26fe4ae7d52b11b8377a1d560b74b05b6d177",
        "086ca4d13c2aef5c5d945fef5a7123b411430d16dc038c2d8019df3c2f23a17e",
        "0ecf524a62251e4e76f72b9bef4cf7c12fe1f632f8843dbf3e829d3210e2f03a"
    ));
}

#[test]
fn test_fp4_mul() {
    assert_eq!(A * B, fp4!(
        "2072854b18be081ec043ca25a85a1fc5e39134ebfbba5cbbe4f3d2e8544e351b",
        "0df94b403baec4effd7d68879fd318b2d56c1186e15b9f8c7a1f48211e98cb29",
        "06e4b415c7feb9a8fb56f25a82be00b6e2eec9218fc3c403409eef986b2814ad",
        "01d6e39b209f0a47c5337eb79e51602c2312d339cc64905ec5a30bc554515046"
    ));
}

#[test]
fn test_fp4_inv() {
    assert_eq!(A.mul_inv(), fp4!(
        "0b22fe30b0074676ad810f94455dcfa37ba18fba0e04c083cc935b22101b3589",
        "09521f94e05f575ca1c3b180c80281f9d9eb04c260552027f863063dc826a48b",
        "2312c97087cb80e161bcab057044c8aa25ceaf1d9c81f747e2764b270dcbea43",
        "1605519a88af82b748b0f5f4d4e18656c9ba4dcb41c9c6976212608217f09cf1"
    ))
}

#[test]
fn test_fp4_exp() {
    let e: FpElement<U256, BN254> = fp!("01");
    assert_eq!(A.mont_exp(e.repr()), A);
    
    let e: FpElement<U256, BN254> = fp!("02");
    assert_eq!(A.mont_exp(e.repr()), A * A);

    let e: FpElement<U256, BN254> = fp!("0c6e2e0f78060cb86e7764d8b2097b345318e96a7744fb01f0d76265949ea374");
    assert_eq!(A.mont_exp(e.repr()), fp4!(
        "0435ae85c13ac349f35a2f0deae4a613585217654990fa733ac57137fd099ecd",
        "0c6d290dbe65009f24265eabaae98aa3cf0a06f34b507671fd77b9ff054cefe7",
        "051996c3fcc7dd1ce7b91ae4ec4fe1e716195f6c5ccc66c532221005583bb0df",
        "1b62104763baee7be7961cdff55116ec2d8ec30c4acfcef804708bcb2d2fc6b4"
    ))
}

#[cfg(feature = "rand")]
#[test]
fn test_fp4_sqr_mul() {
    use rand::SeedableRng;
    let mut rng = rand_chacha::ChaCha8Rng::from_seed(SEED);
    for _ in 0..TEST_COUNT {
        let a = Fp4Element::rand(&mut rng).mont_form();
        let c = a.mont_sqr().mont_rdc();
        let d = a.mont_mul(a).mont_rdc();

        assert_eq!(c, d)
    }
}

#[cfg(feature = "rand")]
#[test]
fn test_fp4_mont_inv() {
    use rand::SeedableRng;
    use cryptix_field::field::MulIdentity;
    let mut rng = rand_chacha::ChaCha8Rng::from_seed(SEED);
    for _ in 0..TEST_COUNT {
        let a = Fp4Element::rand(&mut rng).mont_form();
        let b = a.mont_inv();
        assert_eq!(a.mont_mul(b).mont_rdc(), Fp4Element::ONE);
    }
}

#[cfg(feature = "rand")]
#[test]
fn test_fp4_map() {
    use rand::SeedableRng;
    use cryptix_field::Modular;
    let mut rng = rand_chacha::ChaCha8Rng::from_seed(SEED);
    for _ in 0..20 {
        let a = Fp4Element::rand(&mut rng).mont_form();
        let c = a.map_frob();
        let b = a.mont_exp(BN254::P).mont_rdc();
        let c = c.mont_rdc();

        assert_eq!(b, c);
    }
}