#[cfg(test)]
pub mod test;

pub mod boolean;
pub mod uint32;
pub mod blake2s;
pub mod num;
pub mod lookup;
pub mod ecc;
pub mod pedersen_hash;

use pairing::{
    PrimeField,
    PrimeFieldRepr,
};

use bellman::{
    SynthesisError,
    ConstraintSystem,
    Circuit
};

use jubjub::{
    JubjubEngine,
    Unknown,
    FixedGenerators,
    edwards
};

trait Assignment<T> {
    fn get(&self) -> Result<&T, SynthesisError>;
}

impl<T> Assignment<T> for Option<T> {
    fn get(&self) -> Result<&T, SynthesisError> {
        match *self {
            Some(ref v) => Ok(v),
            None => Err(SynthesisError::AssignmentMissing)
        }
    }
}

const MERKLE_TREE_DEPTH: usize = 29;

pub struct Spend<'a, E: JubjubEngine> {
    pub params: &'a E::Params,
    /// Value of the note being spent
    pub value: Option<u64>,
    /// Randomness that will hide the value
    pub value_randomness: Option<E::Fs>,
    /// Key which allows the proof to be constructed
    /// as defense-in-depth against a flaw in the
    /// protocol that would otherwise be exploitable
    /// by a holder of a viewing key.
    pub rsk: Option<E::Fs>,
    /// The public key that will be re-randomized for
    /// use as a nullifier and signing key for the
    /// transaction.
    pub ak: Option<edwards::Point<E, Unknown>>,
    /// The diversified base used to compute pk_d.
    pub g_d: Option<edwards::Point<E, Unknown>>,
    /// The randomness used to hide the note commitment data
    pub commitment_randomness: Option<E::Fs>,
    /// The authentication path of the commitment in the tree
    pub auth_path: Vec<Option<(E::Fr, bool)>>
}

