#include "libsnark_headers.hpp"

using namespace libsnark;

// There are two types of values:
//   _constants_ are values known at circuit generation time; they
//       are global constants known to everyone
//   _variables_ are values that change in each use of the circuit;
//       they have two subtypes:
//
//   _public variables_ are values known to both the prover
//       and verifier but change in each use of the circuit
//   _private variables_ are values known only to the prover
//       and change in each use of the circuit

// Double a constant EC point (inx,iny) to yield (outx,outy).  The input
// point must not be the point at infinity.
template<typename FieldT>
static void ec_double_point(FieldT &outx, FieldT &outy,
    const FieldT &inx, const FieldT &iny)
{
    FieldT xsq = inx.squared();
    FieldT lambda = (xsq * 3 - 3) * (iny * 2).inverse();
    outx = lambda.squared() - inx * 2;
    outy = lambda * (inx - outx) - iny;
}

// Add constant EC points (inx, iny) and (addx, addy) to yield (outx, outy).
// inx and addx must not be equal.
template<typename FieldT>
static void ec_add_points(FieldT &outx, FieldT &outy,
    const FieldT &inx, const FieldT &iny,
    const FieldT &addx, const FieldT &addy)
{
    FieldT lambda = (addy - iny) * (addx - inx).inverse();
    outx = lambda.squared() - (addx + inx);
    outy = lambda * (inx - outx) - iny;
}

// Double the variable EC point (inx,iny) to yield (outx,outy).  The
// input point must not be the point at infinity.
template<typename FieldT>
class ec_double_gadget : public gadget<FieldT> {
private:
  pb_variable<FieldT> lambda, inxsq;
public:
  const pb_variable<FieldT> outx, outy, inx, iny;

  ec_double_gadget(protoboard<FieldT> &pb,
              const pb_variable<FieldT> &outx,
              const pb_variable<FieldT> &outy,
              const pb_linear_combination<FieldT> &inx,
              const pb_linear_combination<FieldT> &iny) :
    gadget<FieldT>(pb, "ec_double_gadget"), outx(outx), outy(outy),
        inx(inx), iny(iny)
  {
    // Allocate variables to protoboard
    // The strings (like "x") are only for debugging purposes
	  
    lambda.allocate(this->pb, "lambda");
    inxsq.allocate(this->pb, "inxsq");
  }

  void generate_r1cs_constraints()
  {
    // inxsq = inx * inx
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(inx, inx, inxsq));

    // 2 * iny * lambda = 3 * inxsq - 3  (a = -3 on our curve)
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(2 * iny, lambda, 3 * inxsq - 3));

    // outx = lambda^2 - 2 * inx
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(lambda, lambda, outx + 2 * inx));

    // outy = lambda * (inx - outx) - iny
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(lambda, inx - outx, outy + iny));

  }

  void generate_r1cs_witness()
  {
    this->pb.val(inxsq) = this->pb.lc_val(inx) * this->pb.lc_val(inx);
    this->pb.val(lambda) = (this->pb.val(inxsq) * 3 - 3) * (this->pb.lc_val(iny) * 2).inverse();
    this->pb.val(outx) = this->pb.val(lambda).squared() - this->pb.lc_val(inx) * 2;
    this->pb.val(outy) = this->pb.val(lambda) * (this->pb.lc_val(inx) - this->pb.val(outx)) - this->pb.lc_val(iny);
  }
};

// Add the variable EC point (addx,addy) to the variable EC point
// (inx,iny) to yield (outx,outy).  The input point must not be the
// point at infinity.
template<typename FieldT>
class ec_add_gadget : public gadget<FieldT> {
private:
  pb_variable<FieldT> lambda;
public:
  const pb_variable<FieldT> outx, outy;
  const pb_linear_combination<FieldT> inx, iny, addx, addy;

