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

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

using namespace libsnark;
using namespace std;

typedef enum {
    MODE_NONE,
    MODE_PRIV,
    MODE_PUB,
    MODE_CONST
} Mode;

// If a is a quadratic residue, set sqrt_a to one of its square roots
// (-sqrt_a will be the other) and return true.  Otherwise return false.
template<typename FieldT>
bool sqrt_if_possible(FieldT &sqrt_a, const FieldT &a)
{
  // A modification of the Tonelli-Shanks implementation from libff
  // to catch the case when you encounter a nonresidue

  const FieldT one = FieldT::one();

  size_t v = FieldT::s;
  FieldT z = FieldT::nqr_to_t;

  FieldT w = a^FieldT::t_minus_1_over_2;
  FieldT x = a * w;
  FieldT b = x * w; // b = a^t

  while (b != one)
  {
      size_t m = 0;
      FieldT b2m = b;
      while (b2m != one)
      {
          /* invariant: b2m = b^(2^m) after entering this loop */
          b2m = b2m.squared();
          m += 1;
      }
      if (m == v) {
          // Not a quadratic residue
          return false;
      }

      int j = v-m-1;
      w = z;
      while (j > 0)
      {
          w = w.squared();
          --j;
      } // w = z^2^(v-m-1)

      z = w.squared();
      b = b * z;
      x = x * w;
      v = m;
  }
  sqrt_a = x;
  return true;
}

template<typename FieldT>
class verified_encryption_gadget : public gadget<FieldT> {
private:
  const size_t numbits;
  FieldT curve_b, Gx, Gy, Hx, Hy;
  pb_variable<FieldT> r;
  pb_variable<FieldT> xsquared, ysquared;
  pb_variable_array<FieldT> kbits, rbits;
  pb_variable<FieldT> elgx, elgy;
  pb_linear_combination<FieldT> x;
  pb_variable<FieldT> s, y;
  vector<packing_gadget<FieldT> > packers;
  vector<ec_constant_scalarmul_vec_gadget<FieldT> > constmuls;
  vector<ec_scalarmul_vec_gadget<FieldT> > muls;
  vector<ec_add_gadget<FieldT> > adders;

public:
  const Mode mode;
  const pb_variable<FieldT> C1x, C1y, C2x, C2y, Kx, Ky;
  const pb_variable<FieldT> Px, Py;
  const pb_variable_array<FieldT> Ptable;
  const pb_variable<FieldT> k;

  verified_encryption_gadget(protoboard<FieldT> &pb,
              Mode mode,
              const pb_variable<FieldT> &C1x,
              const pb_variable<FieldT> &C1y,
              const pb_variable<FieldT> &C2x,
              const pb_variable<FieldT> &C2y,
              const pb_variable<FieldT> &Kx,
              const pb_variable<FieldT> &Ky,
              const pb_variable<FieldT> &Px,
              const pb_variable<FieldT> &Py,
              const pb_variable_array<FieldT> &Ptable,
              const pb_variable<FieldT> &k) :
    gadget<FieldT>(pb, "verified_encryption_gadget"),
    // Curve parameters and generators
    numbits(FieldT::num_bits),
    curve_b("7950939520449436327800262930799465135910802758673292356620796789196167463969"),
    Gx(0), Gy("11977228949870389393715360594190192321220966033310912010610740966317727761886"),
    Hx(1), Hy("21803877843449984883423225223478944275188924769286999517937427649571474907279"),
    mode(mode), C1x(C1x), C1y(C1y), C2x(C2x), C2y(C2y),
    Kx(Kx), Ky(Ky), Px(Px), Py(Py), Ptable(Ptable), k(k)
  {
    s.allocate(pb, "s");
    y.allocate(pb, "y");
    r.allocate(pb, "r");
    xsquared.allocate(pb, "xsquared");
    ysquared.allocate(pb, "ysquared");
    kbits.allocate(pb, numbits-8, "kbits");
    rbits.allocate(pb, numbits, "rbits");

    // The unpacking gadgets to turn k and r into bits
    packers.emplace_back(pb, kbits, k);
    packers.emplace_back(pb, rbits, r);

    // The El Gamal first component r*G
    constmuls.emplace_back(pb, C1x, C1y, rbits, Gx, Gy);

    // The El Gamal intermediate value r*P
    elgx.allocate(pb, "elgx");
    elgy.allocate(pb, "elgy");
    if (mode == MODE_CONST) {
        constmuls.emplace_back(pb, elgx, elgy, rbits, Hx, Hy);
    } else {
        muls.emplace_back(pb, elgx, elgy, rbits, Px, Py, Ptable, mode == MODE_PRIV, true);
    }

    // The El Gamal second component r*P + M
    x.assign(pb, k * 256 + s);
    adders.emplace_back(pb, C2x, C2y, elgx, elgy, x, y);

    // The generated public key k*G
    constmuls.emplace_back(pb, Kx, Ky, kbits, Gx, Gy);
  }

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

