/** @file			 prover.hpp
 *****************************************************************************
 * @author		 Janno Siim, Stefanos Chaliasos
 * @copyright  MIT license (see LICENSE file)
 *****************************************************************************/
#ifndef PROVER_H_
#define PROVER_H_

#include <utility> //contains pair
#include "omp.h"
#include "crs.hpp"
#include "types.hpp"
#include "shuffle_util.hpp"
#include "utils.hpp"
#include "nlohmann/json.hpp"


using namespace std;

/*! Data structure for offline part of the proof.
 * Contains proof for the permutation matrix argument and
 * the same-message argument.
 * \brief Offline part of the proof.
 */
template<typename ppT>
struct Offline_proof {
		//! Vector of commitment values a_i for i = 1, ... n - 1
		libff::G1_vector<ppT> a_coms;

		//! Vector of commitment values \hat{a_i} for i = 1, ..., n - 1
		libff::G1_vector<ppT> a_hat_coms;

		//! Vector of commitment values b_i for i = 1, ... n - 1
		libff::G2_vector<ppT> b_coms;

		//! Vector of unit vector proofs (proof for permutation matrix argument)
		libff::G1_vector<ppT> uvs;

		//! Vector of proofs for the same-message argument
		libff::G1_vector<ppT> same_msgs;

		//! Commitment to randomizers with polynomials P^{hat}_i
		libff::G1<ppT> t_com;
};


/*! Data structure for online part of the proof.
 * Contains shuffled ciphertexts and proof for the consistency argument.
 * \brief Online part of the proof.
 */
template<typename ppT>
struct Online_proof {
		//! Shuffled ciphertexts
		ElGamal_vector<ppT> output;

		//! Consistency proof (ElGamal ciphertext N)
		ElGamal_pair<ppT> consist;
};


/*! Pair that keeps both online and offline parts of the proof.
 * \tparam elliptic ppT group.
 */
template<typename ppT>
using Proof = pair<Offline_proof<ppT>, Online_proof<ppT>>;


/*! Prover generates a CRS-based non-interactive zero-knowledge proof
 * for a shuffle.
 * Prover is separated into offline and online phase.
 * In offline phase prover takes as an input a permutation and
 * randomizers for the ciphertexts.
 * Then it computes a proof for the permutation matrix argument and
 * the same-message argument.
 * In online phase prover gets ciphertexts as an input.
 * Prover generates shuffled ciphertexts using permutation and randomizers.
 * After that prover generates a proof for the consistency argument.
 * Verifier is able to separately check the proof for offline part.
 * Offline part of the proof must be generated before the online
 * part can be generated.
 *
 * \brief Computes zero-knowledge proof for a shuffle.
 * \tparam ppT elliptic ppT group
 */
template<typename ppT>
class Prover {

public:
		/*! Constructor for prover.
		 * \param crs common reference string generated by a trusted third party
		 */
		Prover(CRS<ppT> crs): crs {crs} {

#ifdef PARALLEL
		chunks = omp_get_max_threads(); //otherwise == 1
#endif
		}

		Proof<ppT> deserialize (json &y) {
				online_proof.output =
					utils::deserialize_ciphertexts(y["Online_proof"]["output"]);
				online_proof.consist =
					utils::deserialize_ElGamal_pair(y["Online_proof"]["consist"]);
				offline_proof.t_com =
					utils::set_coordinates_G1(y["Offline_proof"]["t_com"]);
				offline_proof.a_coms =
					utils::set_vector_coordinates_G1(y["Offline_proof"]["a_coms"]);
				offline_proof.a_hat_coms =
					utils::set_vector_coordinates_G1(y["Offline_proof"]["a_hat_coms"]);
				offline_proof.b_coms =
					utils::set_vector_coordinates_G2(y["Offline_proof"]["b_coms"]);
				offline_proof.uvs =
					utils::set_vector_coordinates_G1(y["Offline_proof"]["uvs"]);
				offline_proof.same_msgs =
					utils::set_vector_coordinates_G1(y["Offline_proof"]["same_msgs"]);
				Proof<ppT> proof {offline_proof, online_proof};
				return proof;
		}