  ec_add_gadget(protoboard<FieldT> &pb,
              const pb_variable<FieldT> &outx,
              const pb_variable<FieldT> &outy,
              const pb_linear_combination<FieldT> &inx,
              const pb_linear_combination<FieldT> &iny,
              const pb_linear_combination<FieldT> &addx,
              const pb_linear_combination<FieldT> &addy) :
    gadget<FieldT>(pb, "ec_add_gadget"),
    outx(outx), outy(outy), inx(inx), iny(iny), addx(addx), addy(addy)
  {
    // Allocate variables to protoboard
    // The strings (like "x") are only for debugging purposes
	  
    lambda.allocate(this->pb, "lambda");
  }

  void generate_r1cs_constraints()
  {
    // (addx - inx) * lambda = addy - iny
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(addx - inx, lambda, addy - iny));

    // outx = lambda^2 - (addx + inx)
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(lambda, lambda, outx + addx + inx));

    // outy = lambda * (inx - outx) - iny
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(lambda, inx - outx, outy + iny));

  }

  void generate_r1cs_witness()
  {
    this->pb.val(lambda) = (this->pb.lc_val(addy) - this->pb.lc_val(iny)) * (this->pb.lc_val(addx) - this->pb.lc_val(inx)).inverse();
    this->pb.val(outx) = this->pb.val(lambda).squared() - (this->pb.lc_val(addx) + this->pb.lc_val(inx));
    this->pb.val(outy) = this->pb.val(lambda) * (this->pb.lc_val(inx) - this->pb.val(outx)) - this->pb.lc_val(iny);
  }
};

// Add the variable EC point P to the constant EC point (inx,iny) to
// yield (outx,outy).  The input point must not be the point at
// infinity.
template<typename FieldT>
class ec_constant_add_gadget : public gadget<FieldT> {
private:
  pb_variable<FieldT> lambda;
public:
  const pb_variable<FieldT> outx, outy;
  const pb_linear_combination<FieldT> inx, iny;
  const FieldT Px, Py;

  ec_constant_add_gadget(protoboard<FieldT> &pb,
              const pb_variable<FieldT> &outx,
              const pb_variable<FieldT> &outy,
              const pb_linear_combination<FieldT> &inx,
              const pb_linear_combination<FieldT> &iny,
              const FieldT &Px, const FieldT &Py) :
    gadget<FieldT>(pb, "ec_constant_add_gadget"),
    outx(outx), outy(outy), inx(inx), iny(iny), Px(Px), Py(Py)
  {
    // Allocate variables to protoboard
    // The strings (like "x") are only for debugging purposes
	  
    lambda.allocate(this->pb, "lambda");
  }

  void generate_r1cs_constraints()
  {
    // (Px - inx) * lambda = Py - iny
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(Px - inx, lambda, Py - iny));

    // outx = lambda^2 - (Px + inx)
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(lambda, lambda, outx + Px + inx));

    // outy = lambda * (inx - outx) - iny
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(lambda, inx - outx, outy + iny));

  }

  void generate_r1cs_witness()
  {
    this->pb.val(lambda) = (Py - this->pb.lc_val(iny)) * (Px - this->pb.lc_val(inx)).inverse();
    this->pb.val(outx) = this->pb.val(lambda).squared() - (Px + this->pb.lc_val(inx));
    this->pb.val(outy) = this->pb.val(lambda) * (this->pb.lc_val(inx) - this->pb.val(outx)) - this->pb.lc_val(iny);
  }
};

// Add nothing or the constant EC point P to the variable EC point
// (inx,iny) to yield (outx,outy).  The input point must not be the
// point at infinity.  The input bit do_add controls whether the
// addition is done.
template<typename FieldT>
class ec_conditional_constant_add_gadget : public gadget<FieldT> {
private:
  pb_variable<FieldT> sumx, sumy;
  std::vector<ec_constant_add_gadget<FieldT> > adder;
public:
  const pb_variable<FieldT> outx, outy;
  const pb_linear_combination<FieldT> inx, iny;
  const pb_variable<FieldT> do_add;
  const FieldT Px, Py;