impl<'a, E: JubjubEngine> Circuit<E> for Spend<'a, E> {
    fn synthesize<CS: ConstraintSystem<E>>(self, cs: &mut CS) -> Result<(), SynthesisError>
    {
        // Booleanize the value into little-endian bit order
        let value_bits = boolean::u64_into_allocated_bits_be(
            cs.namespace(|| "value"),
            self.value
        )?
        .into_iter()
        .rev() // Little endian bit order
        .map(|e| boolean::Boolean::from(e))
        .collect::<Vec<_>>();

        {
            let gv = ecc::fixed_base_multiplication(
                cs.namespace(|| "compute the value in the exponent"),
                FixedGenerators::ValueCommitmentValue,
                &value_bits,
                self.params
            )?;
        
            // Booleanize the randomness
            let hr = boolean::field_into_allocated_bits_be(
                cs.namespace(|| "hr"),
                self.value_randomness
            )?
            .into_iter()
            .rev() // Little endian bit order
            .map(|e| boolean::Boolean::from(e))
            .collect::<Vec<_>>();

            let hr = ecc::fixed_base_multiplication(
                cs.namespace(|| "computation of randomization for value commitment"),
                FixedGenerators::ValueCommitmentRandomness,
                &hr,
                self.params
            )?;

            let gvhr = gv.add(
                cs.namespace(|| "computation of value commitment"),
                &hr,
                self.params
            )?;

            // Expose the value commitment publicly
            let value_commitment_x = cs.alloc_input(
                || "value commitment x",
                || {
                    Ok(*gvhr.x.get_value().get()?)
                }
            )?;

            cs.enforce(
                || "value commitment x equals input",
                |lc| lc + value_commitment_x,
                |lc| lc + CS::one(),
                |lc| lc + gvhr.x.get_variable()
            );

            let value_commitment_y = cs.alloc_input(
                || "value commitment y",
                || {
                    Ok(*gvhr.y.get_value().get()?)
                }
            )?;

            cs.enforce(
                || "value commitment y equals input",
                |lc| lc + value_commitment_y,
                |lc| lc + CS::one(),
                |lc| lc + gvhr.y.get_variable()
            );
        }

        // Compute rk = [rsk] ProvingPublicKey
        let rk;
        {
            // Witness rsk as bits
            let rsk = boolean::field_into_allocated_bits_be(
                cs.namespace(|| "rsk"),
                self.rsk
            )?
            .into_iter()
            .rev() // We need it in little endian bit order
            .map(|e| boolean::Boolean::from(e)).collect::<Vec<_>>();

            // NB: We don't ensure that the bit representation of rsk
            // is "in the field" (Fs) because it's not used except to
            // demonstrate the prover knows it. If they know a
            // congruency then that's equivalent.

            rk = ecc::fixed_base_multiplication(
                cs.namespace(|| "computation of rk"),
                FixedGenerators::ProvingPublicKey,
                &rsk,
                self.params
            )?;
        }

        // Prover witnesses ak (ensures that it's on the curve)
        let ak = ecc::EdwardsPoint::witness(
            cs.namespace(|| "ak"),
            self.ak,
            self.params
        )?;

        // Unpack ak and rk for input to BLAKE2s
        let mut vk = vec![];
        {
            let mut ak_x = ak.x.into_bits_strict(
                cs.namespace(|| "unpack ak.x")
            )?;
            let mut ak_y = ak.y.into_bits_strict(
                cs.namespace(|| "unpack ak.y")
            )?;

            // We want the representation in little endian bit order
            ak_x.reverse();
            ak_y.reverse();

            vk.extend(ak_y);
            vk.push(ak_x[0].clone());
        }
        let mut rho_preimage = vec![];
        {
            let mut rk_x = rk.x.into_bits_strict(
                cs.namespace(|| "unpack rk.x")
            )?;
            let mut rk_y = rk.y.into_bits_strict(
                cs.namespace(|| "unpack rk.y")
            )?;

            // We want the representation in little endian bit order
            rk_x.reverse();
            rk_y.reverse();

            vk.extend(rk_y.iter().cloned());
            vk.push(rk_x[0].clone());

            rho_preimage.extend(rk_y.iter().cloned());
            rho_preimage.push(rk_x[0].clone());
        }

        assert_eq!(vk.len(), 512);

        // Compute the incoming viewing key
        let mut ivk = blake2s::blake2s(
            cs.namespace(|| "computation of ivk"),
            &vk
        )?;

        // Little endian bit order
        ivk.reverse();
        ivk.truncate(251); // drop_5

        // Witness g_d
        let g_d = ecc::EdwardsPoint::witness(
            cs.namespace(|| "witness g_d"),
            self.g_d,
            self.params
        )?;

        // Compute pk_d
        let pk_d = g_d.mul(
            cs.namespace(|| "compute pk_d"),
            &ivk,
            self.params
        )?;

        // Compute note contents
        let mut note_contents = vec![];
        note_contents.extend(value_bits);
        {
            // Unpack g_d for inclusion in the note.
            let mut g_d_x = g_d.x.into_bits_strict(
                cs.namespace(|| "unpack g_d.x")
            )?;
            let mut g_d_y = g_d.y.into_bits_strict(
                cs.namespace(|| "unpack g_d.y")
            )?;

            // We want the representation in little endian bit order
            g_d_x.reverse();
            g_d_y.reverse();

            note_contents.extend(g_d_y);
            note_contents.push(g_d_x[0].clone());
        }
        {
            // Unpack g_d for inclusion in the note.
            let mut pk_d_x = pk_d.x.into_bits_strict(
                cs.namespace(|| "unpack pk_d.x")
            )?;
            let mut pk_d_y = pk_d.y.into_bits_strict(
                cs.namespace(|| "unpack pk_d.y")
            )?;

            // We want the representation in little endian bit order
            pk_d_x.reverse();
            pk_d_y.reverse();

            note_contents.extend(pk_d_y);
            note_contents.push(pk_d_x[0].clone());
        }

        assert_eq!(
            note_contents.len(),
            64 + // value
            256 + // g_d
            256 // p_d
        );

        // Compute the hash of the note contents
        let mut cm = pedersen_hash::pedersen_hash(
            cs.namespace(|| "note content hash"),
            pedersen_hash::Personalization::NoteCommitment,
            &note_contents,
            self.params
        )?;

        {
            // Booleanize the randomness
            let cmr = boolean::field_into_allocated_bits_be(
                cs.namespace(|| "cmr"),
                self.commitment_randomness
            )?
            .into_iter()
            .rev() // We need it in little endian bit order
            .map(|e| boolean::Boolean::from(e))
            .collect::<Vec<_>>();

            let cmr = ecc::fixed_base_multiplication(
                cs.namespace(|| "computation of commitment randomness"),
                FixedGenerators::NoteCommitmentRandomness,
                &cmr,
                self.params
            )?;

            cm = cm.add(
                cs.namespace(|| "randomization of note commitment"),
                &cmr,
                self.params
            )?;
        }

        assert_eq!(self.auth_path.len(), MERKLE_TREE_DEPTH);

        let mut position_bits = vec![];

        // Injective encoding.
        let mut cur = cm.x.clone();

        for (i, e) in self.auth_path.into_iter().enumerate() {
            let cs = &mut cs.namespace(|| format!("merkle tree hash {}", i));

            let cur_is_right = boolean::Boolean::from(boolean::AllocatedBit::alloc(
                cs.namespace(|| "position bit"),
                e.map(|e| e.1)
            )?);

            position_bits.push(cur_is_right.clone());

            let path_element = num::AllocatedNum::alloc(
                cs.namespace(|| "path element"),
                || {
                    Ok(e.get()?.0)
                }
            )?;

            let (xl, xr) = num::AllocatedNum::conditionally_reverse(
                cs.namespace(|| "conditional reversal of preimage"),
                &cur,
                &path_element,
                &cur_is_right
            )?;

            // We don't need to be strict, because the function is
            // collision-resistant.
            let mut preimage = vec![];
            preimage.extend(xl.into_bits(cs.namespace(|| "xl into bits"))?);
            preimage.extend(xr.into_bits(cs.namespace(|| "xr into bits"))?);

            cur = pedersen_hash::pedersen_hash(
                cs.namespace(|| "computation of pedersen hash"),
                pedersen_hash::Personalization::MerkleTree(MERKLE_TREE_DEPTH - i),
                &preimage,
                self.params
            )?.x; // Injective encoding
        }

        assert_eq!(position_bits.len(), MERKLE_TREE_DEPTH);

        // TODO: cur is now the root of the tree, expose it as public input

        let tmp = ecc::fixed_base_multiplication(
            cs.namespace(|| "g^position"),
            FixedGenerators::NullifierPosition,
            &position_bits,
            self.params
        )?;

        cm = cm.add(
            cs.namespace(|| "faerie gold prevention"),
            &tmp,
            self.params
        )?;

        // Let's compute rho = BLAKE2s(rk || cm + position)
        {
            // Unpack g_d for inclusion in the note.
            let mut cm_x = cm.x.into_bits_strict(
                cs.namespace(|| "unpack (cm + position).x")
            )?;
            let mut cm_y = cm.y.into_bits_strict(
                cs.namespace(|| "unpack (cm + position).y")
            )?;

            // We want the representation in little endian bit order
            cm_x.reverse();
            cm_y.reverse();

            rho_preimage.extend(cm_y);
            rho_preimage.push(cm_x[0].clone());
        }
        
        let mut rho = blake2s::blake2s(
            cs.namespace(|| "rho computation"),
            &rho_preimage
        )?;

        // Little endian bit order
        rho.reverse();
        rho.truncate(251); // drop_5

        // Compute nullifier
        let nf = ak.mul(
            cs.namespace(|| "computation of nf"),
            &rho,
            self.params
        )?;

        // TODO: expose nf as public input

        Ok(())
    }
}

