#![no_std]
#![macro_use]
extern crate ark_relations;

pub mod fields {
    use rand::{self, SeedableRng};
    use rand_xorshift::XorShiftRng;

    use ark_ff::{BitIteratorLE, Field, UniformRand};
    use ark_r1cs_std::prelude::*;
    use ark_relations::r1cs::{ConstraintSystem, SynthesisError};
    use ark_std::test_rng;
    use ark_std::vec::Vec;

    pub fn field_test<F, ConstraintF, AF>() -> Result<(), SynthesisError>
    where
        F: Field,
        ConstraintF: Field,
        AF: FieldVar<F, ConstraintF>,
        AF: TwoBitLookupGadget<ConstraintF, TableConstant = F>,
        for<'a> &'a AF: FieldOpsBounds<'a, F, AF>,
    {
        let cs = ConstraintSystem::<ConstraintF>::new_ref();

        let mut rng = test_rng();
        let a_native = F::rand(&mut rng);
        let b_native = F::rand(&mut rng);
        let a = AF::new_witness(ark_relations::ns!(cs, "generate_a"), || Ok(a_native))?;
        let b = AF::new_witness(ark_relations::ns!(cs, "generate_b"), || Ok(b_native))?;
        let b_const = AF::new_constant(ark_relations::ns!(cs, "b_as_constant"), b_native)?;

        let zero = AF::zero();
        let zero_native = zero.value()?;
        zero.enforce_equal(&zero)?;

        let one = AF::one();
        let one_native = one.value()?;
        one.enforce_equal(&one)?;

        one.enforce_not_equal(&zero)?;

        let one_dup = &zero + &one;
        one_dup.enforce_equal(&one)?;

        let two = &one + &one;
        two.enforce_equal(&two)?;
        two.enforce_equal(&one.double()?)?;
        two.enforce_not_equal(&one)?;
        two.enforce_not_equal(&zero)?;

        // a + 0 = a
        let a_plus_zero = &a + &zero;
        assert_eq!(a_plus_zero.value()?, a_native);
        a_plus_zero.enforce_equal(&a)?;
        a_plus_zero.enforce_not_equal(&a.double()?)?;

        // a - 0 = a
        let a_minus_zero = &a - &zero;
        assert_eq!(a_minus_zero.value()?, a_native);
        a_minus_zero.enforce_equal(&a)?;

        // a - a = 0
        let a_minus_a = &a - &a;
        assert_eq!(a_minus_a.value()?, zero_native);
        a_minus_a.enforce_equal(&zero)?;

        // a + b = b + a
        let a_b = &a + &b;
        let b_a = &b + &a;
        assert_eq!(a_b.value()?, a_native + &b_native);
        a_b.enforce_equal(&b_a)?;

        // (a + b) + a = a + (b + a)
        let ab_a = &a_b + &a;
        let a_ba = &a + &b_a;
        assert_eq!(ab_a.value()?, a_native + &b_native + &a_native);
        ab_a.enforce_equal(&a_ba)?;

        let b_times_a_plus_b = &a_b * &b;
        let b_times_b_plus_a = &b_a * &b;
        assert_eq!(
            b_times_a_plus_b.value()?,
            b_native * &(b_native + &a_native)
        );
        assert_eq!(
            b_times_a_plus_b.value()?,
            (b_native + &a_native) * &b_native
        );
        assert_eq!(
            b_times_a_plus_b.value()?,
            (a_native + &b_native) * &b_native
        );
        b_times_b_plus_a.enforce_equal(&b_times_a_plus_b)?;

        // a * 1 = a
        assert_eq!((&a * &one).value()?, a_native * &one_native);

        // a * b = b * a
        let ab = &a * &b;
        let ba = &b * &a;
        assert_eq!(ab.value()?, ba.value()?);
        assert_eq!(ab.value()?, a_native * &b_native);

        let ab_const = &a * &b_const;
        let b_const_a = &b_const * &a;
        assert_eq!(ab_const.value()?, b_const_a.value()?);
        assert_eq!(ab_const.value()?, ab.value()?);
        assert_eq!(ab_const.value()?, a_native * &b_native);

        // (a * b) * a = a * (b * a)
        let ab_a = &ab * &a;
        let a_ba = &a * &ba;
        assert_eq!(ab_a.value()?, a_ba.value()?);
        assert_eq!(ab_a.value()?, a_native * &b_native * &a_native);

        let aa = &a * &a;
        let a_squared = a.square()?;
        a_squared.enforce_equal(&aa)?;
        assert_eq!(aa.value()?, a_squared.value()?);
        assert_eq!(aa.value()?, a_native.square());

        let aa = &a * a.value()?;
        a_squared.enforce_equal(&aa)?;
        assert_eq!(aa.value()?, a_squared.value()?);
        assert_eq!(aa.value()?, a_native.square());

        let a_b2 = &a + b_native;
        a_b.enforce_equal(&a_b2)?;
        assert_eq!(a_b.value()?, a_b2.value()?);

        let a_inv = a.inverse()?;
        a_inv.mul_equals(&a, &one)?;
        assert_eq!(a_inv.value()?, a.value()?.inverse().unwrap());
        assert_eq!(a_inv.value()?, a_native.inverse().unwrap());

        let a_b_inv = a.mul_by_inverse(&b)?;
        a_b_inv.mul_equals(&b, &a)?;
        assert_eq!(a_b_inv.value()?, a_native * b_native.inverse().unwrap());

        // a * a * a = a^3
        let bits = BitIteratorLE::without_trailing_zeros([3u64])
            .map(Boolean::constant)
            .collect::<Vec<_>>();
        assert_eq!(a_native.pow([0x3]), a.pow_le(&bits)?.value()?);

        // a * a * a = a^3
        assert_eq!(a_native.pow([0x3]), a.pow_by_constant(&[0x3])?.value()?);
        assert!(cs.is_satisfied().unwrap());

        // a * a * a = a^3
        let mut constants = [F::zero(); 4];
        for c in &mut constants {
            *c = UniformRand::rand(&mut test_rng());
        }
        let bits = [
            Boolean::<ConstraintF>::constant(false),
            Boolean::constant(true),
        ];
        let lookup_result = AF::two_bit_lookup(&bits, constants.as_ref())?;
        assert_eq!(lookup_result.value()?, constants[2]);
        assert!(cs.is_satisfied().unwrap());

        let f = F::from(1u128 << 64);
        let f_bits = ark_ff::BitIteratorLE::new(&[0u64, 1u64]).collect::<Vec<_>>();
        let fv = AF::new_witness(ark_relations::ns!(cs, "alloc u128"), || Ok(f))?;
        assert_eq!(fv.to_bits_le()?.value().unwrap()[..128], f_bits[..128]);
        assert!(cs.is_satisfied().unwrap());

        let r_native: F = UniformRand::rand(&mut test_rng());

        let r = AF::new_witness(ark_relations::ns!(cs, "r_native"), || Ok(r_native)).unwrap();
        let _ = r.to_non_unique_bits_le()?;
        assert!(cs.is_satisfied().unwrap());
        let _ = r.to_bits_le()?;
        assert!(cs.is_satisfied().unwrap());

        let bytes = r.to_non_unique_bytes()?;
        assert_eq!(ark_ff::to_bytes!(r_native).unwrap(), bytes.value().unwrap());
        assert!(cs.is_satisfied().unwrap());
        let bytes = r.to_bytes()?;
        assert_eq!(ark_ff::to_bytes!(r_native).unwrap(), bytes.value().unwrap());
        assert!(cs.is_satisfied().unwrap());

        let ab_false = &a + (AF::from(Boolean::Constant(false)) * b_native);
        assert_eq!(ab_false.value()?, a_native);
        let ab_true = &a + (AF::from(Boolean::Constant(true)) * b_native);
        assert_eq!(ab_true.value()?, a_native + &b_native);

        if !cs.is_satisfied().unwrap() {
            panic!("{:?}", cs.which_is_unsatisfied().unwrap());
        }
        assert!(cs.is_satisfied().unwrap());
        Ok(())
    }