  ec_conditional_constant_add_gadget(protoboard<FieldT> &pb,
              const pb_variable<FieldT> &outx,
              const pb_variable<FieldT> &outy,
              const pb_linear_combination<FieldT> &inx,
              const pb_linear_combination<FieldT> &iny,
              const pb_variable<FieldT> &do_add,
              const FieldT &Px, const FieldT &Py) :
    gadget<FieldT>(pb, "ec_conditional_constant_add_gadget"),
    outx(outx), outy(outy), inx(inx), iny(iny), do_add(do_add),
    Px(Px), Py(Py)
  {
    // Allocate variables to protoboard
    // The strings (like "x") are only for debugging purposes
	  
    sumx.allocate(this->pb, "sumx");
    sumy.allocate(this->pb, "sumy");
    adder.emplace_back(this->pb, sumx, sumy, inx, iny, Px, Py);
  }

  void generate_r1cs_constraints()
  {
    // Strategy: we always do the addition, but if do_add = 0, we throw
    // it away later.

    adder[0].generate_r1cs_constraints();

    // Now we want to conditionally move the sum.  We want that
    // outx = do_add ? sumx : inx
    // outy = do_add ? sumy : iny

    // so we compute
    // outx = inx + (sumx - inx) * do_add
    // outy = iny + (sumy - iny) * do_add

    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(sumx - inx, do_add, outx - inx));
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(sumy - iny, do_add, outy - iny));

  }

  void generate_r1cs_witness()
  {
    adder[0].generate_r1cs_witness();

    bool move = this->pb.val(do_add) != FieldT(0);
    this->pb.val(outx) = move ? this->pb.val(sumx) : this->pb.lc_val(inx);
    this->pb.val(outy) = move ? this->pb.val(sumy) : this->pb.lc_val(iny);
  }
};

// Add nothing, or one of the constant EC points P1, P2, or P3 to the EC
// point (inx,iny) to yield (outx,outy).  The input point must not be
// the point at infinity.  The two input bits add1 and add2 control what
// is added.  Typically, P3 will equal P1+P2, in which case this gadget
// does two conditional constant adds simultaneously in just 6 constraints.
template<typename FieldT>
class ec_add_P123_gadget : public gadget<FieldT> {
private:
  pb_variable<FieldT> lambda, sumx, sumy, move;
public:
  const pb_variable<FieldT> outx, outy;
  const pb_linear_combination<FieldT> inx, iny;
  const pb_variable<FieldT> add1, add2;
  const FieldT P1x, P1y, P2x, P2y, P3x, P3y;

  ec_add_P123_gadget(protoboard<FieldT> &pb,
              const pb_variable<FieldT> &outx,
              const pb_variable<FieldT> &outy,
              const pb_linear_combination<FieldT> &inx,
              const pb_linear_combination<FieldT> &iny,
              const pb_variable<FieldT> &add1,
              const pb_variable<FieldT> &add2,
              const FieldT &P1x, const FieldT &P1y,
              const FieldT &P2x, const FieldT &P2y,
              const FieldT &P3x, const FieldT &P3y) : 
    gadget<FieldT>(pb, "ec_add_P123_gadget"),
    outx(outx), outy(outy), inx(inx), iny(iny), add1(add1), add2(add2),
    P1x(P1x), P1y(P1y), P2x(P2x), P2y(P2y), P3x(P3x), P3y(P3y)
  {
    // Allocate variables to protoboard
    // The strings (like "x") are only for debugging purposes
	  
    lambda.allocate(this->pb, "lambda");
    sumx.allocate(this->pb, "sumx");
    sumy.allocate(this->pb, "sumy");
    move.allocate(this->pb, "move");
  }

