use algebra_core::Field;
use r1cs_core::{ConstraintSynthesizer, ConstraintSystem, SynthesisError};

#[derive(Copy, Clone)]
struct Circuit<F: Field> {
    d: Option<F>,
    e: Option<F>,
}

impl<ConstraintF: Field> ConstraintSynthesizer<ConstraintF> for Circuit<ConstraintF> {
    fn generate_constraints<CS: ConstraintSystem<ConstraintF>>(
        self,
        cs: &mut CS,
    ) -> Result<(), SynthesisError> {
        println!("generating constraints for test");
        let z = cs.alloc(|| "z", || Ok(ConstraintF::one()))?;
        let five = ConstraintF::one() + ConstraintF::one() + ConstraintF::one() + ConstraintF::one() + ConstraintF::one(); 
        let a = cs.alloc(|| "a", || Ok(five))?;
        let b = cs.alloc(|| "b", || Ok(ConstraintF::zero()))?;
        let c = cs.alloc(|| "b", || Ok(five*five))?;
        let d = cs.alloc(|| "b", || self.d.ok_or(SynthesisError::AssignmentMissing))?;
        let e = cs.alloc(|| "b", || self.e.ok_or(SynthesisError::AssignmentMissing))?;
        
        cs.enforce(
            || format!("constraint 1"),
            |lc| lc + a,
            |lc| lc + a,
            |lc| lc + d,
        );

        cs.enforce(
            || format!("constraint 2"),
            |lc| lc + b,
            |lc| lc + b,
            |lc| lc + e,
        );

        cs.enforce(
            || format!("constraint 3"),
            |lc| lc + z,
            |lc| lc + d + e,
            |lc| lc + c,
        );

        println!("finished generating test constraints");
        Ok(())
    }
}

mod marlin {
    use super::*;
    use crate::Marlin;

    use algebra::{bls12_381::Fr, Bls12_381};
    use blake2::Blake2s;
    use core::ops::MulAssign;
    use num_traits::identities::{One, Zero};
    use poly_commit::marlin_pc::MarlinKZG10;

    type MultiPC = MarlinKZG10<Bls12_381>;
    type MarlinInst = Marlin<Fr, MultiPC, Blake2s>;

    #[test]
    fn test_circuit() {
        let rng = &mut algebra::test_rng();
        let universal_srs = MarlinInst::universal_setup(100, 25, 100, rng).unwrap();

        //let z = Fr::one();
        let a = Fr::one() + Fr::one() + Fr::one() + Fr::one() + Fr::one();
        let b = Fr::zero();
        let mut c = a;
        c.mul_assign(&a);
        let mut d = a;
        d.mul_assign(&a);
        let e = Fr::zero();

        let circ = Circuit {
            d: Some(d),
            e: Some(e),
        };

        let (index_pk, index_vk) = MarlinInst::index(&universal_srs, circ.clone()).unwrap();
        println!("Called index");

        let proof = MarlinInst::prove(&index_pk, circ, rng).unwrap();
        println!("Called prover");

        //let temp = &[z,a,b,c];
        //println!("usize {} count_ones {}", temp.len(), temp.len().count_ones());

        assert!(MarlinInst::verify(&index_vk, &[a,b,c], &proof, rng).unwrap());
        println!("Called verifier");
        println!("\nShould not verify (i.e. verifier messages should print below):");
        assert!(!MarlinInst::verify(&index_vk, &[a], &proof, rng).unwrap());
    }

}