    pub fn frobenius_tests<F: Field, ConstraintF, AF>(maxpower: usize) -> Result<(), SynthesisError>
    where
        F: Field,
        ConstraintF: Field,
        AF: FieldVar<F, ConstraintF>,
        for<'a> &'a AF: FieldOpsBounds<'a, F, AF>,
    {
        let cs = ConstraintSystem::<ConstraintF>::new_ref();
        let mut rng = XorShiftRng::seed_from_u64(1231275789u64);
        for i in 0..=maxpower {
            let mut a = F::rand(&mut rng);
            let mut a_gadget = AF::new_witness(ark_relations::ns!(cs, "a"), || Ok(a))?;
            a_gadget.frobenius_map_in_place(i)?;
            a.frobenius_map(i);

            assert_eq!(a_gadget.value()?, a);
        }

        assert!(cs.is_satisfied().unwrap());
        Ok(())
    }
}

pub mod curves {
    use ark_ec::{
        short_weierstrass_jacobian::GroupProjective as SWProjective,
        twisted_edwards_extended::GroupProjective as TEProjective, AffineCurve, ProjectiveCurve,
    };
    use ark_ff::{test_rng, Field, PrimeField};
    use ark_relations::r1cs::{ConstraintSystem, SynthesisError};
    use ark_std::vec::Vec;