#[test]
fn test_input_circuit_with_bls12_381() {
    use pairing::bls12_381::*;
    use rand::{SeedableRng, Rng, XorShiftRng};
    use ::circuit::test::*;
    use jubjub::{JubjubBls12, fs};

    let params = &JubjubBls12::new();
    let rng = &mut XorShiftRng::from_seed([0x3dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);

    let value: u64 = 1;
    let value_randomness: fs::Fs = rng.gen();
    let ak: edwards::Point<Bls12, Unknown> = edwards::Point::rand(rng, params);
    let g_d: edwards::Point<Bls12, Unknown> = edwards::Point::rand(rng, params);
    let p_d: edwards::Point<Bls12, Unknown> = edwards::Point::rand(rng, params);
    let commitment_randomness: fs::Fs = rng.gen();
    let esk: fs::Fs = rng.gen();
    let rsk: fs::Fs = rng.gen();
    let auth_path = vec![Some((rng.gen(), false)); 29];

    {
        let mut cs = TestConstraintSystem::<Bls12>::new();

        let instance = Spend {
            params: params,
            value: Some(1),
            value_randomness: Some(value_randomness),
            rsk: Some(rsk),
            ak: Some(ak),
            g_d: Some(g_d),
            commitment_randomness: Some(commitment_randomness),
            auth_path: auth_path
        };

        instance.synthesize(&mut cs).unwrap();

        assert!(cs.is_satisfied());

        assert_eq!(cs.num_constraints(), 99816);
    }

    // use bellman::groth16::*;

    // let groth_params = generate_random_parameters::<Bls12, _, _>(Spend {
    //     params: params,
    //     value: None,
    //     value_randomness: None,
    //     rsk: None,
    //     ak: None,
    //     g_d: None,
    //     commitment_randomness: None,
    //     auth_path: vec![None; 29]
    // }, rng).unwrap();

    // let pvk = prepare_verifying_key(&groth_params.vk);

    // use std::time::{Duration, Instant};

    // // Let's benchmark stuff!
    // const SAMPLES: u32 = 50;
    // let mut total_proving = Duration::new(0, 0);

    // for _ in 0..SAMPLES {
    //     let start = Instant::now();
    //     {
    //         let c = Spend {
    //             params: params,
    //             value: Some(1),
    //             value_randomness: Some(value_randomness.clone()),
    //             rsk: Some(rsk.clone()),
    //             ak: Some(ak.clone()),
    //             g_d: Some(g_d.clone()),
    //             commitment_randomness: Some(commitment_randomness.clone()),
    //             auth_path: auth_path.clone()
    //         };

    //         create_random_proof(c, &groth_params, rng).unwrap();
    //     }
    //     total_proving += start.elapsed();
    // }

    // let proving_avg = total_proving / SAMPLES;
    // let proving_avg = proving_avg.subsec_nanos() as f64 / 1_000_000_000f64
    //                   + (proving_avg.as_secs() as f64);

    // panic!("Average proving time: {:?} seconds", proving_avg);
}

/// This is an output circuit instance.
pub struct Output<'a, E: JubjubEngine> {
    pub params: &'a E::Params,
    /// Value of the note being created
    pub value: Option<u64>,
    /// Randomness that will hide the value
    pub value_randomness: Option<E::Fs>,
    /// The diversified base, computed by GH(d)
    pub g_d: Option<edwards::Point<E, Unknown>>,
    /// The diversified address point, computed by GH(d)^ivk
    pub p_d: Option<edwards::Point<E, Unknown>>,
    /// The randomness used to hide the note commitment data
    pub commitment_randomness: Option<E::Fs>,
    /// The ephemeral secret key for DH with recipient
    pub esk: Option<E::Fs>
}