		Proof<ppT> to_affine_proof () {
				utils::to_affine_coordinates(&offline_proof.t_com);
				utils::to_affine_coordinates(&offline_proof.a_coms);
				utils::to_affine_coordinates(&offline_proof.a_hat_coms);
				utils::to_affine_coordinates(&offline_proof.b_coms);
				utils::to_affine_coordinates(&offline_proof.uvs);
				utils::to_affine_coordinates(&offline_proof.same_msgs);
				Proof<ppT> proof {offline_proof, online_proof};
				return proof;
		}


		void print() {
				cout << "\n***** Offline ******\n";
				utils::print_field <libff::G1<ppT>> ("t_com", offline_proof.t_com);
				utils::print_field <libff::G1<ppT>> ("a_coms", offline_proof.a_coms);
				utils::print_field <libff::G1<ppT>> ("a_hat_coms",
																				offline_proof.a_hat_coms);
				utils::print_field <libff::G1<ppT>> ("uvs", offline_proof.uvs);
				utils::print_field <libff::G1<ppT>> ("same_msgs",
																				offline_proof.same_msgs);
				utils::print_field <libff::G2<ppT>> ("b_coms", offline_proof.b_coms);
		}

		/*! Computes offline part of the proof.
		 * This includes permutation matrix argument and the same-message argument.
		 *
		 * \param permutation permutation for shuffling the ciphertexts.
		 *	 Permutation is represented as vector containing numbers from 0 to n-1
		 *	 where n is the number of ciphertexts.
		 *	 Numbers should be in a uniformly random order.
		 * \param randomizers_cipher vector containing n uniformly
		 *	 random field values.
		 *	 Elements are used in the consistency argument and shuffle precomputation.
		 * \param randomizers_proof vector containing n - 1 uniformly
		 *	 random field values.
		 *	 Elements are used in the PM argument and SM argument.
		 * \param randomizer_t uniformly random field value.
		 *	 Element is used in the consistency argument.
		 * \returns offline part of the proof.
		 */
		Offline_proof<ppT> prove_offline(vector<long> permutation,
																		 libff::Fr_vector<ppT> randomizers_cipher,
																		 libff::Fr_vector<ppT> randomizers_proof,
																		 libff::Fr<ppT> randomizer_t) {
				libff::enter_block("Offline prover", true);

				libff::enter_block("Commitment to permutation matrix", false);
				libff::inhibit_profiling_info = true;

				// Precomputation for fixed-based multi-exponentiation of [rho]_1.
				const size_t g1_rho_exp_count = crs.n - 1;
				size_t g1_rho_window_size =
						libff::get_exp_window_size<libff::G1<ppT>>(g1_rho_exp_count);
				libff::window_table<libff::G1<ppT>> g1_rho_table =
						get_window_table(libff::Fr<ppT>::size_in_bits(),
														 g1_rho_window_size,
														 crs.g1_rho);

				// vector of elements (r_i * rho) * g1
				libff::G1_vector<ppT> g1_rho_ris = batch_exp(
																								libff::Fr<ppT>::size_in_bits(),
																								g1_rho_window_size,
																								g1_rho_table,
																								randomizers_proof);

				// Precomputation for fixed-based multi-exponentiation of [rho]_2.
				const size_t g2_rho_exp_count = crs.n - 1;
				size_t g2_rho_window_size = libff::get_exp_window_size<libff::G2<ppT>>(
																				g2_rho_exp_count);
				libff::window_table<libff::G2<ppT>> g2_rho_table =
														get_window_table(libff::Fr<ppT>::size_in_bits(),
																						 g2_rho_window_size,
																						 crs.g2_rho);

				// vector of elements [(r_i * rho)]_2
				libff::G2_vector<ppT> g2_rho_ris = batch_exp(
																				       libff::Fr<ppT>::size_in_bits(),
																							 g2_rho_window_size,
																						   g2_rho_table,
																						   randomizers_proof);

				// Precomputation for fixed-based multi-exponentiation of [\hat{rho}]_1.
				const size_t g1_hat_rho_exp_count = crs.n;
				size_t g1_hat_rho_window_size =
						libff::get_exp_window_size<libff::G1<ppT>>(g1_hat_rho_exp_count);
				libff::window_table<libff::G1<ppT>> g1_hat_rho_table =
														get_window_table(libff::Fr<ppT>::size_in_bits(),
																						 g1_hat_rho_window_size,
																						 crs.g1_rho_hat);

				// vector of elements [r_i * \hat{rho}]_1
				libff::G1_vector<ppT> g1_hat_rho_ris = batch_exp(
																								libff::Fr<ppT>::size_in_bits(),
																								g1_hat_rho_window_size,
																								g1_hat_rho_table,
																								randomizers_proof);

				// Compute commitment for i = 1, ..., n - 1
				for (long i = 0; i < crs.n - 1; i++) {
						long pos = permutation.at(i);
						offline_proof.a_coms.push_back(crs.g1_Pis.at(pos) +
																				   g1_rho_ris.at(i));
						offline_proof.b_coms.push_back(crs.g2_Pis.at(pos) +
																				   g2_rho_ris.at(i));
						offline_proof.a_hat_coms.push_back(crs.g1_Pi_hats.at(pos) +
																				   g1_hat_rho_ris.at(i));
				}


				// Compute commitment for i = n
				libff::G1<ppT> a_com_n = crs.g1_sum_Pi -
						vector_sum<libff::G1<ppT>>(offline_proof.a_coms);
				libff::G2<ppT> b_com_n = crs.g2_sum_Pi -
						vector_sum<libff::G2<ppT>>(offline_proof.b_coms);
				libff::G1<ppT> a_hat_com_n = crs.g1_sum_Pi_hat -
						vector_sum<libff::G1<ppT>>(offline_proof.a_hat_coms);

				libff::inhibit_profiling_info = false;
				libff::leave_block("Commitment to permutation matrix", false);

				// Compute the last proof randomizer r_n
				libff::Fr<ppT> rn = -vector_sum<libff::Fr<ppT>>(randomizers_proof);

				libff::enter_block("Permutation matrix proof", false);
				libff::inhibit_profiling_info = true;
				const libff::Fr<ppT> two = 2;
				// Compute unit vector proof
				for (long i = 0; i < crs.n - 1; i++) {
						long pos = permutation.at(i);
						libff::G1<ppT> temp = two * (offline_proof.a_coms.at(i) +
										crs.g1_P0);
						offline_proof.uvs.push_back( (randomizers_proof.at(i) *
												(temp - g1_rho_ris.at(i))) + crs.g1_Pi_longs.at(pos));
				}
				long pos = permutation.at(crs.n - 1);
				libff::G1<ppT> temp = (two * rn) * (a_com_n + crs.g1_P0);
				offline_proof.uvs.push_back( temp + ((- rn * rn) * crs.g1_rho) +
								crs.g1_Pi_longs.at(pos));
				libff::inhibit_profiling_info = false;
				libff::leave_block("Permutation matrix proof", false);


				libff::enter_block("Same-message proof", false);
				libff::inhibit_profiling_info = true;
				// Precomputation for fixed-base multi-exponentiation of
				// [beta * rho + \hat{beta} * \hat{rho}]_1.
				const size_t g1_beta_rho_exp_count = crs.n;
				size_t g1_beta_rho_window_size =
						libff::get_exp_window_size<libff::G1<ppT>>(g1_beta_rho_exp_count);
				libff::window_table<libff::G1<ppT>> g1_beta_rho_table =
						get_window_table(libff::Fr<ppT>::size_in_bits(),
														 g1_beta_rho_window_size,
														 crs.g1_beta_rho);

				// Compute [r_i * (beta * rho + \hat{beta} * \hat{rho}]_1
				libff::G1_vector<ppT> g1_beta_rho_r =
						batch_exp(libff::Fr<ppT>::size_in_bits(),
											g1_beta_rho_window_size,
											g1_beta_rho_table,
											randomizers_proof);
				g1_beta_rho_r.push_back(rn * crs.g1_beta_rho);

				// Compute same-message proof
				for (long i = 0; i < crs.n; i++) {
						long pos = permutation.at(i);
						offline_proof.same_msgs.push_back(crs.g1_Pi_longs2.at(pos) +
										g1_beta_rho_r.at(i));
				}
				libff::inhibit_profiling_info = false;
				libff::leave_block("Same-message proof", false);


				libff::enter_block("Commit to randomizers", false);
				libff::inhibit_profiling_info = true;
				// Commit to ciphertext randomizers
				offline_proof.t_com = libff::multi_exp<libff::G1<ppT>, libff::Fr<ppT>,
																						   libff::multi_exp_method_BDLO12>
																						  (crs.g1_Pi_hats.begin(),
																							 crs.g1_Pi_hats.end(),
																						   randomizers_cipher.begin(),
																							 randomizers_cipher.end(),
																							 chunks) +
														  windowed_exp(libff::Fr<ppT>::size_in_bits(),
																				   g1_hat_rho_window_size,
																				   g1_hat_rho_table,
																				   randomizer_t);

				libff::inhibit_profiling_info = false;
				libff::leave_block("Commit to randomizers", false);

				libff::enter_block("Shuffle precomputation", false);
				libff::inhibit_profiling_info = true;
				// Precomputation for fixed-based multi-exponentiation of g2.
				const size_t g2_exp_count = crs.n;
				size_t g2_window_size = libff::get_exp_window_size<libff::G2<ppT>>(
																		g2_exp_count);
				libff::window_table<libff::G2<ppT>> g2_table = get_window_table(
																								libff::Fr<ppT>::size_in_bits(),
																								g2_window_size,
																								libff::G2<ppT>::one());

				// Compute blinding values [t_i]
				blinders1 = batch_exp(libff::Fr<ppT>::size_in_bits(), g2_window_size,
														  g2_table, randomizers_cipher);

				// Precomputation for fixed-based multi-exponentiation of [sk]_2.
				const size_t g2_sk_exp_count = crs.n;
				size_t g2_sk_window_size = libff::get_exp_window_size<libff::G2<ppT>>(
																       g2_sk_exp_count);
				libff::window_table<libff::G2<ppT>> g2_sk_table = get_window_table(
																								libff::Fr<ppT>::size_in_bits(),
																								g2_sk_window_size,
																								crs.g2_sk);

				// Compute blinding values [sk * t_i]
				blinders2 = batch_exp(libff::Fr<ppT>::size_in_bits(),
														  g2_sk_window_size,
															g2_sk_table, randomizers_cipher);
				libff::inhibit_profiling_info = false;
				libff::leave_block("Shuffle precomputation", false);

				libff::leave_block("Offline prover", true);
				return offline_proof;
		} // prove_offline