    use ark_r1cs_std::prelude::*;

    pub fn group_test<C, ConstraintF, GG>() -> Result<(), SynthesisError>
    where
        C: ProjectiveCurve,
        ConstraintF: Field,
        GG: CurveVar<C, ConstraintF>,
        for<'a> &'a GG: GroupOpsBounds<'a, C, GG>,
    {
        let cs = ConstraintSystem::<ConstraintF>::new_ref();

        let mut rng = test_rng();
        let a_native = C::rand(&mut rng);
        let b_native = C::rand(&mut rng);
        let a = GG::new_witness(ark_relations::ns!(cs, "generate_a"), || Ok(a_native)).unwrap();
        let b = GG::new_witness(ark_relations::ns!(cs, "generate_b"), || Ok(b_native)).unwrap();

        let zero = GG::zero();
        assert_eq!(zero.value()?, zero.value()?);

        // a == a
        assert_eq!(a.value()?, a.value()?);
        // a + 0 = a
        assert_eq!((&a + &zero).value()?, a.value()?);
        // a - 0 = a
        assert_eq!((&a - &zero).value()?, a.value()?);
        // a - a = 0
        assert_eq!((&a - &a).value()?, zero.value()?);
        // a + b = b + a
        let a_b = &a + &b;
        let b_a = &b + &a;
        assert_eq!(a_b.value()?, b_a.value()?);
        a_b.enforce_equal(&b_a)?;
        assert!(cs.is_satisfied().unwrap());

        // (a + b) + a = a + (b + a)
        let ab_a = &a_b + &a;
        let a_ba = &a + &b_a;
        assert_eq!(ab_a.value()?, a_ba.value()?);
        ab_a.enforce_equal(&a_ba)?;
        assert!(cs.is_satisfied().unwrap());

        // a.double() = a + a
        let a_a = &a + &a;
        let mut a2 = a.clone();
        a2.double_in_place()?;
        a2.enforce_equal(&a_a)?;
        assert_eq!(a2.value()?, a_native.double());
        assert_eq!(a_a.value()?, a_native.double());
        assert_eq!(a2.value()?, a_a.value()?);
        assert!(cs.is_satisfied().unwrap());

        // b.double() = b + b
        let mut b2 = b.clone();
        b2.double_in_place()?;
        let b_b = &b + &b;
        b2.enforce_equal(&b_b)?;
        assert!(cs.is_satisfied().unwrap());
        assert_eq!(b2.value()?, b_b.value()?);

        let _ = a.to_bytes()?;
        assert!(cs.is_satisfied().unwrap());
        let _ = a.to_non_unique_bytes()?;
        assert!(cs.is_satisfied().unwrap());

        let _ = b.to_bytes()?;
        let _ = b.to_non_unique_bytes()?;
        if !cs.is_satisfied().unwrap() {
            panic!("{:?}", cs.which_is_unsatisfied().unwrap());
        }
        assert!(cs.is_satisfied().unwrap());
        Ok(())
    }