impl<'a, E: JubjubEngine> Circuit<E> for Output<'a, E> {
    fn synthesize<CS: ConstraintSystem<E>>(self, cs: &mut CS) -> Result<(), SynthesisError>
    {
        // Booleanize the value into little-endian bit order
        let value_bits = boolean::u64_into_allocated_bits_be(
            cs.namespace(|| "value"),
            self.value
        )?
        .into_iter()
        .rev() // Little endian bit order
        .map(|e| boolean::Boolean::from(e))
        .collect::<Vec<_>>();

        {
            let gv = ecc::fixed_base_multiplication(
                cs.namespace(|| "compute the value in the exponent"),
                FixedGenerators::ValueCommitmentValue,
                &value_bits,
                self.params
            )?;
        
            // Booleanize the randomness
            let hr = boolean::field_into_allocated_bits_be(
                cs.namespace(|| "hr"),
                self.value_randomness
            )?
            .into_iter()
            .rev() // Little endian bit order
            .map(|e| boolean::Boolean::from(e))
            .collect::<Vec<_>>();

            let hr = ecc::fixed_base_multiplication(
                cs.namespace(|| "computation of randomization for value commitment"),
                FixedGenerators::ValueCommitmentRandomness,
                &hr,
                self.params
            )?;

            let gvhr = gv.add(
                cs.namespace(|| "computation of value commitment"),
                &hr,
                self.params
            )?;

            // Expose the value commitment publicly
            let value_commitment_x = cs.alloc_input(
                || "value commitment x",
                || {
                    Ok(*gvhr.x.get_value().get()?)
                }
            )?;

            cs.enforce(
                || "value commitment x equals input",
                |lc| lc + value_commitment_x,
                |lc| lc + CS::one(),
                |lc| lc + gvhr.x.get_variable()
            );

            let value_commitment_y = cs.alloc_input(
                || "value commitment y",
                || {
                    Ok(*gvhr.y.get_value().get()?)
                }
            )?;

            cs.enforce(
                || "value commitment y equals input",
                |lc| lc + value_commitment_y,
                |lc| lc + CS::one(),
                |lc| lc + gvhr.y.get_variable()
            );
        }

        // Let's start to construct our note
        let mut note_contents = vec![];
        note_contents.extend(value_bits);

        // Let's deal with g_d
        {
            let g_d = ecc::EdwardsPoint::witness(
                cs.namespace(|| "witness g_d"),
                self.g_d,
                self.params
            )?;

            // Check that g_d is not of small order
            {
                let g_d = g_d.double(
                    cs.namespace(|| "first doubling of g_d"),
                    self.params
                )?;
                let g_d = g_d.double(
                    cs.namespace(|| "second doubling of g_d"),
                    self.params
                )?;
                let g_d = g_d.double(
                    cs.namespace(|| "third doubling of g_d"),
                    self.params
                )?;

                // (0, -1) is a small order point, but won't ever appear here
                // because cofactor is 2^3, and we performed three doublings.
                // (0, 1) is the neutral element, so checking if x is nonzero
                // is sufficient to prevent small order points here.
                g_d.x.assert_nonzero(cs.namespace(|| "check not inf"))?;
            }

            // Unpack g_d for inclusion in the note.
            let mut g_d_x = g_d.x.into_bits_strict(
                cs.namespace(|| "unpack g_d.x")
            )?;
            let mut g_d_y = g_d.y.into_bits_strict(
                cs.namespace(|| "unpack g_d.y")
            )?;

            // We want the representation in little endian bit order
            g_d_x.reverse();
            g_d_y.reverse();

            note_contents.extend(g_d_y);
            note_contents.push(g_d_x[0].clone());

            // Compute epk from esk
            let esk = boolean::field_into_allocated_bits_be(
                cs.namespace(|| "esk"),
                self.esk
            )?
            .into_iter()
            .rev() // We need it in little endian bit order
            .map(|e| boolean::Boolean::from(e))
            .collect::<Vec<_>>();

            let epk = g_d.mul(
                cs.namespace(|| "epk computation"),
                &esk,
                self.params
            )?;

            // Expose epk publicly
            let epk_x = cs.alloc_input(
                || "epk x",
                || {
                    Ok(*epk.x.get_value().get()?)
                }
            )?;

            cs.enforce(
                || "epk x equals input",
                |lc| lc + epk_x,
                |lc| lc + CS::one(),
                |lc| lc + epk.x.get_variable()
            );

            let epk_y = cs.alloc_input(
                || "epk y",
                || {
                    Ok(*epk.y.get_value().get()?)
                }
            )?;

            cs.enforce(
                || "epk y equals input",
                |lc| lc + epk_y,
                |lc| lc + CS::one(),
                |lc| lc + epk.y.get_variable()
            );
        }

        // Now let's deal with p_d. We don't do any checks and
        // essentially allow the prover to witness any 256 bits
        // they would like.
        {
            let p_d = self.p_d.map(|e| e.into_xy());

            let y_contents = boolean::field_into_allocated_bits_be(
                cs.namespace(|| "p_d bits of y"),
                p_d.map(|e| e.1)
            )?
            .into_iter()
            .rev() // We need it in little endian bit order
            .map(|e| boolean::Boolean::from(e))
            .collect::<Vec<_>>();

            let sign_bit = boolean::Boolean::from(boolean::AllocatedBit::alloc(
                cs.namespace(|| "p_d bit of x"),
                p_d.map(|e| e.0.into_repr().is_odd())
            )?);

            note_contents.extend(y_contents);
            note_contents.push(sign_bit);
        }

        assert_eq!(
            note_contents.len(),
            64 + // value
            256 + // g_d
            256 // p_d
        );

        // Compute the hash of the note contents
        let mut cm = pedersen_hash::pedersen_hash(
            cs.namespace(|| "note content hash"),
            pedersen_hash::Personalization::NoteCommitment,
            &note_contents,
            self.params
        )?;

        {
            // Booleanize the randomness
            let cmr = boolean::field_into_allocated_bits_be(
                cs.namespace(|| "cmr"),
                self.commitment_randomness
            )?
            .into_iter()
            .rev() // We need it in little endian bit order
            .map(|e| boolean::Boolean::from(e))
            .collect::<Vec<_>>();

            let cmr = ecc::fixed_base_multiplication(
                cs.namespace(|| "computation of commitment randomness"),
                FixedGenerators::NoteCommitmentRandomness,
                &cmr,
                self.params
            )?;

            cm = cm.add(
                cs.namespace(|| "randomization of note commitment"),
                &cmr,
                self.params
            )?;
        }

        // Only the x-coordinate of the output is revealed,
        // since we know it is prime order, and we know that
        // the x-coordinate is an injective encoding for
        // prime-order elements.
        let commitment_input = cs.alloc_input(
            || "commitment input",
            || {
                Ok(*cm.x.get_value().get()?)
            }
        )?;

        cs.enforce(
            || "commitment input correct",
            |lc| lc + commitment_input,
            |lc| lc + CS::one(),
            |lc| lc + cm.x.get_variable()
        );

        Ok(())
    }
}