  void generate_r1cs_constraints()
  {
    // Strategy: if add1 = add2 = 0, we compute some nonsense but throw
    // it away later.  Otherwise, the coordinates of the point to add
    // are a _linear_ function of add1 and add2 (since P1, P2, and P3
    // are public constants)

    // In particular, the point to add is ( (P3x - P2x) * add1 + (P3x -
    // P1x) * add2 + (P1x + P2x - P3x), (P3y - P2y) * add1 + (P3y - P1y) *
    // add2 + (P1y + P2y - P3y))

    // (addx - inx) * lambda = addy - iny
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>((P3x - P2x) * add1 + (P3x - P1x) * add2 + (P1x + P2x - P3x) - inx, lambda, (P3y - P2y) * add1 + (P3y - P1y) * add2 + (P1y + P2y - P3y) - iny));

    // sumx = lambda^2 - (addx + inx)
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(lambda, lambda, sumx + (P3x - P2x) * add1 + (P3x - P1x) * add2 + (P1x + P2x - P3x) + inx));

    // sumy = lambda * (inx - sumx) - iny
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(lambda, inx - sumx, sumy + iny));

    // Now we want to conditionally move the sum.  We want that
    // outx = (add1 || add2) ? sumx : inx
    // outy = (add1 || add2) ? sumy : iny

    // so we compute move = add1 || add2, and then
    // outx = inx + (sumx - inx) * move
    // outy = iny + (sumy - iny) * move

    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(1 - add1, 1 - add2, 1 - move));
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(sumx - inx, move, outx - inx));
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(sumy - iny, move, outy - iny));

  }

  void generate_r1cs_witness()
  {
    FieldT addxval = (P3x - P2x) * this->pb.val(add1) + (P3x - P1x) * this->pb.val(add2) + (P1x + P2x - P3x);
    FieldT addyval = (P3y - P2y) * this->pb.val(add1) + (P3y - P1y) * this->pb.val(add2) + (P1y + P2y - P3y);
    this->pb.val(lambda) = (addyval - this->pb.lc_val(iny)) * (addxval - this->pb.lc_val(inx)).inverse();
    this->pb.val(sumx) = this->pb.val(lambda).squared() - (addxval + this->pb.lc_val(inx));
    this->pb.val(sumy) = this->pb.val(lambda) * (this->pb.lc_val(inx) - this->pb.val(sumx)) - this->pb.lc_val(iny);

    bool a1 = this->pb.val(add1) != FieldT(0);
    bool a2 = this->pb.val(add2) != FieldT(0);
    this->pb.val(move) = a1 || a2;
    this->pb.val(outx) = (a1 || a2) ? this->pb.val(sumx) : this->pb.lc_val(inx);
    this->pb.val(outy) = (a1 || a2) ? this->pb.val(sumy) : this->pb.lc_val(iny);
  }
};

// Compute a*P as (outx, outy) for a given constant point P, given a
// as a bit vector.  The _caller_ is responsible for proving that the
// elements of avec are bits.
template<typename FieldT>
class ec_constant_scalarmul_vec_gadget : public gadget<FieldT> {
private:
  FieldT Cx, Cy, CPx, CPy;
  pb_variable_array<FieldT> accumx, accumy;
  std::vector<ec_add_P123_gadget<FieldT> > conddoubleadders;
  std::vector<ec_conditional_constant_add_gadget<FieldT> > condsingleadders;
  std::vector<ec_constant_add_gadget<FieldT> > singleadders;
public:
  const pb_variable<FieldT> outx, outy;
  const pb_variable_array<FieldT> avec;
  const FieldT Px, Py;

  // Strategy: We compute (as compile-time constants) (powers of 2)
  // times P, and then conditionally add them into an accumulator.
  // Because our adder cannot handle the point at infinity O, we start
  // the accumulator with a value of C, whose discrete log with respect
  // to P should be unknown, so that we won't encounter O along the way.
  // (Also, a should not be 0 or the group order.)  We actually start
  // the accumulator with either C or C+P depending on avec[0], so we
  // get the first conditional add "for free".  Then we use the
  // ec_add_P123_gadget to do the conditional adds two at a time (at a
  // cost of 6 constraints per pair, as opposed to 5 for a single
  // conditional add).  If the length of avec is even, then there will
  // be one left over, and we do a single conditional add for that one.
  // Finally, we add the public point -C.