    pub fn sw_test<P, GG>() -> Result<(), SynthesisError>
    where
        P: ark_ec::SWModelParameters,
        GG: CurveVar<SWProjective<P>, <P::BaseField as Field>::BasePrimeField>,
        for<'a> &'a GG: GroupOpsBounds<'a, SWProjective<P>, GG>,
    {
        use ark_ec::group::Group;
        use ark_ff::{BitIteratorLE, UniformRand};
        use ark_r1cs_std::prelude::*;

        group_test::<SWProjective<P>, _, GG>()?;

        let mut rng = test_rng();

        let cs = ConstraintSystem::<<P::BaseField as Field>::BasePrimeField>::new_ref();

        let a = SWProjective::<P>::rand(&mut rng);
        let b = SWProjective::<P>::rand(&mut rng);
        let a_affine = a.into_affine();
        let b_affine = b.into_affine();

        let ns = ark_relations::ns!(cs, "allocating variables");
        let mut gadget_a = GG::new_witness(cs.clone(), || Ok(a))?;
        let gadget_b = GG::new_witness(cs.clone(), || Ok(b))?;
        drop(ns);
        assert_eq!(gadget_a.value()?.into_affine().x, a_affine.x);
        assert_eq!(gadget_a.value()?.into_affine().y, a_affine.y);
        assert_eq!(gadget_b.value()?.into_affine().x, b_affine.x);
        assert_eq!(gadget_b.value()?.into_affine().y, b_affine.y);
        assert_eq!(cs.which_is_unsatisfied().unwrap(), None);

        // Check addition
        let ab = a + &b;
        let ab_affine = ab.into_affine();
        let gadget_ab = &gadget_a + &gadget_b;
        let gadget_ba = &gadget_b + &gadget_a;
        gadget_ba.enforce_equal(&gadget_ab)?;

        let ab_val = gadget_ab.value()?.into_affine();
        assert_eq!(ab_val, ab_affine, "Result of addition is unequal");
        assert!(cs.is_satisfied().unwrap());

        // Check doubling
        let aa = Group::double(&a);
        let aa_affine = aa.into_affine();
        gadget_a.double_in_place()?;
        let aa_val = gadget_a.value()?.into_affine();
        assert_eq!(
            aa_val, aa_affine,
            "Gadget and native values are unequal after double."
        );
        assert!(cs.is_satisfied().unwrap());

        // Check mul_bits
        let scalar = P::ScalarField::rand(&mut rng);
        let native_result = aa.into_affine().mul(scalar);
        let native_result = native_result.into_affine();

        let scalar: Vec<bool> = BitIteratorLE::new(scalar.into_repr()).collect();
        let input: Vec<Boolean<_>> =
            Vec::new_witness(ark_relations::ns!(cs, "bits"), || Ok(scalar)).unwrap();
        let result = gadget_a.scalar_mul_le(input.iter())?;
        let result_val = result.value()?.into_affine();
        assert_eq!(
            result_val, native_result,
            "gadget & native values are diff. after scalar mul"
        );
        assert!(cs.is_satisfied().unwrap());

        if !cs.is_satisfied().unwrap() {
            panic!("{:?}", cs.which_is_unsatisfied().unwrap());
        }

        assert!(cs.is_satisfied().unwrap());
        Ok(())
    }