#[test]
fn test_output_circuit_with_bls12_381() {
    use pairing::bls12_381::*;
    use rand::{SeedableRng, Rng, XorShiftRng};
    use ::circuit::test::*;
    use jubjub::{JubjubBls12, fs};

    let params = &JubjubBls12::new();
    let rng = &mut XorShiftRng::from_seed([0x3dbe6259, 0x8d313d76, 0x3237db17, 0xe5bc0654]);

    let value: u64 = 1;
    let value_randomness: fs::Fs = rng.gen();
    let g_d: edwards::Point<Bls12, Unknown> = edwards::Point::rand(rng, params);
    let p_d: edwards::Point<Bls12, Unknown> = edwards::Point::rand(rng, params);
    let commitment_randomness: fs::Fs = rng.gen();
    let esk: fs::Fs = rng.gen();

    {
        let mut cs = TestConstraintSystem::<Bls12>::new();

        let instance = Output {
            params: params,
            value: Some(1),
            value_randomness: Some(value_randomness),
            g_d: Some(g_d.clone()),
            p_d: Some(p_d.clone()),
            commitment_randomness: Some(commitment_randomness),
            esk: Some(esk.clone())
        };

        instance.synthesize(&mut cs).unwrap();

        assert!(cs.is_satisfied());

        assert_eq!(cs.num_constraints(), 8315);
    }

    // use bellman::groth16::*;

    // let groth_params = generate_random_parameters::<Bls12, _, _>(Output {
    //     params: params,
    //     value: None,
    //     value_randomness: None,
    //     g_d: None,
    //     p_d: None,
    //     commitment_randomness: None,
    //     esk: None
    // }, rng).unwrap();

    // let pvk = prepare_verifying_key(&groth_params.vk);

    // use std::time::{Duration, Instant};

    // // Let's benchmark stuff!
    // const SAMPLES: u32 = 50;
    // let mut total_proving = Duration::new(0, 0);

    // for _ in 0..SAMPLES {
    //     let start = Instant::now();
    //     {
    //         let c = Output {
    //             params: params,
    //             value: Some(1),
    //             value_randomness: Some(value_randomness),
    //             g_d: Some(g_d.clone()),
    //             p_d: Some(p_d.clone()),
    //             commitment_randomness: Some(commitment_randomness),
    //             esk: Some(esk.clone())
    //         };

    //         create_random_proof(c, &groth_params, rng).unwrap();
    //     }
    //     total_proving += start.elapsed();
    // }

    // let proving_avg = total_proving / SAMPLES;
    // let proving_avg = proving_avg.subsec_nanos() as f64 / 1_000_000_000f64
    //                   + (proving_avg.as_secs() as f64);

    // panic!("Average proving time: {:?} seconds", proving_avg);
}