//! Methods to preprocess the constraint system for use in a proof

use crate::commitment_scheme::kzg10::CommitKey;
use crate::constraint_system::cs_errors::PreProcessingError;
use crate::constraint_system::StandardComposer;

use crate::fft::{EvaluationDomain, Evaluations, Polynomial};
use crate::proof_system::widget;
use dusk_bls12_381::Scalar;
use failure::Error;
use merlin::Transcript;

/// Struct that contains all of the selector and permutation polynomials in PLONK
/// These polynomials are in coefficient form
pub(crate) struct SelectorPolynomials {
    q_m: Polynomial,
    q_l: Polynomial,
    q_r: Polynomial,
    q_o: Polynomial,
    q_c: Polynomial,
    q_4: Polynomial,
    q_arith: Polynomial,
    q_range: Polynomial,
    q_logic: Polynomial,
    q_ecc: Polynomial,
    left_sigma: Polynomial,
    right_sigma: Polynomial,
    out_sigma: Polynomial,
    fourth_sigma: Polynomial,
}

impl StandardComposer {
    /// Pads the circuit to the next power of two
    /// `diff` is the difference between circuit size and next power of two.
    fn pad(&mut self, diff: usize) {
        // Add a zero variable to circuit
        let zero_scalar = Scalar::zero();
        let zero_var = self.add_input(zero_scalar);

        let zeroes_scalar = vec![zero_scalar; diff];
        let zeroes_var = vec![zero_var; diff];

        self.q_m.extend(zeroes_scalar.iter());
        self.q_l.extend(zeroes_scalar.iter());
        self.q_r.extend(zeroes_scalar.iter());
        self.q_o.extend(zeroes_scalar.iter());
        self.q_c.extend(zeroes_scalar.iter());
        self.q_4.extend(zeroes_scalar.iter());
        self.q_arith.extend(zeroes_scalar.iter());
        self.q_range.extend(zeroes_scalar.iter());
        self.q_logic.extend(zeroes_scalar.iter());
        self.q_ecc.extend(zeroes_scalar.iter());

        self.w_l.extend(zeroes_var.iter());
        self.w_r.extend(zeroes_var.iter());
        self.w_o.extend(zeroes_var.iter());
        self.w_4.extend(zeroes_var.iter());

        self.n += diff;
    }
    /// Checks that all of the wires of the composer have the same
    /// length.
    fn check_poly_same_len(&self) -> Result<(), PreProcessingError> {
        let k = self.q_m.len();

        if self.q_o.len() == k
            && self.q_l.len() == k
            && self.q_r.len() == k
            && self.q_c.len() == k
            && self.q_4.len() == k
            && self.q_arith.len() == k
            && self.q_range.len() == k
            && self.q_logic.len() == k
            && self.q_ecc.len() == k
            && self.w_l.len() == k
            && self.w_r.len() == k
            && self.w_o.len() == k
        {
            Ok(())
        } else {
            Err(PreProcessingError::MissmatchedPolyLen)
        }
    }
    /// These are the parts of preprocessing that the prover must compute
    /// Although the prover does not need the verification key, he must compute the commitments
    /// in order to seed the transcript, allowing both the prover and verifier to have the same view
    pub fn preprocess_prover(
        &mut self,
        commit_key: &CommitKey,
        transcript: &mut Transcript,
    ) -> Result<widget::ProverKey, Error> {
        let (_, selectors, domain) = self.preprocess_shared(commit_key, transcript)?;

        let domain_4n = EvaluationDomain::new(4 * domain.size())?;
        let q_m_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.q_m), domain_4n);
        let q_l_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.q_l), domain_4n);
        let q_r_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.q_r), domain_4n);
        let q_o_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.q_o), domain_4n);
        let q_c_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.q_c), domain_4n);
        let q_4_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.q_4), domain_4n);
        let q_arith_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.q_arith), domain_4n);
        let q_range_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.q_range), domain_4n);
        let q_logic_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.q_logic), domain_4n);
        let q_ecc_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.q_ecc), domain_4n);

        let left_sigma_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.left_sigma), domain_4n);
        let right_sigma_eval_4n = Evaluations::from_vec_and_domain(
            domain_4n.coset_fft(&selectors.right_sigma),
            domain_4n,
        );
        let out_sigma_eval_4n =
            Evaluations::from_vec_and_domain(domain_4n.coset_fft(&selectors.out_sigma), domain_4n);
        let fourth_sigma_eval_4n = Evaluations::from_vec_and_domain(
            domain_4n.coset_fft(&selectors.fourth_sigma),
            domain_4n,
        );
        // XXX: Remove this and compute it on the fly
        let linear_eval_4n = Evaluations::from_vec_and_domain(
            domain_4n.coset_fft(&[Scalar::zero(), Scalar::one()]),
            domain_4n,
        );

        // Prover Key for arithmetic circuits
        let arithmetic_prover_key = widget::arithmetic::ProverKey {
            q_m: (selectors.q_m, q_m_eval_4n),
            q_l: (selectors.q_l.clone(), q_l_eval_4n.clone()),
            q_r: (selectors.q_r.clone(), q_r_eval_4n.clone()),
            q_o: (selectors.q_o, q_o_eval_4n),
            q_c: (selectors.q_c.clone(), q_c_eval_4n.clone()),
            q_4: (selectors.q_4, q_4_eval_4n),
            q_arith: (selectors.q_arith, q_arith_eval_4n),
        };

        // Prover Key for range circuits
        let range_prover_key = widget::range::ProverKey {
            q_range: (selectors.q_range, q_range_eval_4n),
        };

        // Prover Key for logic circuits
        let logic_prover_key = widget::logic::ProverKey {
            q_c: (selectors.q_c.clone(), q_c_eval_4n.clone()),
            q_logic: (selectors.q_logic, q_logic_eval_4n),
        };

        // Prover Key for ecc circuits
        let ecc_prover_key = widget::ecc::ProverKey {
            q_l: (selectors.q_l, q_l_eval_4n),
            q_r: (selectors.q_r, q_r_eval_4n),
            q_c: (selectors.q_c, q_c_eval_4n),
            q_ecc: (selectors.q_ecc, q_ecc_eval_4n),
        };

        // Prover Key for permutation argument
        let permutation_prover_key = widget::permutation::ProverKey {
            left_sigma: (selectors.left_sigma, left_sigma_eval_4n),
            right_sigma: (selectors.right_sigma, right_sigma_eval_4n),
            out_sigma: (selectors.out_sigma, out_sigma_eval_4n),
            fourth_sigma: (selectors.fourth_sigma, fourth_sigma_eval_4n),
            linear_evaluations: linear_eval_4n,
        };

        let prover_key = widget::ProverKey {
            arithmetic: arithmetic_prover_key,
            logic: logic_prover_key,
            range: range_prover_key,
            permutation: permutation_prover_key,
            ecc: ecc_prover_key,
            // Compute 4n evaluations for X^n -1
            v_h_coset_4n: domain_4n.compute_vanishing_poly_over_coset(domain.size() as u64),
        };

        Ok(prover_key)
    }
    /// The verifier only requires the commitments in order to verify a proof
    /// We can therefore speed up preprocessing for the verifier by skipping the FFTs
    /// needed to compute the 4n evaluations
    pub fn preprocess_verifier(
        &mut self,
        commit_key: &CommitKey,
        transcript: &mut Transcript,
    ) -> Result<widget::VerifierKey, Error> {
        let (verifier_key, _, _) = self.preprocess_shared(commit_key, transcript)?;
        Ok(verifier_key)
    }
    // Both the prover and verifier must perform IFFTs on the selector polynomials and permutation polynomials
    // In order to commit to them and have the same transcript view
    fn preprocess_shared(
        &mut self,
        commit_key: &CommitKey,
        transcript: &mut Transcript,
    ) -> Result<(widget::VerifierKey, SelectorPolynomials, EvaluationDomain), Error> {
        let domain = EvaluationDomain::new(self.circuit_size())?;

        // Check that the length of the wires is consistent.
        self.check_poly_same_len()?;

        // 1. Pad circuit to a power of two
        self.pad(domain.size as usize - self.n);

        let q_m_poly = Polynomial::from_coefficients_slice(&domain.ifft(&self.q_m));
        let q_l_poly = Polynomial::from_coefficients_slice(&domain.ifft(&self.q_l));
        let q_r_poly = Polynomial::from_coefficients_slice(&domain.ifft(&self.q_r));
        let q_o_poly = Polynomial::from_coefficients_slice(&domain.ifft(&self.q_o));
        let q_c_poly = Polynomial::from_coefficients_slice(&domain.ifft(&self.q_c));
        let q_4_poly = Polynomial::from_coefficients_slice(&domain.ifft(&self.q_4));
        let q_arith_poly = Polynomial::from_coefficients_slice(&domain.ifft(&self.q_arith));
        let q_range_poly = Polynomial::from_coefficients_slice(&domain.ifft(&self.q_range));
        let q_logic_poly = Polynomial::from_coefficients_slice(&domain.ifft(&self.q_logic));
        let q_ecc_poly = Polynomial::from_coefficients_slice(&domain.ifft(&self.q_ecc));

        // 2. Compute the sigma polynomials
        let (left_sigma_poly, right_sigma_poly, out_sigma_poly, fourth_sigma_poly) =
            self.perm.compute_sigma_polynomials(self.n, &domain);

        let q_m_poly_commit = commit_key.commit(&q_m_poly).unwrap_or_default();
        let q_l_poly_commit = commit_key.commit(&q_l_poly).unwrap_or_default();
        let q_r_poly_commit = commit_key.commit(&q_r_poly).unwrap_or_default();
        let q_o_poly_commit = commit_key.commit(&q_o_poly).unwrap_or_default();
        let q_c_poly_commit = commit_key.commit(&q_c_poly).unwrap_or_default();
        let q_4_poly_commit = commit_key.commit(&q_4_poly).unwrap_or_default();
        let q_arith_poly_commit = commit_key.commit(&q_arith_poly).unwrap_or_default();
        let q_range_poly_commit = commit_key.commit(&q_range_poly).unwrap_or_default();
        let q_logic_poly_commit = commit_key.commit(&q_logic_poly).unwrap_or_default();
        let q_ecc_poly_commit = commit_key.commit(&q_ecc_poly).unwrap_or_default();

        let left_sigma_poly_commit = commit_key.commit(&left_sigma_poly)?;
        let right_sigma_poly_commit = commit_key.commit(&right_sigma_poly)?;
        let out_sigma_poly_commit = commit_key.commit(&out_sigma_poly)?;
        let fourth_sigma_poly_commit = commit_key.commit(&fourth_sigma_poly)?;

        // Verifier Key for arithmetic circuits
        let arithmetic_verifier_key = widget::arithmetic::VerifierKey {
            q_m: q_m_poly_commit,
            q_l: q_l_poly_commit,
            q_r: q_r_poly_commit,
            q_o: q_o_poly_commit,
            q_c: q_c_poly_commit,
            q_4: q_4_poly_commit,
            q_arith: q_arith_poly_commit,
        };
        // Verifier Key for range circuits
        let range_verifier_key = widget::range::VerifierKey {
            q_range: q_range_poly_commit,
        };
        // Verifier Key for logic circuits
        let logic_verifier_key = widget::logic::VerifierKey {
            q_c: q_c_poly_commit,
            q_logic: q_logic_poly_commit,
        };
        // Verifier Key for ecc circuits
        let ecc_verifier_key = widget::ecc::VerifierKey {
            q_l: q_l_poly_commit,
            q_r: q_r_poly_commit,
            q_ecc: q_ecc_poly_commit,
        };
        // Verifier Key for permutation argument
        let permutation_verifier_key = widget::permutation::VerifierKey {
            left_sigma: left_sigma_poly_commit,
            right_sigma: right_sigma_poly_commit,
            out_sigma: out_sigma_poly_commit,
            fourth_sigma: fourth_sigma_poly_commit,
        };

        let verifier_key = widget::VerifierKey {
            n: self.circuit_size(),
            arithmetic: arithmetic_verifier_key,
            logic: logic_verifier_key,
            range: range_verifier_key,
            ecc: ecc_verifier_key,
            permutation: permutation_verifier_key,
        };

        let selectors = SelectorPolynomials {
            q_m: q_m_poly,
            q_l: q_l_poly,
            q_r: q_r_poly,
            q_o: q_o_poly,
            q_c: q_c_poly,
            q_4: q_4_poly,
            q_arith: q_arith_poly,
            q_range: q_range_poly,
            q_logic: q_logic_poly,
            q_ecc: q_ecc_poly,
            left_sigma: left_sigma_poly,
            right_sigma: right_sigma_poly,
            out_sigma: out_sigma_poly,
            fourth_sigma: fourth_sigma_poly,
        };

        // Add the circuit description to the transcript
        verifier_key.seed_transcript(transcript);

        Ok((verifier_key, selectors, domain))
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::constraint_system::helper::*;
    #[test]
    /// Tests that the circuit gets padded to the correct length
    /// XXX: We can do this test without dummy_gadget method
    fn test_pad() {
        let mut composer: StandardComposer = StandardComposer::new();
        dummy_gadget(100, &mut composer);

        // Pad the circuit to next power of two
        let next_pow_2 = composer.n.next_power_of_two() as u64;
        composer.pad(next_pow_2 as usize - composer.n);

        let size = composer.n;
        assert!(size.is_power_of_two());
        assert!(composer.q_m.len() == size);
        assert!(composer.q_l.len() == size);
        assert!(composer.q_o.len() == size);
        assert!(composer.q_r.len() == size);
        assert!(composer.q_c.len() == size);
        assert!(composer.q_arith.len() == size);
        assert!(composer.q_range.len() == size);
        assert!(composer.q_logic.len() == size);
        assert!(composer.w_l.len() == size);
        assert!(composer.w_r.len() == size);
        assert!(composer.w_o.len() == size);
    }
}