    pub fn te_test<P, GG>() -> Result<(), SynthesisError>
    where
        P: ark_ec::TEModelParameters,
        GG: CurveVar<TEProjective<P>, <P::BaseField as Field>::BasePrimeField>,
        for<'a> &'a GG: GroupOpsBounds<'a, TEProjective<P>, GG>,
    {
        use ark_ec::group::Group;
        use ark_ff::{BitIteratorLE, UniformRand};

        group_test::<TEProjective<P>, _, GG>()?;

        let mut rng = test_rng();

        let cs = ConstraintSystem::<<P::BaseField as Field>::BasePrimeField>::new_ref();

        let a = TEProjective::<P>::rand(&mut rng);
        let b = TEProjective::<P>::rand(&mut rng);
        let a_affine = a.into_affine();
        let b_affine = b.into_affine();

        let ns = ark_relations::ns!(cs, "allocating variables");
        let mut gadget_a = GG::new_witness(cs.clone(), || Ok(a))?;
        let gadget_b = GG::new_witness(cs.clone(), || Ok(b))?;
        drop(ns);

        assert_eq!(gadget_a.value()?.into_affine().x, a_affine.x);
        assert_eq!(gadget_a.value()?.into_affine().y, a_affine.y);
        assert_eq!(gadget_b.value()?.into_affine().x, b_affine.x);
        assert_eq!(gadget_b.value()?.into_affine().y, b_affine.y);
        assert_eq!(cs.which_is_unsatisfied()?, None);

        // Check addition
        let ab = a + &b;
        let ab_affine = ab.into_affine();
        let gadget_ab = &gadget_a + &gadget_b;
        let gadget_ba = &gadget_b + &gadget_a;
        gadget_ba.enforce_equal(&gadget_ab)?;

        let ab_val = gadget_ab.value()?.into_affine();
        assert_eq!(ab_val, ab_affine, "Result of addition is unequal");
        assert!(cs.is_satisfied().unwrap());

        // Check doubling
        let aa = Group::double(&a);
        let aa_affine = aa.into_affine();
        gadget_a.double_in_place()?;
        let aa_val = gadget_a.value()?.into_affine();
        assert_eq!(
            aa_val, aa_affine,
            "Gadget and native values are unequal after double."
        );
        assert!(cs.is_satisfied().unwrap());

        // Check mul_bits
        let scalar = P::ScalarField::rand(&mut rng);
        let native_result = AffineCurve::mul(&aa.into_affine(), scalar);
        let native_result = native_result.into_affine();

        let scalar: Vec<bool> = BitIteratorLE::new(scalar.into_repr()).collect();
        let input: Vec<Boolean<_>> =
            Vec::new_witness(ark_relations::ns!(cs, "bits"), || Ok(scalar)).unwrap();
        let result = gadget_a.scalar_mul_le(input.iter())?;
        let result_val = result.value()?.into_affine();
        assert_eq!(
            result_val, native_result,
            "gadget & native values are diff. after scalar mul"
        );
        assert!(cs.is_satisfied().unwrap());

        if !cs.is_satisfied().unwrap() {
            panic!("{:?}", cs.which_is_unsatisfied().unwrap());
        }

        assert!(cs.is_satisfied().unwrap());
        Ok(())
    }
}

pub mod pairing {
    use ark_ec::{PairingEngine, ProjectiveCurve};
    use ark_ff::{test_rng, BitIteratorLE, Field, PrimeField, UniformRand};
    use ark_r1cs_std::prelude::*;
    use ark_relations::r1cs::{ConstraintSystem, SynthesisError};
    use ark_std::vec::Vec;

