use pairing::*;
use pairing::wnaf::*;
use ::{
    Input,
    Error,
    LinearCombination,
    Index,
    Circuit,
    Variable,
    ConstraintSystem,
    PublicConstraintSystem
};
use super::{VerifyingKey, Parameters};
use domain::{Scalar, EvaluationDomain};
use rand::Rng;
use multicore;
use std::sync::Arc;

pub fn generate_random_parameters<E, C, R>(
    circuit: C,
    rng: &mut R
) -> Result<Parameters<E>, Error>
    where E: Engine, C: Circuit<E>, R: Rng
{
    let g1 = rng.gen();
    let g2 = rng.gen();
    let alpha = rng.gen();
    let beta = rng.gen();
    let gamma = rng.gen();
    let delta = rng.gen();
    let tau = rng.gen();

    generate_parameters::<E, C>(
        circuit,
        g1,
        g2,
        alpha,
        beta,
        gamma,
        delta,
        tau
    )
}

/// Create parameters for a circuit, given some trapdoors.
pub fn generate_parameters<E, C>(
    circuit: C,
    g1: E::G1,
    g2: E::G2,
    alpha: E::Fr,
    beta: E::Fr,
    gamma: E::Fr,
    delta: E::Fr,
    tau: E::Fr
) -> Result<Parameters<E>, Error>
    where E: Engine, C: Circuit<E>
{
    // This is our assembly structure that we'll use to synthesize the
    // circuit into a QAP.
    struct KeypairAssembly<E: Engine> {
        num_inputs: usize,
        num_aux: usize,
        num_constraints: usize,
        at_inputs: Vec<Vec<(E::Fr, usize)>>,
        bt_inputs: Vec<Vec<(E::Fr, usize)>>,
        ct_inputs: Vec<Vec<(E::Fr, usize)>>,
        at_aux: Vec<Vec<(E::Fr, usize)>>,
        bt_aux: Vec<Vec<(E::Fr, usize)>>,
        ct_aux: Vec<Vec<(E::Fr, usize)>>
    }

    impl<E: Engine> PublicConstraintSystem<E> for KeypairAssembly<E> {
        fn alloc_input<F: FnOnce() -> Result<E::Fr, Error>>(&mut self, f: F) -> Result<Variable, Error> {
            // In this context, we don't have an assignment.
            let _ = f();

            let index = self.num_inputs;
            self.num_inputs += 1;

            self.at_inputs.push(vec![]);
            self.bt_inputs.push(vec![]);
            self.ct_inputs.push(vec![]);

            Ok(Variable(Index::Input(index)))
        }
    }

    impl<E: Engine> ConstraintSystem<E> for KeypairAssembly<E> {
        fn alloc<F: FnOnce() -> Result<E::Fr, Error>>(&mut self, f: F) -> Result<Variable, Error> {
            // In this context, we don't have an assignment.
            let _ = f();

            let index = self.num_aux;
            self.num_aux += 1;

            self.at_aux.push(vec![]);
            self.bt_aux.push(vec![]);
            self.ct_aux.push(vec![]);

            Ok(Variable(Index::Aux(index)))
        }

        fn enforce(
            &mut self,
            a: LinearCombination<E>,
            b: LinearCombination<E>,
            c: LinearCombination<E>
        )
        {
            fn qap_eval<E: Engine>(
                l: LinearCombination<E>,
                inputs: &mut [Vec<(E::Fr, usize)>],
                aux: &mut [Vec<(E::Fr, usize)>],
                this_constraint: usize
            )
            {
                for (index, coeff) in l.0 {
                    match index {
                        Index::Input(id) => inputs[id].push((coeff, this_constraint)),
                        Index::Aux(id) => aux[id].push((coeff, this_constraint))
                    }
                }
            }

            qap_eval(a, &mut self.at_inputs, &mut self.at_aux, self.num_constraints);
            qap_eval(b, &mut self.bt_inputs, &mut self.bt_aux, self.num_constraints);
            qap_eval(c, &mut self.ct_inputs, &mut self.ct_aux, self.num_constraints);

            self.num_constraints += 1;
        }
    }

    let mut assembly = KeypairAssembly {
        num_inputs: 0,
        num_aux: 0,
        num_constraints: 0,
        at_inputs: vec![],
        bt_inputs: vec![],
        ct_inputs: vec![],
        at_aux: vec![],
        bt_aux: vec![],
        ct_aux: vec![]
    };

    // Allocate the "one" input variable
    assembly.alloc_input(|| Ok(E::Fr::one()))?;

    // Synthesize the circuit.
    circuit.synthesize(&mut assembly)?.synthesize(&mut assembly)?;

    // Input consistency constraints: x * 0 = 0
    for i in 0..assembly.num_inputs {
        assembly.enforce(LinearCombination::zero() + Variable(Index::Input(i)),
                         LinearCombination::zero(),
                         LinearCombination::zero());
    }

    // Ensure that all auxillary variables are constrained
    for i in 0..assembly.num_aux {
        if assembly.at_aux[i].len() == 0 &&
           assembly.bt_aux[i].len() == 0 &&
           assembly.ct_aux[i].len() == 0
        {
            return Err(Error::UnconstrainedVariable(Variable(Index::Aux(i))));
        }
    }

    // Create evaluation domain for the QAP
    let powers_of_tau = vec![Scalar::<E>(E::Fr::zero()); assembly.num_constraints];
    let mut powers_of_tau = EvaluationDomain::from_coeffs(powers_of_tau)?;

    // Compute G1 window table
    let mut g1_table = vec![];
    let g1_table_size = E::G1::recommended_wnaf_for_num_scalars(
        // H query
        (powers_of_tau.as_ref().len() - 1)
        // IC/L queries
        + assembly.num_inputs + assembly.num_aux
        // A query
        + assembly.num_inputs + assembly.num_aux
        // B query
        + assembly.num_inputs + assembly.num_aux
    );
    wnaf_table(&mut g1_table, g1, g1_table_size);

    // Compute G2 window table
    let mut g2_table = vec![];
    let g2_table_size = E::G2::recommended_wnaf_for_num_scalars(
        // B query
        assembly.num_inputs + assembly.num_aux
    );
    wnaf_table(&mut g2_table, g2, g2_table_size);

    let gamma_inverse = gamma.inverse().ok_or(Error::UnexpectedIdentity)?;
    let delta_inverse = delta.inverse().ok_or(Error::UnexpectedIdentity)?;

    // Compute the H query
    let mut h = vec![E::G1::zero(); powers_of_tau.as_ref().len() - 1];
    {
        // Compute the powers of tau
        {
            let powers_of_tau = powers_of_tau.as_mut();
            multicore::scope(powers_of_tau.len(), |scope, chunk| {
                for (i, powers_of_tau) in powers_of_tau.chunks_mut(chunk).enumerate()
                {
                    scope.spawn(move || {
                        let mut current_tau_power = tau.pow(&[(i*chunk) as u64]);

                        for p in powers_of_tau {
                            p.0 = current_tau_power;
                            current_tau_power.mul_assign(&tau);
                        }
                    });
                }
            });
        }

        // coeff = t(x) / delta
        let mut coeff = powers_of_tau.z(&tau);
        coeff.mul_assign(&delta_inverse);

        // Compute the H query with multiple threads
        multicore::scope(h.len(), |scope, chunk| {
            for (h, p) in h.chunks_mut(chunk).zip(powers_of_tau.as_ref().chunks(chunk))
            {
                let g1_table = &g1_table;

                scope.spawn(move || {
                    // Create wNAF form storage location for this thread
                    let mut wnaf = vec![];

                    // Set values of the H query to g1^{(tau^i * t(tau)) / delta}
                    for (h, p) in h.iter_mut().zip(p.iter())
                    {
                        // Compute final exponent
                        let mut exp = p.0;
                        exp.mul_assign(&coeff);

                        // Compute wNAF form of exponent
                        wnaf_form(&mut wnaf, exp.into_repr(), g1_table_size);

                        // Exponentiate
                        *h = wnaf_exp(g1_table, &wnaf);
                    }

                    // Batch normalize
                    E::G1::batch_normalization(h);
                });
            }
        });
    }

    // Use inverse FFT to convert powers of tau to Lagrange coefficients
    powers_of_tau.ifft();
    let powers_of_tau = powers_of_tau.into_coeffs();

    let mut a = vec![E::G1::zero(); assembly.num_inputs + assembly.num_aux];
    let mut b_g1 = vec![E::G1::zero(); assembly.num_inputs + assembly.num_aux];
    let mut b_g2 = vec![E::G2::zero(); assembly.num_inputs + assembly.num_aux];
    let mut ic = vec![E::G1::zero(); assembly.num_inputs];
    let mut l = vec![E::G1::zero(); assembly.num_aux];

    fn eval<E: Engine>(
        // wNAF window tables
        g1_table: &[E::G1],
        g1_table_size: usize,
        g2_table: &[E::G2],
        g2_table_size: usize,

        // Lagrange coefficients for tau
        powers_of_tau: &[Scalar<E>],

        // QAP polynomials
        at: &[Vec<(E::Fr, usize)>],
        bt: &[Vec<(E::Fr, usize)>],
        ct: &[Vec<(E::Fr, usize)>],

        // Resulting evaluated QAP polynomials
        a: &mut [E::G1],
        b_g1: &mut [E::G1],
        b_g2: &mut [E::G2],
        ext: &mut [E::G1],

        // Inverse coefficient for ext elements
        inv: &E::Fr,

        // Trapdoors
        alpha: &E::Fr,
        beta: &E::Fr
    )
    {
        // Sanity check
        assert_eq!(a.len(), at.len());
        assert_eq!(a.len(), bt.len());
        assert_eq!(a.len(), ct.len());
        assert_eq!(a.len(), b_g1.len());
        assert_eq!(a.len(), b_g2.len());
        assert_eq!(a.len(), ext.len());

        // Evaluate polynomials in multiple threads
        multicore::scope(a.len(), |scope, chunk| {
            for ((((((a, b_g1), b_g2), ext), at), bt), ct) in a.chunks_mut(chunk)
                                                               .zip(b_g1.chunks_mut(chunk))
                                                               .zip(b_g2.chunks_mut(chunk))
                                                               .zip(ext.chunks_mut(chunk))
                                                               .zip(at.chunks(chunk))
                                                               .zip(bt.chunks(chunk))
                                                               .zip(ct.chunks(chunk))
            {
                scope.spawn(move || {
                    // Create wNAF form storage location for this thread
                    let mut wnaf = vec![];

                    for ((((((a, b_g1), b_g2), ext), at), bt), ct) in a.iter_mut()
                                                                       .zip(b_g1.iter_mut())
                                                                       .zip(b_g2.iter_mut())
                                                                       .zip(ext.iter_mut())
                                                                       .zip(at.iter())
                                                                       .zip(bt.iter())
                                                                       .zip(ct.iter())
                    {
                        fn eval_at_tau<E: Engine>(
                            powers_of_tau: &[Scalar<E>],
                            p: &[(E::Fr, usize)]
                        ) -> E::Fr
                        {
                            let mut acc = E::Fr::zero();

                            for &(ref coeff, index) in p {
                                let mut n = powers_of_tau[index].0;
                                n.mul_assign(coeff);
                                acc.add_assign(&n);
                            }

                            acc
                        }

                        // Evaluate QAP polynomials at tau
                        let mut at = eval_at_tau(powers_of_tau, at);
                        let mut bt = eval_at_tau(powers_of_tau, bt);
                        let ct = eval_at_tau(powers_of_tau, ct);

                        // Compute A query (in G1)
                        if !at.is_zero() {
                            wnaf_form(&mut wnaf, at.into_repr(), g1_table_size);
                            *a = wnaf_exp(&g1_table, &wnaf);
                        }

                        // Compute B query (in G1/G2)
                        if !bt.is_zero() {
                            // Normalize the field element once
                            let bt_repr = bt.into_repr();
                            wnaf_form(&mut wnaf, bt_repr, g1_table_size);
                            *b_g1 = wnaf_exp(&g1_table, &wnaf);

                            // G1 window table might use the same window size
                            // as the G2 window table, so we wouldn't need to
                            // recompute the wNAF form of the exponent.
                            if g1_table_size != g2_table_size {
                                wnaf_form(&mut wnaf, bt_repr, g2_table_size);
                            }
                            *b_g2 = wnaf_exp(&g2_table, &wnaf);
                        }

                        at.mul_assign(&beta);
                        bt.mul_assign(&alpha);

                        let mut e = at;
                        e.add_assign(&bt);
                        e.add_assign(&ct);
                        e.mul_assign(inv);

                        wnaf_form(&mut wnaf, e.into_repr(), g1_table_size);
                        *ext = wnaf_exp(&g1_table, &wnaf);
                    }

                    // Batch normalize
                    E::G1::batch_normalization(a);
                    E::G1::batch_normalization(b_g1);
                    E::G2::batch_normalization(b_g2);
                    E::G1::batch_normalization(ext);
                });
            }
        });
    }

    // Evaluate for inputs.
    eval(
        &g1_table,
        g1_table_size,
        &g2_table,
        g2_table_size,
        &powers_of_tau,
        &assembly.at_inputs,
        &assembly.bt_inputs,
        &assembly.ct_inputs,
        &mut a[0..assembly.num_inputs],
        &mut b_g1[0..assembly.num_inputs],
        &mut b_g2[0..assembly.num_inputs],
        &mut ic,
        &gamma_inverse,
        &alpha,
        &beta
    );

    // Evaluate for auxillary variables.
    eval(
        &g1_table,
        g1_table_size,
        &g2_table,
        g2_table_size,
        &powers_of_tau,
        &assembly.at_aux,
        &assembly.bt_aux,
        &assembly.ct_aux,
        &mut a[assembly.num_inputs..],
        &mut b_g1[assembly.num_inputs..],
        &mut b_g2[assembly.num_inputs..],
        &mut l,
        &delta_inverse,
        &alpha,
        &beta
    );

    let g1 = g1.into_affine();
    let g2 = g2.into_affine();

    let vk = VerifyingKey::<E> {
        alpha_g1: g1.mul(alpha).into_affine(),
        beta_g1: g1.mul(beta).into_affine(),
        beta_g2: g2.mul(beta).into_affine(),
        gamma_g2: g2.mul(gamma).into_affine(),
        delta_g1: g1.mul(delta).into_affine(),
        delta_g2: g2.mul(delta).into_affine(),
        ic: ic.into_iter().map(|e| e.into_affine()).collect()
    };

    Ok(Parameters {
        vk: vk,
        h: Arc::new(h.into_iter().map(|e| e.into_affine()).collect()),
        l: Arc::new(l.into_iter().map(|e| e.into_affine()).collect()),

        // Filter points at infinity away from A/B queries
        a: Arc::new(a.into_iter().filter(|e| !e.is_zero()).map(|e| e.into_affine()).collect()),
        b_g1: Arc::new(b_g1.into_iter().filter(|e| !e.is_zero()).map(|e| e.into_affine()).collect()),
        b_g2: Arc::new(b_g2.into_iter().filter(|e| !e.is_zero()).map(|e| e.into_affine()).collect())
    })
}