#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 = MODE_NONE;

  if (argc == 2) {
    if (!strcmp(argv[1], "priv")) {
        mode = MODE_PRIV;
    } else if (!strcmp(argv[1], "pub")) {
        mode = MODE_PUB;
    }
  }
  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;
    exit(1);
  }

  // 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> outx, outy;
  pb_variable<FieldT> Px, Py;
  pb_variable<FieldT> s;
  pb_variable_array<FieldT> Ptable;


  // Allocate variables

  size_t numbits = FieldT::num_bits;
  outx.allocate(pb, "outx");
  outy.allocate(pb, "outy");
  Px.allocate(pb, "Px");
  Py.allocate(pb, "Py");
  Ptable.allocate(pb, 2*numbits, "Ptable");
  s.allocate(pb, "s");

  // 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(4);
  } else {
      pb.set_input_sizes(4+2*numbits);
  }

  // Initialize the gadget
  ec_scalarmul_gadget<FieldT> sm(pb, outx, outy, s, Px, Py, Ptable, mode == MODE_PRIV, true);
  sm.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(s) = FieldT::random_element();
  // A variable base point P
  pb.val(Px) = FieldT("1095194319010475832867263440470707690447963461907735667341232728633587089702");
  pb.val(Py) = FieldT("9185463202887631101218413269806857706246311016297504828581985913021301344974");
  cout << "Computing " << pb.val(s) << "*G" << endl;

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

  cout << pb.val(s) << "*P" << " = (" << pb.val(outx) << ", " << pb.val(outy) << ")" << endl;


  return 0;
}