    #[allow(dead_code)]
    pub fn bilinearity_test<E: PairingEngine, P: PairingVar<E>>() -> Result<(), SynthesisError>
    where
        for<'a> &'a P::G1Var: GroupOpsBounds<'a, E::G1Projective, P::G1Var>,
        for<'a> &'a P::G2Var: GroupOpsBounds<'a, E::G2Projective, P::G2Var>,
        for<'a> &'a P::GTVar: FieldOpsBounds<'a, E::Fqk, P::GTVar>,
    {
        let cs = ConstraintSystem::<E::Fq>::new_ref();

        let mut rng = test_rng();
        let a = E::G1Projective::rand(&mut rng);
        let b = E::G2Projective::rand(&mut rng);
        let s = E::Fr::rand(&mut rng);

        let mut sa = a;
        sa *= s;
        let mut sb = b;
        sb *= s;

        let a_g = P::G1Var::new_witness(cs.clone(), || Ok(a.into_affine()))?;
        let b_g = P::G2Var::new_witness(cs.clone(), || Ok(b.into_affine()))?;
        let sa_g = P::G1Var::new_witness(cs.clone(), || Ok(sa.into_affine()))?;
        let sb_g = P::G2Var::new_witness(cs.clone(), || Ok(sb.into_affine()))?;

        let mut _preparation_num_constraints = cs.num_constraints();
        let a_prep_g = P::prepare_g1(&a_g)?;
        let b_prep_g = P::prepare_g2(&b_g)?;
        _preparation_num_constraints = cs.num_constraints() - _preparation_num_constraints;

        let sa_prep_g = P::prepare_g1(&sa_g)?;
        let sb_prep_g = P::prepare_g2(&sb_g)?;

        let (ans1_g, ans1_n) = {
            let _ml_constraints = cs.num_constraints();
            let ml_g = P::miller_loop(&[sa_prep_g], &[b_prep_g.clone()])?;
            let _fe_constraints = cs.num_constraints();
            let ans_g = P::final_exponentiation(&ml_g)?;
            let ans_n = E::pairing(sa, b);
            (ans_g, ans_n)
        };

        let (ans2_g, ans2_n) = {
            let ans_g = P::pairing(a_prep_g.clone(), sb_prep_g)?;
            let ans_n = E::pairing(a, sb);
            (ans_g, ans_n)
        };

        let (ans3_g, ans3_n) = {
            let s_iter = BitIteratorLE::without_trailing_zeros(s.into_repr())
                .map(Boolean::constant)
                .collect::<Vec<_>>();

            let mut ans_g = P::pairing(a_prep_g, b_prep_g)?;
            let mut ans_n = E::pairing(a, b);
            ans_n = ans_n.pow(s.into_repr());
            ans_g = ans_g.pow_le(&s_iter)?;

            (ans_g, ans_n)
        };

        ans1_g.enforce_equal(&ans2_g)?;
        ans2_g.enforce_equal(&ans3_g)?;

        assert_eq!(ans1_g.value()?, ans1_n, "Failed native test 1");
        assert_eq!(ans2_g.value()?, ans2_n, "Failed native test 2");
        assert_eq!(ans3_g.value()?, ans3_n, "Failed native test 3");

        assert_eq!(ans1_n, ans2_n, "Failed ans1_native == ans2_native");
        assert_eq!(ans2_n, ans3_n, "Failed ans2_native == ans3_native");
        assert_eq!(ans1_g.value()?, ans3_g.value()?, "Failed ans1 == ans3");
        assert_eq!(ans1_g.value()?, ans2_g.value()?, "Failed ans1 == ans2");
        assert_eq!(ans2_g.value()?, ans3_g.value()?, "Failed ans2 == ans3");

        if !cs.is_satisfied().unwrap() {
            panic!("Unsatisfied: {:?}", cs.which_is_unsatisfied());
        }

        assert!(cs.is_satisfied().unwrap(), "cs is not satisfied");
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn it_works() {
        assert_eq!(2 + 2, 4);
    }
}