#include <stdlib.h>
#include <iostream>
#include <fstream>

#include "ecgadget.hpp"
#include "scalarmul.hpp"

using namespace libsnark;
using namespace std;

// A gadget for, given private input scalar r0,
// computing private scalars r1, d0, d1 and public
// EC points (D0x,D0y) and (D1x,D1y), such that:
// r1 = h_2(r0)
// d0 = h_1(r0)
// d1 = h_1(r1)
// D0 = d0 * G
// D1 = d1 * G
//
// where h_1(r) is the x-coordinate of r*H1 and
// h_2(r) is the x-coordinate of r*H2.
template<typename FieldT>
class ratchet_commit_gadget : public gadget<FieldT> {
private:
  FieldT Gx, Gy, H1x, H1y, H2x, H2y;
  pb_variable<FieldT> r1y, d0y, d1y;
  vector<ec_constant_scalarmul_gadget<FieldT> > constmuls;

public:
  const pb_variable<FieldT> r0, r1, d0, d1, D0x, D0y, D1x, D1y;

  ratchet_commit_gadget(protoboard<FieldT> &pb,
              const pb_variable<FieldT> &r0,
              const pb_variable<FieldT> &r1,
              const pb_variable<FieldT> &d0,
              const pb_variable<FieldT> &d1,
              const pb_variable<FieldT> &D0x,
              const pb_variable<FieldT> &D0y,
              const pb_variable<FieldT> &D1x,
              const pb_variable<FieldT> &D1y) :
    gadget<FieldT>(pb, "ratchet_commit_gadget"),
    // Curve parameters and generators
    Gx(0), Gy("11977228949870389393715360594190192321220966033310912010610740966317727761886"),
    H1x(1), H1y("21803877843449984883423225223478944275188924769286999517937427649571474907279"),
    H2x(3), H2y("5020743718369453748575779309408113228867962046286774659221819240049391841511"),
    r0(r0), r1(r1), d0(d0), d1(d1), D0x(D0x), D0y(D0y), D1x(D1x), D1y(D1y)
  {
    r1y.allocate(pb, "r1y");
    d0y.allocate(pb, "d0y");
    d1y.allocate(pb, "d1y");

    // r1 = [r0*H2]_x
    constmuls.emplace_back(pb, r1, r1y, r0, H2x, H2y);
    // d0 = [r0*H1]_x
    constmuls.emplace_back(pb, d0, d0y, r0, H1x, H1y);
    // d1 = [r1*H1]_x
    constmuls.emplace_back(pb, d1, d1y, r1, H1x, H1y);
    // D0 = d0*G
    constmuls.emplace_back(pb, D0x, D0y, d0, Gx, Gy);
    // D1 = d1*G
    constmuls.emplace_back(pb, D1x, D1y, d1, Gx, Gy);
  }

  void generate_r1cs_constraints()
  {
    for (auto&& gadget : constmuls) {
        gadget.generate_r1cs_constraints();
    }
  }

  void generate_r1cs_witness()
  {
    for (auto&& gadget : constmuls) {
        gadget.generate_r1cs_witness();
    }
  }
};

int main(int argc, char **argv)
{
  // Initialize the curve parameters

  default_r1cs_gg_ppzksnark_pp::init_public_params();
  init_curveparams();

  typedef libff::Fr<default_r1cs_gg_ppzksnark_pp> FieldT;
  
  // Create protoboard

  libff::start_profiling();

  cout << "Keypair" << endl;

  protoboard<FieldT> pb;
  pb_variable<FieldT> r0, r1, d0, d1, D0x, D0y, D1x, D1y;

  // Allocate variables

  // Public outputs:

  D0x.allocate(pb, "D0x");
  D0y.allocate(pb, "D0y");
  D1x.allocate(pb, "D1x");
  D1y.allocate(pb, "D1y");

  // Private inputs:

  r0.allocate(pb, "r0");

  // Private outputs:

  r1.allocate(pb, "r1");
  d0.allocate(pb, "d0");
  d1.allocate(pb, "d1");

  // This sets up the protoboard variables so that the first n of them
  // represent the public values and the rest is private

  pb.set_input_sizes(4);

  // Initialize the gadgets
  ratchet_commit_gadget<FieldT> rcom(pb, r0, r1, d0, d1, D0x, D0y, D1x, D1y);

  rcom.generate_r1cs_constraints();

  const r1cs_constraint_system<FieldT> constraint_system = pb.get_constraint_system();

  const r1cs_gg_ppzksnark_keypair<default_r1cs_gg_ppzksnark_pp> keypair = r1cs_gg_ppzksnark_generator<default_r1cs_gg_ppzksnark_pp>(constraint_system);

  // Add witness values

  cout << "Prover" << endl;
  
  pb.val(r0) = FieldT::random_element();

  libff::enter_block("PROVER TIME");

  rcom.generate_r1cs_witness();

  cout << "r0 = " << pb.val(r0) << endl;
  cout << "r1 = " << pb.val(r1) << endl;
  cout << "d0 = " << pb.val(d0) << endl;
  cout << "d1 = " << pb.val(d1) << endl;
  cout << "D0 = (" << pb.val(D0x) << ", " << pb.val(D0y) << ")" << endl;
  cout << "D1 = (" << pb.val(D1x) << ", " << pb.val(D1y) << ")" << endl;

  const r1cs_gg_ppzksnark_proof<default_r1cs_gg_ppzksnark_pp> proof = r1cs_gg_ppzksnark_prover<default_r1cs_gg_ppzksnark_pp>(keypair.pk, pb.primary_input(), pb.auxiliary_input());

  libff::leave_block("PROVER TIME");

  cout << "Verifier" << endl;

  libff::enter_block("VERIFIER TIME");

  bool verified = r1cs_gg_ppzksnark_verifier_strong_IC<default_r1cs_gg_ppzksnark_pp>(keypair.vk, pb.primary_input(), proof);

  libff::leave_block("VERIFIER TIME");

  cout << "Number of R1CS constraints: " << constraint_system.num_constraints() << endl;
  cout << "Primary (public) input length: " << pb.primary_input().size() << endl;
//  cout << "Primary (public) input: " << pb.primary_input() << endl;
  cout << "Auxiliary (private) input length: " << pb.auxiliary_input().size() << endl;
//  cout << "Auxiliary (private) input: " << pb.auxiliary_input() << endl;
  cout << "Verification status: " << verified << endl;

  ofstream pkfile(string("pk_ratchetcom"));
  pkfile << keypair.pk;
  pkfile.close();
  ofstream vkfile(string("vk_ratchetcom"));
  vkfile << keypair.vk;
  vkfile.close();
  ofstream pffile(string("proof_ratchetcom"));
  pffile << proof;
  pffile.close();

  return 0;
}