		/*! Computes online part of the proof.
		 * This includes shuffling the ciphertexts and the consistency argument.
		 *
		 * \param ciphertexts input ciphertexts
		 * \param permutation permutation for shuffling the ciphertexts.
		 *	 Permutation is represented as vector containing numbers from 0 to n-1
		 *	 where n is the number of ciphertexts.
		 *	 Numbers should be in a uniformly random order.
		 * \param randomizers_proof vector containing n - 1 uniformly
		 *	 random field values. Elements are used in the consistency argument.
		 * \param randomizer_t uniformly random field value.
		 *	 Element is used in the consistency argument.
		 * \returns online part of the proof.
		 */
		Online_proof<ppT> prove_online(ElGamal_vector<ppT> ciphertexts,
																   vector<long> permutation,
																   libff::Fr_vector<ppT> randomizers_proof,
																	 libff::Fr<ppT> randomizer_t) {
				libff::enter_block("Online prover", false);
				libff::inhibit_profiling_info = true;

				online_proof.output = shuffle(ciphertexts, permutation);

				// Compute last pm randomizer r_n
				libff::Fr<ppT> rn = -vector_sum<libff::Fr<ppT>>(randomizers_proof);
				randomizers_proof.push_back(rn);

				libff::G2_vector<ppT> first_elems;
				libff::G2_vector<ppT> second_elems;
				for (ElGamal_pair<ppT> ciphertext : ciphertexts) {
						first_elems.push_back(ciphertext.first);
						second_elems.push_back(ciphertext.second);
				}

				libff::G2<ppT> consist1 = libff::multi_exp<libff::G2<ppT>,
																								   libff::Fr<ppT>,
																									 libff::multi_exp_method_BDLO12>
																								  (first_elems.begin(),
																									 first_elems.end(),
																									 randomizers_proof.begin(),
																									 randomizers_proof.end(),
																									 chunks);
				consist1 = consist1 + (randomizer_t * libff::G2<ppT>::one());

				libff::G2<ppT> consist2 = libff::multi_exp<libff::G2<ppT>,
																									 libff::Fr<ppT>,
																									 libff::multi_exp_method_BDLO12>
																									(second_elems.begin(),
																									 second_elems.end(),
																									 randomizers_proof.begin(),
																									 randomizers_proof.end(),
																									 chunks);
				consist2 = consist2 + (randomizer_t * crs.g2_sk);

				online_proof.consist = ElGamal_pair<ppT>(consist1, consist2);

				libff::inhibit_profiling_info = false;
				libff::leave_block("Online prover", false);
				return online_proof;
		} // prove_online