    for (auto&& gadget : packers) {
        gadget.generate_r1cs_constraints(true);
    }

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

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

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

  void find_s_y(const FieldT &kval, FieldT &sval, FieldT &yval)
  {

    FieldT s_candidate = 0;
    while (true) {
        FieldT x_candidate = kval*256+s_candidate;
        FieldT ysq_candidate = (x_candidate.squared() - 3)*x_candidate + curve_b;

        if (sqrt_if_possible<FieldT>(yval, ysq_candidate)) {
            sval = s_candidate;
            return;
        }
        s_candidate += 1;
    }
  }

  void generate_r1cs_witness()
  {
    // Find an s and y such that x^3 - 3*x + b = y^2, where x = 256*k + s
    FieldT sval, yval;
    find_s_y(this->pb.val(k), sval, yval);
    this->pb.val(s) = sval;
    this->pb.val(y) = yval;

    this->pb.val(r) = FieldT::random_element();
    this->pb.val(xsquared) = (this->pb.val(k) * 256 + this->pb.val(s)).squared();
    this->pb.val(ysquared) = this->pb.val(y).squared();
    x.evaluate(this->pb);
    for (auto&& gadget : packers) {
        gadget.generate_r1cs_witness_from_packed();
    }

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

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

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

int main(int argc, char **argv)
{
  Mode mode = MODE_NONE;
  size_t numverifencs = 1;

  if (argc == 2 || argc == 3) {
    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 (argc == 3) {
        numverifencs = atoi(argv[2]);
    }
  }
  if (mode == MODE_NONE || numverifencs < 1) {
    cerr << "Usage: " << argv[0] << " mode n" << 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 << endl;
    cerr << "and where n is the number of verifencs in the circuit" << 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> C1x[numverifencs], C1y[numverifencs];
  pb_variable<FieldT> C2x[numverifencs], C2y[numverifencs];
  pb_variable<FieldT> Kx[numverifencs], Ky[numverifencs];
  pb_variable<FieldT> Px[numverifencs], Py[numverifencs];
  pb_variable_array<FieldT> Ptable[numverifencs];
  pb_variable<FieldT> k[numverifencs];

  const size_t numbits = FieldT::num_bits;

  // Allocate variables

  // Public outputs:

  for (size_t i = 0; i < numverifencs; ++i) {
      // 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[i].allocate(pb, "C1x");
      C1y[i].allocate(pb, "C1y");
      C2x[i].allocate(pb, "C2x");
      C2y[i].allocate(pb, "C2y");

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

      // Public inputs:

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

  if (mode != MODE_CONST) {
      for (size_t i = 0; i < numverifencs; ++i) {
        // The Ptable might be public or private, according to the mode
        Ptable[i].allocate(pb, 2*numbits, "Ptable");
      }
  }

  for (size_t i = 0; i < numverifencs; ++i) {
      // Private inputs:
      // k is a 246-bit random number
      k[i].allocate(pb, "k");
  }

  // 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*numverifencs);
  } else if (mode == MODE_PUB) {
      pb.set_input_sizes(8*numverifencs+2*numbits*numverifencs);
  } else if (mode == MODE_CONST) {
      pb.set_input_sizes(6*numverifencs);
  }

  // Initialize the gadgets
  vector<verified_encryption_gadget<FieldT> > vencs;

  for (size_t i = 0; i < numverifencs; ++i) {
      vencs.emplace_back(pb, mode, C1x[i], C1y[i], C2x[i], C2y[i], Kx[i], Ky[i], Px[i], Py[i], Ptable[i], k[i]);
  }
  for (auto&& gadget : vencs) {
      gadget.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) {
      FieldT curve_b("7950939520449436327800262930799465135910802758673292356620796789196167463969");
      for (size_t i = 0; i < numverifencs; ++i) {
        // A variable base point P
        FieldT x, y, ysq;
        do {
            x = FieldT::random_element();
            ysq = (x.squared() - 3)*x + curve_b;
        } while (!sqrt_if_possible<FieldT>(y, ysq));

        pb.val(Px[i]) = x;
        pb.val(Py[i]) = y;
      }
  }
  gmp_randstate_t randstate;
  gmp_randinit_default(randstate);
  FieldT seed = FieldT::random_element();
  mpz_t seed_mpz;
  mpz_init(seed_mpz);
  seed.mont_repr.to_mpz(seed_mpz);
  gmp_randseed(randstate, seed_mpz);
  mpz_clear(seed_mpz);
  for (size_t i = 0; i < numverifencs; ++i) {
      mpz_t kval;
      mpz_init(kval);
      mpz_urandomb(kval, randstate, 246);
      pb.val(k[i]) = FieldT(kval);
      mpz_clear(kval);
  }

  libff::enter_block("PROVER TIME");

  for (auto&& gadget : vencs) {
      gadget.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());

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

  return 0;
}