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

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

using namespace libsnark;
using namespace std;

int main(int argc, char **argv)
{
  enum {
    MODE_NONE,
    MODE_PRIV,
    MODE_PUB,
    MODE_CONST
  } mode = MODE_NONE;

  if (argc == 2) {
    if (!strcmp(argv[1], "priv")) {
        mode = MODE_PRIV;
    } else if (!strcmp(argv[1], "pub")) {
        mode = MODE_PUB;
    } else if (!strcmp(argv[1], "const")) {
        mode = MODE_CONST;
    }
  }
  if (mode == MODE_NONE) {
    cerr << "Usage: " << argv[0] << " mode" << endl << endl;
    cerr << "Where mode is one of:" << endl;
    cerr << " priv:  use private Ptable" << endl;
    cerr << " pub:   use public Ptable" << endl;
    cerr << " const: use constant public key (no Ptable)" << endl;
    exit(1);
  }

  // Initialize the curve parameters

  default_r1cs_gg_ppzksnark_pp::init_public_params();

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

  libff::start_profiling();

  cout << "Keypair" << endl;

  protoboard<FieldT> pb;
  pb_variable<FieldT> C1x, C1y, C2x, C2y, Kx, Ky;
  pb_variable<FieldT> Px, Py;
  pb_variable_array<FieldT> Ptable;
  pb_variable<FieldT> k, s, y, r;


  // Allocate variables

  size_t numbits = FieldT::num_bits;

  // Public outputs:

  // El Gamal encryption of k under public key P (or H if MODE_CONST)
  // C1 = r*G, C2 = r*P + M  (where M=(256*k+s,y))
  C1x.allocate(pb, "C1x");
  C1y.allocate(pb, "C1y");
  C2x.allocate(pb, "C2x");
  C2y.allocate(pb, "C2y");

  // Public key corresponding to private key k
  // K = k*G
  Kx.allocate(pb, "Kx");
  Ky.allocate(pb, "Ky");

  // Public inputs:

  // The public key P (if not MODE_CONST)
  if (mode != MODE_CONST) {
    Px.allocate(pb, "Px");
    Py.allocate(pb, "Py");

    // The Ptable might be public or private, according to the mode
    Ptable.allocate(pb, 2*numbits, "Ptable");
  }

  // Private inputs:
  // k is a 246-bit random number
  k.allocate(pb, "k");
  // s and y are such that M = (256*k+s,y) is a point on the curve
  s.allocate(pb, "s");
  y.allocate(pb, "y");
  // r is the randomness for the El Gamal encryption
  r.allocate(pb, "r");

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

  if (mode == MODE_PRIV) {
      pb.set_input_sizes(8);
  } else if (mode == MODE_PUB) {
      pb.set_input_sizes(8+2*numbits);
  } else if (mode == MODE_CONST) {
      pb.set_input_sizes(6);
  }

  // Initialize the gadgets

  // Curve parameters and generators
  FieldT curve_b("7950939520449436327800262930799465135910802758673292356620796789196167463969");
  FieldT Gx(0), Gy("11977228949870389393715360594190192321220966033310912010610740966317727761886");
  FieldT Hx(1), Hy("21803877843449984883423225223478944275188924769286999517937427649571474907279");

  // Prove (256*k+s,y) is on the curve
  pb_variable<FieldT> xsquared, ysquared;
  xsquared.allocate(pb, "xsquared");
  ysquared.allocate(pb, "ysquared");
  pb.add_r1cs_constraint(r1cs_constraint<FieldT>(y, y, ysquared));
  pb.add_r1cs_constraint(r1cs_constraint<FieldT>(k * 256 + s, k * 256 + s, xsquared));
  pb.add_r1cs_constraint(r1cs_constraint<FieldT>(xsquared - 3, k * 256 + s, ysquared - curve_b));

  // The unpacking gadgets to turn k and r into bits
  pb_variable_array<FieldT> kbits, rbits;
  kbits.allocate(pb, numbits-8, "kbits");
  rbits.allocate(pb, numbits, "rbits");
  packing_gadget<FieldT> kpacker(pb, kbits, k);
  packing_gadget<FieldT> rpacker(pb, rbits, r);
  kpacker.generate_r1cs_constraints(true);
  rpacker.generate_r1cs_constraints(true);

  // The El Gamal first component r*G
  ec_constant_scalarmul_vec_gadget<FieldT> C1gadget(pb, C1x, C1y, rbits, Gx, Gy);
  C1gadget.generate_r1cs_constraints();

  // The El Gamal intermediate value r*P
  pb_variable<FieldT> elgx, elgy;
  elgx.allocate(pb, "elgx");
  elgy.allocate(pb, "elgy");

  gadget<FieldT> *ElGgadgetp = NULL;
  if (mode == MODE_CONST) {
      ElGgadgetp = new ec_constant_scalarmul_vec_gadget<FieldT> (pb, elgx, elgy, rbits, Hx, Hy);
      (static_cast<ec_constant_scalarmul_vec_gadget<FieldT>*>(ElGgadgetp))->generate_r1cs_constraints();
  } else {
      ElGgadgetp = new ec_scalarmul_vec_gadget<FieldT> (pb, elgx, elgy, rbits, Px, Py, Ptable, mode == MODE_PRIV, true);
      (static_cast<ec_scalarmul_vec_gadget<FieldT>*>(ElGgadgetp))->generate_r1cs_constraints();
  }

  // The El Gamal second component r*P + M
  pb_linear_combination<FieldT> x;
  x.assign(pb, k * 256 + s);
  ec_add_gadget<FieldT> ElGfinal(pb, C2x, C2y, elgx, elgy, x, y);
  ElGfinal.generate_r1cs_constraints();

  // The generated public key k*G
  ec_constant_scalarmul_vec_gadget<FieldT> Kgadget(pb, Kx, Ky, kbits, Gx, Gy);
  Kgadget.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;
  
  if (mode != MODE_CONST) {
    // A variable base point P
    pb.val(Px) = FieldT("1095194319010475832867263440470707690447963461907735667341232728633587089702");
    pb.val(Py) = FieldT("9185463202887631101218413269806857706246311016297504828581985913021301344974");
  }
  pb.val(k) = FieldT("31329510635628557928212225120518124937732397714111203844965919301557399521");
  pb.val(s) = FieldT(1);
  pb.val(y) = FieldT("4364798287654239504994818950156019747851405522689486598132350453516910863367");
  pb.val(r) = FieldT::random_element();

  pb.val(xsquared) = (pb.val(k) * 256 + pb.val(s)).squared();
  pb.val(ysquared) = pb.val(y).squared();

  kpacker.generate_r1cs_witness_from_packed();
  rpacker.generate_r1cs_witness_from_packed();

  C1gadget.generate_r1cs_witness();
  if (mode == MODE_CONST) {
      (static_cast<ec_constant_scalarmul_vec_gadget<FieldT>*>(ElGgadgetp))->generate_r1cs_witness();
  } else {
      (static_cast<ec_scalarmul_vec_gadget<FieldT>*>(ElGgadgetp))->generate_r1cs_witness();
  }
  delete ElGgadgetp;

  x.evaluate(pb);
  ElGfinal.generate_r1cs_witness();

  Kgadget.generate_r1cs_witness();

  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());

  cout << "Verifier" << endl;

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

  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_verifenc_") + argv[1]);
  pkfile << keypair.pk;
  pkfile.close();
  ofstream vkfile(string("vk_verifenc_") + argv[1]);
  vkfile << keypair.vk;
  vkfile.close();
  ofstream pffile(string("proof_verifenc_") + argv[1]);
  pffile << proof;
  pffile.close();

  return 0;
}