		/*! Shuffles the ciphertexts i.e. blinds the ciphertexts
		 * using randomizers and permutes with the permutation.
		 *
		 * \param ciphertexts input ciphertexts
		 * \param permutation permutation as a vector of numbers from 0, ..., n - 1,
		 *	 where n is the number of ciphertexts
		 * \returns shuffled vector of ciphertexts
		 */
		ElGamal_vector<ppT> shuffle(ElGamal_vector<ppT> ciphertexts,
																vector<long> permutation) {

				ElGamal_vector<ppT> shuffled_ciphertexts = ElGamal_vector<ppT>(crs.n);
				for (long i = 0; i < crs.n; i++) {
						long pos = permutation.at(i);
						shuffled_ciphertexts.at(pos) = ElGamal_pair<ppT>(
																								(ciphertexts.at(i)).first +
																										blinders1.at(pos),
																								(ciphertexts.at(i)).second +
																										blinders2.at(pos));
						}
				return shuffled_ciphertexts;
		} // shuffle


		/* Computes both offline and online parts of the proof.
		 *
		 * \param permutation permutation for shuffling the ciphertexts.
		 *	 Permutation is represented as vector containing numbers from 0 to n-1
		 *	 where n is the number of ciphertexts.
		 *	 Numbers should be in a uniformly random order.
		 * \param ciphertexts input ciphertexts
		 * \param randomizers_cipher vector containing n uniformly
		 *	 random field values. Elements are used to blind the permuted ciphertexts.
		 * \param randomizers_pm vector containing n uniformly random field values.
		 * 	Elements are used in the PM and PSP arguments.
		 * \param randomizers_sm vector containing n uniformly random field values.
		 *	 Elements are used in the SM and PSP arguments.
		 * \param randomizers_EU vector containing 2 uniformly random field values.
		 *	 Elements are used in the PSP argument.
		 * \return complete proof for the shuffle
		 */
		Proof<ppT> prove(ElGamal_vector<ppT> ciphertexts, vector<long> permutation,
										 libff::Fr_vector<ppT> randomizers_cipher,
									   libff::Fr_vector<ppT> randomizers_proof,
									   libff::Fr<ppT> randomizer_t) {
				libff::enter_block("Prover", true);
				offline_proof = prove_offline(permutation, randomizers_cipher,
																			randomizers_proof, randomizer_t);
				online_proof = prove_online(ciphertexts, permutation, randomizers_proof,
																		randomizer_t);
				Proof<ppT> proof {offline_proof, online_proof};
				libff::leave_block("Prover", true);
				return proof;
		}


private:
		//! Common reference string
		CRS<ppT> crs;

		//! Offline proof
		Offline_proof<ppT> offline_proof;

		//! Online proof
		Online_proof<ppT> online_proof;

		//! Number of threads that can be used
		size_t chunks = 1;

		//! Precomputation of [t_i]_2 for online prover
		libff::G2_vector<ppT> blinders1;

		//! Precomputation of [t_i * sk]_2 for online prover
		libff::G2_vector<ppT> blinders2;
}; // Prover

#endif /* PROVER_H_ */