  ec_constant_scalarmul_vec_gadget(protoboard<FieldT> &pb,
              const pb_variable<FieldT> &outx,
              const pb_variable<FieldT> &outy,
              const pb_variable_array<FieldT> &avec,
              const FieldT &Px, const FieldT &Py) :
    gadget<FieldT>(pb, "ec_constant_scalarmul_vec_gadget"),
    // Precomputed coordinates of C
    Cx(2),
    Cy("4950745124018817972378217179409499695353526031437053848725554590521829916331"),
    outx(outx), outy(outy), avec(avec), Px(Px), Py(Py)
  {
    size_t numbits = avec.size();
    accumx.allocate(this->pb, numbits/2+1, "accumx");
    accumy.allocate(this->pb, numbits/2+1, "accumy");

    ec_add_points(CPx, CPy, Cx, Cy, Px, Py);

    FieldT twoiPx, twoiPy, twoi1Px, twoi1Py, twoi3Px, twoi3Py;
    size_t i = 1;

    ec_double_point(twoiPx, twoiPy, Px, Py);

    while(i < numbits) {
        // Invariants: i is odd, and twoiP = 2^i * P
        // Compute twoi1P = 2^{i+1} * P = 2 * twoiP and
        //         twoi3P = 2^i * 3 * P = 3 * twoiP
        ec_double_point(twoi1Px, twoi1Py, twoiPx, twoiPy);
        ec_add_points(twoi3Px, twoi3Py, twoi1Px, twoi1Py, twoiPx, twoiPy);

        if (i == numbits-1) {
            // There's only one bit of avec left; use a single conditional
            // add.
            condsingleadders.emplace_back(this->pb,
                accumx[(i+1)/2], accumy[(i+1)/2],
                accumx[(i-1)/2], accumy[(i-1)/2],
                avec[i],
                twoiPx, twoiPy);
        } else {
            conddoubleadders.emplace_back(this->pb,
                accumx[(i+1)/2], accumy[(i+1)/2],
                accumx[(i-1)/2], accumy[(i-1)/2],
                avec[i], avec[i+1],
                twoiPx, twoiPy, twoi1Px, twoi1Py, twoi3Px, twoi3Py);
        }

        ec_double_point(twoiPx, twoiPy, twoi1Px, twoi1Py);
        i += 2;
    }

    // If numbits is even, the output so far is in accum[(numbits)/2].
    // If numbits is odd, it is in accum[(numbits-1)/2].  So in either
    // case, it is in accum[numbits/2].
    singleadders.emplace_back(this->pb,
        outx, outy, accumx[numbits/2], accumy[numbits/2],
        Cx, -Cy);
  }

  void generate_r1cs_constraints()
  {
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(Cx + (CPx-Cx) * avec[0], 1, accumx[0]));
    this->pb.add_r1cs_constraint(r1cs_constraint<FieldT>(Cy + (CPy-Cy) * avec[0], 1, accumy[0]));

    for (auto&& gadget : conddoubleadders) {
        gadget.generate_r1cs_constraints();
    }
    for (auto&& gadget : condsingleadders) {
        gadget.generate_r1cs_constraints();
    }
    for (auto&& gadget : singleadders) {
        gadget.generate_r1cs_constraints();
    }
  }

  void generate_r1cs_witness()
  {
    this->pb.val(accumx[0]) = Cx + (CPx-Cx) * this->pb.val(avec[0]);
    this->pb.val(accumy[0]) = Cy + (CPy-Cy) * this->pb.val(avec[0]);

    for (auto&& gadget : conddoubleadders) {
        gadget.generate_r1cs_witness();
    }
    for (auto&& gadget : condsingleadders) {
        gadget.generate_r1cs_witness();
    }
    for (auto&& gadget : singleadders) {
        gadget.generate_r1cs_witness();
    }
  }
};

// Compute a*P as (outx, outy) for a given constant point P, given a
// as a field element.
template<typename FieldT>
class ec_constant_scalarmul_gadget : public gadget<FieldT> {
private:
  pb_variable_array<FieldT> avec;
  std::vector<packing_gadget<FieldT> > packers;
  std::vector<ec_constant_scalarmul_vec_gadget<FieldT> > vecgadget;

public:
  const pb_variable<FieldT> outx, outy, a;
  const FieldT Px, Py;

  ec_constant_scalarmul_gadget(protoboard<FieldT> &pb,
              const pb_variable<FieldT> &outx,
              const pb_variable<FieldT> &outy,
              const pb_variable<FieldT> &a,
              const FieldT &Px, const FieldT &Py) :
    gadget<FieldT>(pb, "ec_constant_scalarmul_gadget"),
    outx(outx), outy(outy), a(a), Px(Px), Py(Py)
  {
    // Allocate variables to protoboard
    // The strings (like "x") are only for debugging purposes
	  
    size_t numbits = FieldT::num_bits;
    avec.allocate(this->pb, numbits, "avec");
    packers.emplace_back(this->pb, avec, a);
    vecgadget.emplace_back(this->pb, outx, outy, avec, Px, Py);
  }

  void generate_r1cs_constraints()
  {
    packers[0].generate_r1cs_constraints(true);
    vecgadget[0].generate_r1cs_constraints();
  }

  void generate_r1cs_witness()
  {
    packers[0].generate_r1cs_witness_from_packed();
    vecgadget[0].generate_r1cs_witness();
  }
};

// Compute a*G + b*H as (outx, outy), given a and b as field elements.
template<typename FieldT>
class ec_pedersen_gadget : public gadget<FieldT> {
private:
  pb_variable<FieldT> aoutx, aouty, boutx, bouty;
  std::vector<ec_constant_scalarmul_gadget<FieldT> > mulgadgets;
  std::vector<ec_add_gadget<FieldT> > addgadget;
  const FieldT Gx, Gy, Hx, Hy;

public:
  const pb_variable<FieldT> outx, outy, a, b;

  ec_pedersen_gadget(protoboard<FieldT> &pb,
              const pb_variable<FieldT> &outx,
              const pb_variable<FieldT> &outy,
              const pb_variable<FieldT> &a,
              const pb_variable<FieldT> &b) :
    gadget<FieldT>(pb, "ec_pedersen_gadget"),
    outx(outx), outy(outy), a(a), b(b),
  // Precomputed coordinates of G and H
  Gx(0),
  Gy("11977228949870389393715360594190192321220966033310912010610740966317727761886"),
  Hx(1),
  Hy("21803877843449984883423225223478944275188924769286999517937427649571474907279")
  {
    // Allocate variables to protoboard
    // The strings (like "x") are only for debugging purposes
	  
    aoutx.allocate(this->pb, "aoutx");
    aouty.allocate(this->pb, "aouty");
    boutx.allocate(this->pb, "boutx");
    bouty.allocate(this->pb, "bouty");
    mulgadgets.emplace_back(this->pb, aoutx, aouty, a, Gx, Gy);
    mulgadgets.emplace_back(this->pb, boutx, bouty, b, Hx, Hy);
    addgadget.emplace_back(this->pb, outx, outy, aoutx, aouty, boutx, bouty);
  }

  void generate_r1cs_constraints()
  {
    mulgadgets[0].generate_r1cs_constraints();
    mulgadgets[1].generate_r1cs_constraints();
    addgadget[0].generate_r1cs_constraints();
  }

  void generate_r1cs_witness()
  {
    mulgadgets[0].generate_r1cs_witness();
    mulgadgets[1].generate_r1cs_witness();
    addgadget[0].generate_r1cs_witness();
  }
};