/* Copyright (C) 2012-2020 IBM Corp.
 * This program is Licensed under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance
 * with the License. You may obtain a copy of the License at
 *   http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License. See accompanying LICENSE file.
 */
#include <NTL/BasicThreadPool.h>

#include <helib/recryption.h>
#include <helib/EncryptedArray.h>
#include <helib/EvalMap.h>
#include <helib/powerful.h>
#include <helib/CtPtrs.h>
#include <helib/intraSlot.h>
#include <helib/norms.h>
#include <helib/sample.h>
#include <helib/debugging.h>
#include <helib/fhe_stats.h>
#include <helib/log.h>

#ifdef HELIB_DEBUG

#include <helib/debugging.h>

namespace helib {

long printFlag = FLAG_PRINT_VEC;
/************************ Some local functions ***********************/
/*********************************************************************/

static void checkCriticalValue(const std::vector<NTL::ZZX>& zzParts,
                               const DoubleCRT& sKey,
                               const RecryptData& rcData,
                               long q);

static void checkRecryptBounds(const std::vector<NTL::ZZX>& zzParts,
                               const DoubleCRT& sKey,
                               const Context& context,
                               long q);

static void checkRecryptBounds_v(const std::vector<NTL::ZZX>& v,
                                 const DoubleCRT& sKey,
                                 const Context& context,
                                 long q);
} // namespace helib

#endif // HELIB_DEBUG

namespace helib {

// Return in poly a polynomial with X^i encoded in all the slots
static void x2iInSlots(NTL::ZZX& poly,
                       long i,
                       std::vector<NTL::ZZX>& xVec,
                       const EncryptedArray& ea)
{
  xVec.resize(ea.size());
  NTL::ZZX x2i = NTL::ZZX(i, 1);
  for (long j = 0; j < (long)xVec.size(); j++)
    xVec[j] = x2i;
  ea.encode(poly, xVec);
}

// Make every entry of vec divisible by p2e by adding/subtracting q, while
// keeping the added multiples small.  Specifically, for q = 1 mod p2e and any
// integer z can be made divisible by p2e via z' = z + v*q, with |v| <= p2e/2.

static void newMakeDivisible(NTL::ZZX& poly,
                             long p2e,
                             long q,
                             const Context& context,
                             NTL::ZZX& vpoly)
{
  if (p2e == 1) {
    vpoly = 0;
    return;
  }

  assertTrue<InvalidArgument>(q > 0l, "q must be positive");
  assertTrue<InvalidArgument>(p2e > 0l, "p2e must be positive");

  assertEq<InvalidArgument>(q % p2e, 1l, "q must equal 1 modulo p2e");

  long p = context.getP();

  const RecryptData& rcData = context.getRcData();
  const PowerfulDCRT& p2d_conv = *rcData.p2dConv;

  NTL::Vec<NTL::ZZ> pwrfl;
  p2d_conv.ZZXtoPowerful(pwrfl, poly);

#ifdef HELIB_DEBUG
  NTL::Vec<NTL::ZZ> vvec(NTL::INIT_SIZE, pwrfl.length());
#endif

  for (long i : range(pwrfl.length())) {
    NTL::ZZ& z = pwrfl[i];
    long v;

    // What to add to z to make it divisible by p2e?
    long zMod = rem(z, p2e); // zMod is in [0,p2e-1]
    // NOTE: this makes sure we get a truly balanced remainder
    if (zMod > p2e / 2 || (p == 2 && zMod == p2e / 2 && NTL::RandomBnd(2))) {
      // randomize so that v has expected value 0
      zMod = p2e - zMod;
    } else {
      // need to add a negative number
      zMod = -zMod;
    }
    v = zMod;
    z += NTL::to_ZZ(q) * v; // make z divisible by p2e

    if (rem(z, p2e) != 0) { // sanity check
      std::cerr << "**error: original z[" << i
                << "]=" << (z - (NTL::to_ZZ(q) * v)) << std::dec
                << ", p^e=" << p2e << std::endl;
      std::cerr << "z' = z + " << v << "*q = " << z << std::endl;
      exit(1);
    }

#ifdef HELIB_DEBUG
    vvec[i] = v;
#endif
  }

  p2d_conv.powerfulToZZX(poly, pwrfl);

#ifdef HELIB_DEBUG
  p2d_conv.powerfulToZZX(vpoly, vvec);
#endif
}

/*********************************************************************/
/*********************************************************************/

/**
 * Summary of Appendix A from https://ia.cr/2014/873 (version from 2019):
 * Assume that we already chosen e, e' and t.
 *
 * Based in this analysis, we need
 *    (1) (f*p^{e'} + 2*p^r+2))*B <= p^e/2
 * where B is a certain high-probability bound and f is a certain
 * fudge factor.
 *
 **/

// the routine compute_fudge is used to correct for the fact that
// the v-coeffs are not quite uniform

static double compute_fudge(long p2ePrime, long p2e)
{
  double eps = 0;

  if (p2ePrime > 1) {

    if (p2ePrime % 2 == 0) {
      eps = 1 / fsquare(p2ePrime);

      // The exact variance in this case is at most the variance
      // of a random variable that is distributed over
      //    -N..+N
      // where N = 2^{e'}/2.
      // Each endpoint occurs with probability 1/(4*N),
      // and the remaining values each occur with the same probability
      // 1/(2*N)

      // This variance is exactly computed as
      //    (N^2)/3 + 1/6 = ((N^2)/3)*(1 + 1/(2*N^2)), where N = 2^{e'}/2
      // So the std dev is at most
      //    N/sqrt(3)*(1 + 1/(4*N^2))

    } else {
      eps = 1 / double(p2e);

      // We are computing X + Y mod p^{e'}, where
      // X and Y are independent.
      // Y is uniformly distributed over
      //    -floor(p^{r}/2)..floor(p^{r}/2)
      // X is distributed over
      //    -floor(p^e/2)-1..floor(p^e/2)+1,
      // where each endpoint occurs with probability 1 / (2*(p^e+1)),
      // and the remaining p^e values are equally likely

      // The variance in this case is bounded by
      //   (N^2)/3*(1-eps) + (N^2)*eps = (N^2)/3*(1+2*eps),
      //       where N = p^{e'}/2 and eps < 1/p^e
      // So the std dev is bounded by
      //    N/sqrt(3)*sqrt(1+2*eps) <= N/sqrt(3)*(1+eps)
    }
  }

  return 1 + eps;
}

void RecryptData::setAE(long& e, long& ePrime, const Context& context)
{
  double coeff_bound = context.boundForRecryption();
  // coeff_bound is ultimately a high prob bound on |w0+w1*s|,
  // the coeffs of w0, w1 are chosen uniformly on [-1/2,1/2]

  long p = context.getP();
  long p2r = context.getAlMod().getPPowR();
  long r = context.getAlMod().getR();
  long frstTerm = 2 * p2r + 2;

  long e_bnd = 0;
  long p2e_bnd = 1;
  while (p2e_bnd <= ((1L << 30) - 2) / p) { // NOTE: this avoids overflow
    e_bnd++;
    p2e_bnd *= p;
  }
  // e_bnd is largest e such that p^e+1 < 2^30

  // Start with the smallest e s.t. p^e/2 >= frstTerm*coeff_bound
  ePrime = 0;
  e = r + 1;
  while (e <= e_bnd && NTL::power_long(p, e) < frstTerm * coeff_bound * 2)
    e++;

  //  if (e > e_bnd) Error("setAE: cannot find suitable e");
  assertFalse<RuntimeError>(e > e_bnd, "setAE: cannot find suitable e");

  // long ePrimeTry = r+1;
  long ePrimeTry = 1;

  while (ePrimeTry <= e_bnd) {
    long p2ePrimeTry = NTL::power_long(p, ePrimeTry);
    // long eTry = ePrimeTry+1;
    long eTry = std::max(r + 1, ePrimeTry + 1);
    while (eTry <= e_bnd && eTry - ePrimeTry < e - ePrime) {
      long p2eTry = NTL::power_long(p, eTry);
      double fudge = compute_fudge(p2ePrimeTry, p2eTry);
      if (p2eTry >= (p2ePrimeTry * fudge + frstTerm) * coeff_bound * 2)
        break;

      eTry++;
    }

    if (eTry <= e_bnd && eTry - ePrimeTry < e - ePrime) {
      e = eTry;
      ePrime = ePrimeTry;
    }

    ePrimeTry++;
  }

#ifdef HELIB_DEBUG
  std::cerr << "RecryptData::setAE(): e=" << e << ", e'=" << ePrime
            << std::endl;
#endif
}

bool RecryptData::operator==(const RecryptData& other) const
{
  if (mvec != other.mvec)
    return false;

  if (skHwt != other.skHwt)
    return false;

  return true;
}

// The main method
void RecryptData::init(const Context& context,
                       const NTL::Vec<long>& mvec_,
                       bool enableThick,
                       bool build_cache_,
                       bool minimal)
{
  if (alMod != nullptr) { // were we called for a second time?
    std::cerr << "@Warning: multiple calls to RecryptData::init\n";
    return;
  }

  // sanity check
  assertEq(computeProd(mvec_),
           context.getM(),
           "Cyclotomic polynomial mismatch");

  // Record the arguments to this function
  mvec = mvec_;
  build_cache = build_cache_;
  alsoThick = enableThick;

  bool mvec_ok = true;
  for (long i : range(mvec.length())) {
    NTL::Vec<NTL::Pair<long, long>> factors;
    factorize(factors, mvec[i]);
    if (factors.length() > 1)
      mvec_ok = false;
  }

  if (!mvec_ok) {
    Warning("prime power factorization recommended for bootstrapping");
  }

  skHwt = context.getHwt();
  e = context.getE();
  ePrime = context.getEPrime();

  long r = context.getAlMod().getR();

  // First part of Bootstrapping works wrt plaintext space p^{r'}
  alMod = std::make_shared<PAlgebraMod>(context.getZMStar(), e - ePrime + r);
  ea = std::make_shared<EncryptedArray>(context, *alMod);
  // Polynomial defaults to F0, PAlgebraMod explicitly given

  p2dConv = std::make_shared<PowerfulDCRT>(context, mvec);

  if (!enableThick)
    return;

  // Initialize the linear polynomial for unpacking the slots
  NTL::zz_pBak bak;
  bak.save();
  ea->getAlMod().restoreContext();
  long nslots = ea->size();
  long d = ea->getDegree();

  const NTL::Mat<NTL::zz_p>& CBi =
      ea->getDerived(PA_zz_p()).getNormalBasisMatrixInverse();

  std::vector<NTL::ZZX> LM;
  LM.resize(d);
  for (long i = 0; i < d; i++) // prepare the linear polynomial
    LM[i] = rep(CBi[i][0]);

  std::vector<NTL::ZZX> C;
  ea->buildLinPolyCoeffs(C, LM); // "build" the linear polynomial

  unpackSlotEncoding.resize(d); // encode the coefficients

  for (long j = 0; j < d; j++) {
    std::vector<NTL::ZZX> v(nslots);
    for (long k = 0; k < nslots; k++)
      v[k] = C[j];
    ea->encode(unpackSlotEncoding[j], v);
  }
  firstMap = std::make_shared<EvalMap>(*ea, minimal, mvec, true, build_cache);
  secondMap = std::make_shared<EvalMap>(context.getEA(),
                                        minimal,
                                        mvec,
                                        false,
                                        build_cache);
}

/********************************************************************/
/********************************************************************/

// Extract digits from fully packed slots
void extractDigitsPacked(Ctxt& ctxt,
                         long botHigh,
                         long r,
                         long ePrime,
                         const std::vector<NTL::ZZX>& unpackSlotEncoding);

// Extract digits from unpacked slots
void extractDigitsThin(Ctxt& ctxt, long botHigh, long r, long ePrime);

// bootstrap a ciphertext to reduce noise
void PubKey::reCrypt(Ctxt& ctxt) const
{
  HELIB_TIMER_START;

  // Some sanity checks for dummy ciphertext
  long ptxtSpace = ctxt.getPtxtSpace();
  if (ctxt.isEmpty())
    return;
  if (ctxt.parts.size() == 1 && ctxt.parts[0].skHandle.isOne()) {
    // Dummy encryption, just ensure that it is reduced mod p
    NTL::ZZX poly = to_ZZX(ctxt.parts[0]);
    for (long i = 0; i < poly.rep.length(); i++)
      poly[i] = NTL::to_ZZ(rem(poly[i], ptxtSpace));
    poly.normalize();
    ctxt.DummyEncrypt(poly);
    return;
  }

  // check that we have bootstrapping data
  assertTrue(recryptKeyID >= 0l, "No bootstrapping data");

  long p = getContext().getP();
  long r = getContext().getAlMod().getR();
  long p2r = getContext().getAlMod().getPPowR();

  long intFactor = ctxt.intFactor;

  // the bootstrapping key is encrypted relative to plaintext space p^{e-e'+r}.
  const RecryptData& rcData = getContext().getRcData();
  long e = rcData.e;
  long ePrime = rcData.ePrime;
  long p2ePrime = NTL::power_long(p, ePrime);
  long q = NTL::power_long(p, e) + 1;
  assertTrue(e >= r, "rcData.e must be at least alMod.r");

#ifdef HELIB_DEBUG
  std::cerr << "reCrypt: p=" << p << ", r=" << r << ", e=" << e
            << " ePrime=" << ePrime << ", q=" << q << std::endl;
  CheckCtxt(ctxt, "init");
#endif

  // can only bootstrap ciphertext with plaintext-space dividing p^r
  assertEq(p2r % ptxtSpace, 0l, "ptxtSpace must divide p^r when bootstrapping");

  ctxt.dropSmallAndSpecialPrimes();

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after mod down");
#endif

  HELIB_NTIMER_START(AAA_preProcess);

  // Make sure that this ciphertext is in canonical form
  if (!ctxt.inCanonicalForm())
    ctxt.reLinearize();

  // Mod-switch down if needed
  IndexSet s = ctxt.getPrimeSet() / context.getSpecialPrimes();
  assertTrue(s <= context.getCtxtPrimes(), "prime set is messed up");
  if (s.card() > 3) { // leave only first three ciphertext primes
    long first = s.first();
    IndexSet s3(first, first + 2);
    s.retain(s3);
  }
  ctxt.modDownToSet(s);

  // key-switch to the bootstrapping key
  ctxt.reLinearize(recryptKeyID);

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after key switching");
#endif

  // "raw mod-switch" to the bootstrapping modulus q=p^e+1.
  std::vector<NTL::ZZX> zzParts; // the mod-switched parts, in ZZX format

  double mfac = ctxt.getContext().getZMStar().getNormBnd();
  double noise_est = ctxt.rawModSwitch(zzParts, q) * mfac;
  // noise_est is an upper bound on the L-infty norm of the scaled noise
  // in the pwrfl basis
  double noise_bnd =
      HELIB_MIN_CAP_FRAC * p2r * ctxt.getContext().boundForRecryption();
  // noise_bnd is the bound assumed in selecting the parameters
  double noise_rat = noise_est / noise_bnd;

  HELIB_STATS_UPDATE("raw-mod-switch-noise", noise_rat);

  if (noise_rat > 1) {
    // TODO: Turn the following preprocessor logics into a warnOrThrow function
    std::string message =
        "rawModSwitch scaled noise exceeds bound: " + std::to_string(noise_rat);
#ifdef HELIB_DEBUG
    Warning(message);
#else
    throw LogicError(message);
#endif
  }

  assertEq(zzParts.size(),
           (std::size_t)2,
           "Exactly 2 parts required for mod-switching in thin bootstrapping");

#ifdef HELIB_DEBUG
  if (dbgKey) {
    checkRecryptBounds(zzParts, dbgKey->getRecryptKey(), ctxt.getContext(), q);
  }
#endif

  std::vector<NTL::ZZX> v;
  v.resize(2);

  // Add multiples of q to make the zzParts divisible by p^{e'}
  for (long i : range(2)) {
    // make divisible by p^{e'}

    newMakeDivisible(zzParts[i], p2ePrime, q, ctxt.getContext(), v[i]);
  }

#ifdef HELIB_DEBUG
  if (dbgKey) {
    checkRecryptBounds_v(v, dbgKey->getRecryptKey(), ctxt.getContext(), q);
    checkCriticalValue(zzParts,
                       dbgKey->getRecryptKey(),
                       ctxt.getContext().getRcData(),
                       q);
  }
#endif

  for (long i : range(zzParts.size())) {
    zzParts[i] /= p2ePrime; // divide by p^{e'}
  }

  // NOTE: here we lose the intFactor associated with ctxt.
  // We will restore it below.
  ctxt = recryptEkey;

  ctxt.multByConstant(zzParts[1]);
  ctxt.addConstant(zzParts[0]);

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after preProcess");
#endif
  HELIB_NTIMER_STOP(AAA_preProcess);

  // Move the powerful-basis coefficients to the plaintext slots
  HELIB_NTIMER_START(AAA_LinearTransform1);
  ctxt.getContext().getRcData().firstMap->apply(ctxt);
  HELIB_NTIMER_STOP(AAA_LinearTransform1);

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after LinearTransform1");
#endif

  // Extract the digits e-e'+r-1,...,e-e' (from fully packed slots)
  HELIB_NTIMER_START(AAA_extractDigitsPacked);
  extractDigitsPacked(ctxt,
                      e - ePrime,
                      r,
                      ePrime,
                      context.getRcData().unpackSlotEncoding);
  HELIB_NTIMER_STOP(AAA_extractDigitsPacked);

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after extractDigitsPacked");
#endif

  // Move the slots back to powerful-basis coefficients
  HELIB_NTIMER_START(AAA_LinearTransform2);
  ctxt.getContext().getRcData().secondMap->apply(ctxt);
  HELIB_NTIMER_STOP(AAA_LinearTransform2);

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after linearTransform2");
#endif

  // restore intFactor
  if (intFactor != 1)
    ctxt.intFactor = NTL::MulMod(ctxt.intFactor, intFactor, ptxtSpace);
}

#ifdef HELIB_BOOT_THREADS

// Extract digits from fully packed slots, multithreaded version
void extractDigitsPacked(Ctxt& ctxt,
                         long botHigh,
                         long r,
                         long ePrime,
                         const std::vector<NTL::ZZX>& unpackSlotEncoding)
{
  HELIB_TIMER_START;

  // Step 1: unpack the slots of ctxt
  HELIB_NTIMER_START(unpack);
  ctxt.cleanUp();

  // Apply the d automorphisms and store them in scratch area
  long d = ctxt.getContext().getOrdP();

  std::vector<Ctxt> unpacked(d, Ctxt(ZeroCtxtLike, ctxt));
  { // explicit scope to force all temporaries to be released
    std::vector<std::shared_ptr<DoubleCRT>> coeff_vector;
    std::vector<double> coeff_vector_sz;
    coeff_vector.resize(d);
    coeff_vector_sz.resize(d);

    HELIB_NTIMER_START(unpack1);
    for (long i = 0; i < d; i++) {
      coeff_vector[i] = std::make_shared<DoubleCRT>(unpackSlotEncoding[i],
                                                    ctxt.getContext(),
                                                    ctxt.getPrimeSet());
      coeff_vector_sz[i] = NTL::conv<double>(
          embeddingLargestCoeff(unpackSlotEncoding[i],
                                ctxt.getContext().getZMStar()));
    }
    HELIB_NTIMER_STOP(unpack1);

    HELIB_NTIMER_START(unpack2);
    std::vector<Ctxt> frob(d, Ctxt(ZeroCtxtLike, ctxt));

    NTL_EXEC_RANGE(d, first, last)
    // FIXME: implement using hoisting!
    for (long j = first; j < last; j++) { // process jth Frobenius
      frob[j] = ctxt;
      frob[j].frobeniusAutomorph(j);
      frob[j].cleanUp();
      // FIXME: not clear if we should call cleanUp here
    }
    NTL_EXEC_RANGE_END

    HELIB_NTIMER_STOP(unpack2);

    HELIB_NTIMER_START(unpack3);
    Ctxt tmp1(ZeroCtxtLike, ctxt);
    for (long i = 0; i < d; i++) {
      for (long j = 0; j < d; j++) {
        tmp1 = frob[j];
        tmp1.multByConstant(*coeff_vector[mcMod(i + j, d)],
                            coeff_vector_sz[mcMod(i + j, d)]);
        unpacked[i] += tmp1;
      }
    }
    HELIB_NTIMER_STOP(unpack3);
  }
  HELIB_NTIMER_STOP(unpack);

  //#ifdef HELIB_DEBUG
  //  CheckCtxt(unpacked[0], "after unpack");
  //#endif

  NTL_EXEC_RANGE(d, first, last)
  for (long i = first; i < last; i++) {
    extractDigitsThin(unpacked[i], botHigh, r, ePrime);
  }
  NTL_EXEC_RANGE_END

  //#ifdef HELIB_DEBUG
  // CheckCtxt(unpacked[0], "before repack");
  //#endif

  // Step 3: re-pack the slots
  HELIB_NTIMER_START(repack);
  const EncryptedArray& ea2 = ctxt.getContext().getEA();
  NTL::ZZX xInSlots;
  std::vector<NTL::ZZX> xVec(ea2.size());
  ctxt = unpacked[0];
  for (long i = 1; i < d; i++) {
    x2iInSlots(xInSlots, i, xVec, ea2);
    unpacked[i].multByConstant(xInSlots);
    ctxt += unpacked[i];
  }
  HELIB_NTIMER_STOP(repack);
  //#ifdef HELIB_DEBUG
  // CheckCtxt(ctxt, "after repack");
  //#endif
}

#else

// Extract digits from fully packed slots
void extractDigitsPacked(Ctxt& ctxt,
                         long botHigh,
                         long r,
                         long ePrime,
                         const std::vector<NTL::ZZX>& unpackSlotEncoding)
{
  HELIB_TIMER_START;

  // Step 1: unpack the slots of ctxt
  HELIB_NTIMER_START(unpack);
  ctxt.cleanUp();

  // Apply the d automorphisms and store them in scratch area
  long d = ctxt.getContext().getOrdP();

  std::vector<Ctxt> unpacked(d, Ctxt(ZeroCtxtLike, ctxt));
  { // explicit scope to force all temporaries to be released
    std::vector<std::shared_ptr<DoubleCRT>> coeff_vector;
    std::vector<double> coeff_vector_sz;
    coeff_vector.resize(d);
    coeff_vector_sz.resize(d);
    for (long i = 0; i < d; i++) {
      coeff_vector[i] = std::make_shared<DoubleCRT>(unpackSlotEncoding[i],
                                                    ctxt.getContext(),
                                                    ctxt.getPrimeSet());
      coeff_vector_sz[i] = NTL::conv<double>(
          embeddingLargestCoeff(unpackSlotEncoding[i],
                                ctxt.getContext().getZMStar()));
    }

    Ctxt tmp1(ZeroCtxtLike, ctxt);
    Ctxt tmp2(ZeroCtxtLike, ctxt);

    // FIXME: implement using hoisting!
    for (long j = 0; j < d; j++) { // process jth Frobenius
      tmp1 = ctxt;
      tmp1.frobeniusAutomorph(j);
      tmp1.cleanUp();
      // FIXME: not clear if we should call cleanUp here

      for (long i = 0; i < d; i++) {
        tmp2 = tmp1;
        tmp2.multByConstant(*coeff_vector[mcMod(i + j, d)],
                            coeff_vector_sz[mcMod(i + j, d)]);
        unpacked[i] += tmp2;
      }
    }
  }
  HELIB_NTIMER_STOP(unpack);

  //#ifdef HELIB_DEBUG
  //  CheckCtxt(unpacked[0], "after unpack");
  //#endif

  for (long i = 0; i < (long)unpacked.size(); i++) {
    extractDigitsThin(unpacked[i], botHigh, r, ePrime);
  }

  //#ifdef HELIB_DEBUG
  //  CheckCtxt(unpacked[0], "before repack");
  //#endif

  // Step 3: re-pack the slots
  HELIB_NTIMER_START(repack);
  const EncryptedArray& ea2 = *ctxt.getContext().ea;
  NTL::ZZX xInSlots;
  std::vector<NTL::ZZX> xVec(ea2.size());
  ctxt = unpacked[0];
  for (long i = 1; i < d; i++) {
    x2iInSlots(xInSlots, i, xVec, ea2);
    unpacked[i].multByConstant(xInSlots);
    ctxt += unpacked[i];
  }
  HELIB_NTIMER_STOP(repack);
}

#endif

// Use packed bootstrapping, so we can bootstrap all in just one go.
void packedRecrypt(const CtPtrs& cPtrs,
                   const std::vector<zzX>& unpackConsts,
                   const EncryptedArray& ea)
{
  PubKey& pKey = (PubKey&)cPtrs[0]->getPubKey();

  // Allocate temporary ciphertexts for the recryption
  int nPacked = divc(cPtrs.size(), ea.getDegree()); // ceil(totalNum/d)
  std::vector<Ctxt> cts(nPacked, Ctxt(pKey));

  repack(CtPtrs_vectorCt(cts), cPtrs, ea); // pack ciphertexts
  //  cout << "@"<< lsize(cts)<<std::flush;
  for (Ctxt& c : cts) {   // then recrypt them
    c.reducePtxtSpace(2); // we only have recryption data for binary ctxt
    pKey.reCrypt(c);
  }
  unpack(cPtrs, CtPtrs_vectorCt(cts), ea, unpackConsts);
}

// recrypt all ctxt at level < belowLvl
void packedRecrypt(const CtPtrs& array,
                   const std::vector<zzX>& unpackConsts,
                   const EncryptedArray& ea,
                   long belowLvl)
{
  std::vector<Ctxt*> v;
  for (long i = 0; i < array.size(); i++)
    if (array.isSet(i) && !array[i]->isEmpty() &&
        array[i]->bitCapacity() < belowLvl * (array[i]->getContext().BPL()))
      v.push_back(array[i]);
  packedRecrypt(CtPtrs_vectorPt(v), unpackConsts, ea);
}
void packedRecrypt(const CtPtrMat& m,
                   const std::vector<zzX>& unpackConsts,
                   const EncryptedArray& ea,
                   long belowLvl)
{
  std::vector<Ctxt*> v;
  for (long i = 0; i < m.size(); i++)
    for (long j = 0; j < m[i].size(); j++)
      if (m[i].isSet(j) && !m[i][j]->isEmpty() &&
          m[i][j]->bitCapacity() < belowLvl * (m[i][j]->getContext().BPL()))
        v.push_back(m[i][j]);
  packedRecrypt(CtPtrs_vectorPt(v), unpackConsts, ea);
}

//===================== Thin Bootstrapping stuff ==================

void ThinRecryptData::init(const Context& context,
                           const NTL::Vec<long>& mvec_,
                           bool alsoThick,
                           bool build_cache_,
                           bool minimal)
{
  RecryptData::init(context, mvec_, alsoThick, build_cache_, minimal);
  coeffToSlot =
      std::make_shared<ThinEvalMap>(*ea, minimal, mvec, true, build_cache);
  slotToCoeff = std::make_shared<ThinEvalMap>(context.getEA(),
                                              minimal,
                                              mvec,
                                              false,
                                              build_cache);
}

// Extract digits from thinly packed slots

long fhe_force_chen_han = 0;

void extractDigitsThin(Ctxt& ctxt, long botHigh, long r, long ePrime)
{
  HELIB_TIMER_START;

  Ctxt unpacked(ctxt);
  unpacked.cleanUp();

  std::vector<Ctxt> scratch;

  long p = ctxt.getContext().getP();
  long p2r = NTL::power_long(p, r);
  long topHigh = botHigh + r - 1;

  // degree Chen/Han technique is p^{bot-1}(p-1)r
  // degree of basic technique is p^{bot-1}p^r,
  //     or p^{bot-1}p^{r-1} if p==2, r > 1, and bot+r > 2

  bool use_chen_han = false;
  if (r > 1) {
    double chen_han_cost = log(p - 1) + log(r);
    double basic_cost;
    if (p == 2 && r > 2 && botHigh + r > 2)
      basic_cost = (r - 1) * log(p);
    else
      basic_cost = r * log(p);

    // std::cerr << "*** basic: " << basic_cost << "\n";
    // std::cerr << "*** chen/han: " << chen_han_cost << "\n";

    double thresh = 1.5;
    if (p == 2)
      thresh = 1.75;
    // increasing thresh makes chen_han less likely to be chosen.
    // For p == 2, the basic algorithm is just squaring,
    // and so is a bit cheaper, so we raise thresh a bit.
    // This is all a bit heuristic.

    if (basic_cost > thresh * chen_han_cost)
      use_chen_han = true;
  }

  if (fhe_force_chen_han > 0)
    use_chen_han = true;
  else if (fhe_force_chen_han < 0)
    use_chen_han = false;

  if (use_chen_han) {
    // use Chen and Han technique

    extendExtractDigits(scratch, unpacked, botHigh, r);

#if 0
    for (long i: range(scratch.size())) {
      CheckCtxt(scratch[i], "**");
    }
#endif

    for (long j = 0; j < botHigh; j++) {
      unpacked -= scratch[j];
      unpacked.divideByP();
    }

    if (p == 2 && botHigh > 0) // For p==2, subtract also the previous bit
      unpacked += scratch[botHigh - 1];
    unpacked.negate();

    if (r > ePrime) { // Add in digits from the bottom part, if any
      long topLow = r - 1 - ePrime;
      Ctxt tmp = scratch[topLow];
      for (long j = topLow - 1; j >= 0; --j) {
        tmp.multByP();
        tmp += scratch[j];
      }
      if (ePrime > 0)
        tmp.multByP(ePrime); // multiply by p^e'
      unpacked += tmp;
    }
    unpacked.reducePtxtSpace(p2r); // Our plaintext space is now mod p^r

    ctxt = unpacked;
  } else {

    if (p == 2 && r > 2 && topHigh + 1 > 2)
      topHigh--; // For p==2 we sometime get a bit for free

    extractDigits(scratch, unpacked, topHigh + 1);

    // set unpacked = -\sum_{j=botHigh}^{topHigh} scratch[j] * p^{j-botHigh}
    if (topHigh >= LONG(scratch.size())) {
      topHigh = scratch.size() - 1;
      std::cerr << " @ suspect: not enough digits in extractDigitsPacked\n";
    }

    unpacked = scratch[topHigh];
    for (long j = topHigh - 1; j >= botHigh; --j) {
      unpacked.multByP();
      unpacked += scratch[j];
    }
    if (p == 2 && botHigh > 0) // For p==2, subtract also the previous bit
      unpacked += scratch[botHigh - 1];
    unpacked.negate();

    if (r > ePrime) { // Add in digits from the bottom part, if any
      long topLow = r - 1 - ePrime;
      Ctxt tmp = scratch[topLow];
      for (long j = topLow - 1; j >= 0; --j) {
        tmp.multByP();
        tmp += scratch[j];
      }
      if (ePrime > 0)
        tmp.multByP(ePrime); // multiply by p^e'
      unpacked += tmp;
    }
    unpacked.reducePtxtSpace(p2r); // Our plaintext space is now mod p^r
    ctxt = unpacked;
  }
}

// Hack to get at private fields of public key
struct PubKeyHack
{                         // The public key
  const Context& context; // The context

  //! @var Ctxt pubEncrKey
  //! The public encryption key is an encryption of 0,
  //! relative to the first secret key
  Ctxt pubEncrKey;

  std::vector<long> skHwts;            // The Hamming weight of the secret keys
  std::vector<KeySwitch> keySwitching; // The key-switching matrices

  // The keySwitchMap structure contains pointers to key-switching matrices
  // for re-linearizing automorphisms. The entry keySwitchMap[i][n] contains
  // the index j such that keySwitching[j] is the first matrix one needs to
  // use when re-linearizing s_i(X^n).
  std::vector<std::vector<long>> keySwitchMap;

  NTL::Vec<int> KS_strategy; // NTL Vec's support I/O, which is
                             // more convenient

  // bootstrapping data

  long recryptKeyID; // index of the bootstrapping key
  Ctxt recryptEkey;  // the key itself, encrypted under key #0
};

// bootstrap a ciphertext to reduce noise
void PubKey::thinReCrypt(Ctxt& ctxt) const
{
  HELIB_TIMER_START;

  // Some sanity checks for dummy ciphertext
  long ptxtSpace = ctxt.getPtxtSpace();
  if (ctxt.isEmpty())
    return;

  if (ctxt.parts.size() == 1 && ctxt.parts[0].skHandle.isOne()) {
    // Dummy encryption, just ensure that it is reduced mod p
    NTL::ZZX poly = to_ZZX(ctxt.parts[0]);
    for (long i = 0; i < poly.rep.length(); i++)
      poly[i] = NTL::to_ZZ(rem(poly[i], ptxtSpace));
    poly.normalize();
    ctxt.DummyEncrypt(poly);
    return;
  }

  // check that we have bootstrapping data
  assertTrue(recryptKeyID >= 0l, "Bootstrapping data not present");

  long p = ctxt.getContext().getP();
  long r = ctxt.getContext().getAlMod().getR();
  long p2r = ctxt.getContext().getAlMod().getPPowR();

  long intFactor = ctxt.intFactor;

  const ThinRecryptData& trcData = ctxt.getContext().getRcData();

  // the bootstrapping key is encrypted relative to plaintext space p^{e-e'+r}.
  long e = trcData.e;
  long ePrime = trcData.ePrime;
  long p2ePrime = NTL::power_long(p, ePrime);
  long q = NTL::power_long(p, e) + 1;
  assertTrue(e >= r, "trcData.e must be at least alMod.r");

  // can only bootstrap ciphertext with plaintext-space dividing p^r
  assertEq(p2r % ptxtSpace,
           0l,
           "ptxtSpace must divide p^r when thin bootstrapping");

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "init");
#endif

  ctxt.dropSmallAndSpecialPrimes();

#define DROP_BEFORE_THIN_RECRYPT
#define THIN_RECRYPT_NLEVELS (3)
#ifdef DROP_BEFORE_THIN_RECRYPT
  // experimental code...we should drop down to a reasonably low level
  // before doing the first linear map.
  long first = context.getCtxtPrimes().first();
  long last = std::min(context.getCtxtPrimes().last(),
                       first + THIN_RECRYPT_NLEVELS - 1);
  ctxt.bringToSet(IndexSet(first, last));
#endif

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after mod down");
#endif

  // Move the slots to powerful-basis coefficients
  HELIB_NTIMER_START(AAA_slotToCoeff);
  trcData.slotToCoeff->apply(ctxt);
  HELIB_NTIMER_STOP(AAA_slotToCoeff);

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after slotToCoeff");
#endif

  HELIB_NTIMER_START(AAA_bootKeySwitch);

  // Make sure that this ciphertext is in canonical form
  if (!ctxt.inCanonicalForm())
    ctxt.reLinearize();

  // Mod-switch down if needed
  IndexSet s = ctxt.getPrimeSet() / context.getSpecialPrimes();
  assertTrue(s <= context.getCtxtPrimes(), "prime set is messed up");
  if (s.card() > 3) { // leave only first three ciphertext primes
    long first = s.first();
    IndexSet s3(first, first + 2);
    s.retain(s3);
  }
  ctxt.modDownToSet(s);

  // key-switch to the bootstrapping key
  ctxt.reLinearize(recryptKeyID);

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after key switching");
#endif

  // "raw mod-switch" to the bootstrapping mosulus q=p^e+1.
  std::vector<NTL::ZZX> zzParts; // the mod-switched parts, in ZZX format

  double mfac = ctxt.getContext().getZMStar().getNormBnd();
  double noise_est = ctxt.rawModSwitch(zzParts, q) * mfac;
  // noise_est is an upper bound on the L-infty norm of the scaled noise
  // in the pwrfl basis
  double noise_bnd =
      HELIB_MIN_CAP_FRAC * p2r * ctxt.getContext().boundForRecryption();
  // noise_bnd is the bound assumed in selecting the parameters
  double noise_rat = noise_est / noise_bnd;

  HELIB_STATS_UPDATE("raw-mod-switch-noise", noise_rat);

  if (noise_rat > 1) {
    // TODO: Turn the following preprocessor logics into a warnOrThrow function
    std::string message =
        "rawModSwitch scaled noise exceeds bound: " + std::to_string(noise_rat);
#ifdef HELIB_DEBUG
    Warning(message);
#else
    throw LogicError(message);
#endif
  }

  assertEq(zzParts.size(),
           (std::size_t)2,
           "Exactly 2 parts required for mod-switching in thin bootstrapping");

#ifdef HELIB_DEBUG
  if (dbgKey) {
    checkRecryptBounds(zzParts, dbgKey->getRecryptKey(), ctxt.getContext(), q);
  }
#endif

  std::vector<NTL::ZZX> v;
  v.resize(2);

  // Add multiples of q to make the zzParts divisible by p^{e'}
  for (long i : range(2)) {
    // make divisible by p^{e'}

    newMakeDivisible(zzParts[i], p2ePrime, q, ctxt.getContext(), v[i]);
  }

#ifdef HELIB_DEBUG
  if (dbgKey) {
    checkRecryptBounds_v(v, dbgKey->getRecryptKey(), ctxt.getContext(), q);
    checkCriticalValue(zzParts,
                       dbgKey->getRecryptKey(),
                       ctxt.getContext().getRcData(),
                       q);
  }
#endif

  for (long i : range(zzParts.size())) {
    zzParts[i] /= p2ePrime; // divide by p^{e'}
  }

  // NOTE: here we lose the intFactor associated with ctxt.
  // We will restore it below.
  ctxt = recryptEkey;

  ctxt.multByConstant(zzParts[1]);
  ctxt.addConstant(zzParts[0]);

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after bootKeySwitch");
#endif

  HELIB_NTIMER_STOP(AAA_bootKeySwitch);

  // Move the powerful-basis coefficients to the plaintext slots
  HELIB_NTIMER_START(AAA_coeffToSlot);
  trcData.coeffToSlot->apply(ctxt);
  HELIB_NTIMER_STOP(AAA_coeffToSlot);

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after coeffToSlot");
#endif

  // Extract the digits e-e'+r-1,...,e-e' (from fully packed slots)
  HELIB_NTIMER_START(AAA_extractDigitsThin);
  extractDigitsThin(ctxt, e - ePrime, r, ePrime);
  HELIB_NTIMER_STOP(AAA_extractDigitsThin);

#ifdef HELIB_DEBUG
  CheckCtxt(ctxt, "after extractDigitsThin");
#endif

  // restore intFactor
  if (intFactor != 1)
    ctxt.intFactor = NTL::MulMod(ctxt.intFactor, intFactor, ptxtSpace);
}

#ifdef HELIB_DEBUG

static void checkCriticalValue(const std::vector<NTL::ZZX>& zzParts,
                               const DoubleCRT& sKey,
                               const RecryptData& rcData,
                               long q)
{
  NTL::ZZX ptxt;
  rawDecrypt(ptxt, zzParts, sKey); // no mod q

  NTL::Vec<NTL::ZZ> powerful;
  rcData.p2dConv->ZZXtoPowerful(powerful, ptxt);
  NTL::xdouble max_pwrfl = NTL::conv<NTL::xdouble>(largestCoeff(powerful));
  double critical_value = NTL::conv<double>((max_pwrfl / q) / q);

  vecRed(powerful, powerful, q, false);
  max_pwrfl = NTL::conv<NTL::xdouble>(largestCoeff(powerful));
  critical_value += NTL::conv<double>(max_pwrfl / q);

  HELIB_STATS_UPDATE("critical-value", critical_value);

  std::cerr << "=== critical_value=" << critical_value;
  if (critical_value > 0.5)
    std::cerr << " BAD-BOUND";

  std::cerr << "\n";
}

static void checkRecryptBounds(const std::vector<NTL::ZZX>& zzParts,
                               const DoubleCRT& sKey,
                               const Context& context,
                               long q)
{
  const RecryptData& rcData = context.getRcData();
  double coeff_bound = context.boundForRecryption();
  long p2r = context.getAlMod().getPPowR();

  NTL::ZZX ptxt;
  rawDecrypt(ptxt, zzParts, sKey); // no mod q

  NTL::Vec<NTL::ZZ> powerful;
  rcData.p2dConv->ZZXtoPowerful(powerful, ptxt);
  double max_pwrfl = NTL::conv<double>(largestCoeff(powerful));
  double ratio = max_pwrfl / (2 * q * coeff_bound);

  HELIB_STATS_UPDATE("|x|/bound", ratio);

  std::cerr << "=== |x|/bound=" << ratio;
  if (ratio > 1.0)
    std::cerr << " BAD-BOUND";

  vecRed(powerful, powerful, q, false);
  max_pwrfl = NTL::conv<double>(largestCoeff(powerful));
  ratio = max_pwrfl / (2 * p2r * coeff_bound);

  HELIB_STATS_UPDATE("|x%q|/bound", ratio);

  std::cerr << ", (|x%q|)/bound=" << ratio;
  if (ratio > 1.0)
    std::cerr << " BAD-BOUND";

  std::cerr << "\n";
}

static void checkRecryptBounds_v(const std::vector<NTL::ZZX>& v,
                                 const DoubleCRT& sKey,
                                 const Context& context,
                                 UNUSED long q)
{
  const RecryptData& rcData = context.getRcData();

  long p = context.getP();
  long e = rcData.e;
  long p2e = NTL::power_long(p, e);
  long ePrime = rcData.ePrime;
  long p2ePrime = NTL::power_long(p, ePrime);
  long phim = context.getPhiM();

  double fudge = compute_fudge(p2ePrime, p2e);

  double coeff_bound = context.boundForRecryption() * fudge;

  double sigma = context.stdDevForRecryption() * fudge;

  NTL::ZZX ptxt;
  rawDecrypt(ptxt, v, sKey); // no mod q

  NTL::Vec<NTL::ZZ> powerful;
  rcData.p2dConv->ZZXtoPowerful(powerful, ptxt);
  double max_pwrfl = NTL::conv<double>(largestCoeff(powerful));

  double denom = p2ePrime * coeff_bound;
  double ratio = max_pwrfl / denom;

  HELIB_STATS_UPDATE("|v|/bound", ratio);

  std::cerr << "=== |v|/bound=" << ratio;
  if (ratio > 1.0)
    std::cerr << " BAD-BOUND";
  std::cerr << "\n";

  ptxt -= v[0]; // so now ptxt is just sKey * v[1]
  rcData.p2dConv->ZZXtoPowerful(powerful, ptxt);

  assertEq(powerful.length(), phim, "length should be phim");

  double ran_pwrfl = NTL::conv<double>(powerful[NTL::RandomBnd(phim)]);
  // pick a random coefficient in the poweful basis

  double std_devs = fabs(ran_pwrfl) / (p2ePrime * sigma);
  // number of standard deviations away from mean

  // update various indicator variables
  HELIB_STATS_UPDATE("sigma_0_5", double(std_devs <= 0.5)); // 0.383
  HELIB_STATS_UPDATE("sigma_1_0", double(std_devs <= 1.0)); // 0.683
  HELIB_STATS_UPDATE("sigma_1_5", double(std_devs <= 1.5)); // 0.866
  HELIB_STATS_UPDATE("sigma_2_0", double(std_devs <= 2.0)); // 0.954
  HELIB_STATS_UPDATE("sigma_2_5", double(std_devs <= 2.5)); // 0.988
  HELIB_STATS_UPDATE("sigma_3_0", double(std_devs <= 3.0)); // 0.997, 1 in 370
  HELIB_STATS_UPDATE("sigma_3_5",
                     double(std_devs <= 3.5)); // 0.999535, 1 in 2149
  HELIB_STATS_UPDATE("sigma_4_0",
                     double(std_devs <= 4.0)); // 0.999937, 1 in 15787

  // compute sample variance, and scale by the variance we expect
  HELIB_STATS_UPDATE("sigma_calc",
                     fsquare(ran_pwrfl) / fsquare(p2ePrime * sigma));

  // save the scaled value for application of other tests
  HELIB_STATS_SAVE("v_values", ran_pwrfl / (p2ePrime * sigma));
}

#endif

#if 0
void fhe_stats_print(long iter, const Context& context)
{
   long phim = context.getPhiM();

   std::cerr << "||||| recryption stats ||||\n";
   std::cerr << "**** averages ****\n";
   std::cerr << "=== critical_value=" << (fhe_stats_cv_sum/iter) << "\n";
   std::cerr << "=== |x|/bound=" << (fhe_stats_x_sum/iter) << "\n";
   std::cerr << "=== |x%q|/bound=" << (fhe_stats_xmod_sum/iter) << "\n";
   std::cerr << "=== |u|/bound=" << (fhe_stats_u_sum/iter) << "\n";
   std::cerr << "=== |v|/bound=" << (fhe_stats_v_sum/iter) << "\n";
   std::cerr << "**** maxima ****\n";
   std::cerr << "=== critical_value=" << (fhe_stats_cv_max) << "\n";
   std::cerr << "=== |x|/bound=" << (fhe_stats_x_max) << "\n";
   std::cerr << "=== |x%q|/bound=" << (fhe_stats_xmod_max) << "\n";
   std::cerr << "=== |u|/bound=" << (fhe_stats_u_max) << "\n";
   std::cerr << "=== |v|/bound=" << (fhe_stats_v_max) << "\n";
   std::cerr << "**** theoretical bounds ***\n";
   std::cerr << "=== single-max=" << (sqrt(2.0*log(phim))/context.scale) << "\n";
   std::cerr << "=== global-max=" << (sqrt(2.0*(log(iter)+log(phim)))/context.scale) << "\n";

}
#endif

} // namespace helib


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2007-2015 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2015 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2015 Mateusz Loskot, London, UK.
// Copyright (c) 2014-2015 Adam Wulkiewicz, Lodz, Poland.

// This file was modified by Oracle on 2014, 2015.
// Modifications copyright (c) 2014-2015 Oracle and/or its affiliates.

// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle
// Contributed and/or modified by Menelaos Karavelas, on behalf of Oracle

// Parts of Boost.Geometry are redesigned from Geodan's Geographic Library
// (geolib/GGL), copyright (c) 1995-2010 Geodan, Amsterdam, the Netherlands.

// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_GEOMETRY_ALGORITHMS_EQUALS_HPP
#define BOOST_GEOMETRY_ALGORITHMS_EQUALS_HPP

#include <cstddef>
#include <vector>

#include <boost/range.hpp>

#include <boost/variant/apply_visitor.hpp>
#include <boost/variant/static_visitor.hpp>
#include <boost/variant/variant_fwd.hpp>

#include <boost/geometry/core/access.hpp>
#include <boost/geometry/core/coordinate_dimension.hpp>
#include <boost/geometry/core/geometry_id.hpp>
#include <boost/geometry/core/reverse_dispatch.hpp>
#include <boost/geometry/core/tags.hpp>

#include <boost/geometry/geometries/concepts/check.hpp>

#include <boost/geometry/algorithms/detail/equals/point_point.hpp>
#include <boost/geometry/algorithms/detail/not.hpp>
#include <boost/geometry/algorithms/not_implemented.hpp>

// For trivial checks
#include <boost/geometry/algorithms/area.hpp>
#include <boost/geometry/algorithms/length.hpp>
#include <boost/geometry/util/math.hpp>
#include <boost/geometry/util/select_coordinate_type.hpp>
#include <boost/geometry/util/select_most_precise.hpp>

#include <boost/geometry/algorithms/detail/equals/collect_vectors.hpp>
#include <boost/geometry/algorithms/relate.hpp>
#include <boost/geometry/algorithms/detail/relate/relate_impl.hpp>

#include <boost/geometry/views/detail/indexed_point_view.hpp>

namespace boost { namespace geometry
{

#ifndef DOXYGEN_NO_DETAIL
namespace detail { namespace equals
{

template
<
    std::size_t Dimension,
    std::size_t DimensionCount
>
struct box_box
{
    template <typename Box1, typename Box2>
    static inline bool apply(Box1 const& box1, Box2 const& box2)
    {
        if (!geometry::math::equals(get<min_corner, Dimension>(box1), get<min_corner, Dimension>(box2))
            || !geometry::math::equals(get<max_corner, Dimension>(box1), get<max_corner, Dimension>(box2)))
        {
            return false;
        }
        return box_box<Dimension + 1, DimensionCount>::apply(box1, box2);
    }
};

template <std::size_t DimensionCount>
struct box_box<DimensionCount, DimensionCount>
{
    template <typename Box1, typename Box2>
    static inline bool apply(Box1 const& , Box2 const& )
    {
        return true;
    }
};

struct segment_segment
{
    template <typename Segment1, typename Segment2>
    static inline bool apply(Segment1 const& segment1, Segment2 const& segment2)
    {
        return equals::equals_point_point(
                    indexed_point_view<Segment1 const, 0>(segment1),
                    indexed_point_view<Segment2 const, 0>(segment2) )
                ? equals::equals_point_point(
                    indexed_point_view<Segment1 const, 1>(segment1),
                    indexed_point_view<Segment2 const, 1>(segment2) )
                : ( equals::equals_point_point(
                        indexed_point_view<Segment1 const, 0>(segment1),
                        indexed_point_view<Segment2 const, 1>(segment2) )
                 && equals::equals_point_point(
                        indexed_point_view<Segment1 const, 1>(segment1),
                        indexed_point_view<Segment2 const, 0>(segment2) )
                  );
    }
};

struct area_check
{
    template <typename Geometry1, typename Geometry2>
    static inline bool apply(Geometry1 const& geometry1, Geometry2 const& geometry2)
    {
        return geometry::math::equals(
                geometry::area(geometry1),
                geometry::area(geometry2));
    }
};

struct length_check
{
    template <typename Geometry1, typename Geometry2>
    static inline bool apply(Geometry1 const& geometry1, Geometry2 const& geometry2)
    {
        return geometry::math::equals(
                geometry::length(geometry1),
                geometry::length(geometry2));
    }
};

template <typename TrivialCheck>
struct equals_by_collection
{
    template <typename Geometry1, typename Geometry2>
    static inline bool apply(Geometry1 const& geometry1, Geometry2 const& geometry2)
    {
        if (! TrivialCheck::apply(geometry1, geometry2))
        {
            return false;
        }

        typedef typename geometry::select_most_precise
            <
                typename select_coordinate_type
                    <
                        Geometry1, Geometry2
                    >::type,
                double
            >::type calculation_type;

        typedef std::vector<collected_vector<calculation_type> > v;
        v c1, c2;

        geometry::collect_vectors(c1, geometry1);
        geometry::collect_vectors(c2, geometry2);

        if (boost::size(c1) != boost::size(c2))
        {
            return false;
        }

        std::sort(c1.begin(), c1.end());
        std::sort(c2.begin(), c2.end());

        // Just check if these vectors are equal.
        return std::equal(c1.begin(), c1.end(), c2.begin());
    }
};

template<typename Geometry1, typename Geometry2>
struct equals_by_relate
    : detail::relate::relate_impl
        <
            detail::de9im::static_mask_equals_type,
            Geometry1,
            Geometry2
        >
{};

}} // namespace detail::equals
#endif // DOXYGEN_NO_DETAIL

#ifndef DOXYGEN_NO_DISPATCH
namespace dispatch
{

template
<
    typename Geometry1,
    typename Geometry2,
    typename Tag1 = typename tag<Geometry1>::type,
    typename Tag2 = typename tag<Geometry2>::type,
    std::size_t DimensionCount = dimension<Geometry1>::type::value,
    bool Reverse = reverse_dispatch<Geometry1, Geometry2>::type::value
>
struct equals: not_implemented<Tag1, Tag2>
{};

// If reversal is needed, perform it
template
<
    typename Geometry1, typename Geometry2,
    typename Tag1, typename Tag2,
    std::size_t DimensionCount
>
struct equals<Geometry1, Geometry2, Tag1, Tag2, DimensionCount, true>
    : equals<Geometry2, Geometry1, Tag2, Tag1, DimensionCount, false>
{
    static inline bool apply(Geometry1 const& g1, Geometry2 const& g2)
    {
        return equals
            <
                Geometry2, Geometry1,
                Tag2, Tag1,
                DimensionCount,
                false
            >::apply(g2, g1);
    }
};

template <typename P1, typename P2, std::size_t DimensionCount, bool Reverse>
struct equals<P1, P2, point_tag, point_tag, DimensionCount, Reverse>
    : geometry::detail::not_
        <
            detail::disjoint::point_point<P1, P2, 0, DimensionCount>
        >
{};

template <typename Box1, typename Box2, std::size_t DimensionCount, bool Reverse>
struct equals<Box1, Box2, box_tag, box_tag, DimensionCount, Reverse>
    : detail::equals::box_box<0, DimensionCount>
{};

template <typename Ring1, typename Ring2, bool Reverse>
struct equals<Ring1, Ring2, ring_tag, ring_tag, 2, Reverse>
    : detail::equals::equals_by_collection<detail::equals::area_check>
{};

template <typename Polygon1, typename Polygon2, bool Reverse>
struct equals<Polygon1, Polygon2, polygon_tag, polygon_tag, 2, Reverse>
    : detail::equals::equals_by_collection<detail::equals::area_check>
{};

template <typename Polygon, typename Ring, bool Reverse>
struct equals<Polygon, Ring, polygon_tag, ring_tag, 2, Reverse>
    : detail::equals::equals_by_collection<detail::equals::area_check>
{};

template <typename Ring, typename Box, bool Reverse>
struct equals<Ring, Box, ring_tag, box_tag, 2, Reverse>
    : detail::equals::equals_by_collection<detail::equals::area_check>
{};

template <typename Polygon, typename Box, bool Reverse>
struct equals<Polygon, Box, polygon_tag, box_tag, 2, Reverse>
    : detail::equals::equals_by_collection<detail::equals::area_check>
{};

template <typename Segment1, typename Segment2, std::size_t DimensionCount, bool Reverse>
struct equals<Segment1, Segment2, segment_tag, segment_tag, DimensionCount, Reverse>
    : detail::equals::segment_segment
{};

template <typename LineString1, typename LineString2, bool Reverse>
struct equals<LineString1, LineString2, linestring_tag, linestring_tag, 2, Reverse>
    //: detail::equals::equals_by_collection<detail::equals::length_check>
    : detail::equals::equals_by_relate<LineString1, LineString2>
{};

template <typename LineString, typename MultiLineString, bool Reverse>
struct equals<LineString, MultiLineString, linestring_tag, multi_linestring_tag, 2, Reverse>
    : detail::equals::equals_by_relate<LineString, MultiLineString>
{};

template <typename MultiLineString1, typename MultiLineString2, bool Reverse>
struct equals<MultiLineString1, MultiLineString2, multi_linestring_tag, multi_linestring_tag, 2, Reverse>
    : detail::equals::equals_by_relate<MultiLineString1, MultiLineString2>
{};

template <typename MultiPolygon1, typename MultiPolygon2, bool Reverse>
struct equals
    <
        MultiPolygon1, MultiPolygon2,
        multi_polygon_tag, multi_polygon_tag,
        2,
        Reverse
    >
    : detail::equals::equals_by_collection<detail::equals::area_check>
{};

template <typename Polygon, typename MultiPolygon, bool Reverse>
struct equals
    <
        Polygon, MultiPolygon,
        polygon_tag, multi_polygon_tag,
        2,
        Reverse
    >
    : detail::equals::equals_by_collection<detail::equals::area_check>
{};

} // namespace dispatch
#endif // DOXYGEN_NO_DISPATCH

namespace resolve_variant {

template <typename Geometry1, typename Geometry2>
struct equals
{
    static inline bool apply(Geometry1 const& geometry1,
                             Geometry2 const& geometry2)
    {
        concept::check_concepts_and_equal_dimensions
        <
            Geometry1 const,
            Geometry2 const
        >();

        return dispatch::equals<Geometry1, Geometry2>
                       ::apply(geometry1, geometry2);
    }
};

template <BOOST_VARIANT_ENUM_PARAMS(typename T), typename Geometry2>
struct equals<boost::variant<BOOST_VARIANT_ENUM_PARAMS(T)>, Geometry2>
{
    struct visitor: static_visitor<bool>
    {
        Geometry2 const& m_geometry2;

        visitor(Geometry2 const& geometry2)
            : m_geometry2(geometry2)
        {}

        template <typename Geometry1>
        inline bool operator()(Geometry1 const& geometry1) const
        {
            return equals<Geometry1, Geometry2>
                   ::apply(geometry1, m_geometry2);
        }

    };

    static inline bool apply(
        boost::variant<BOOST_VARIANT_ENUM_PARAMS(T)> const& geometry1,
        Geometry2 const& geometry2
    )
    {
        return boost::apply_visitor(visitor(geometry2), geometry1);
    }
};

template <typename Geometry1, BOOST_VARIANT_ENUM_PARAMS(typename T)>
struct equals<Geometry1, boost::variant<BOOST_VARIANT_ENUM_PARAMS(T)> >
{
    struct visitor: static_visitor<bool>
    {
        Geometry1 const& m_geometry1;

        visitor(Geometry1 const& geometry1)
            : m_geometry1(geometry1)
        {}

        template <typename Geometry2>
        inline bool operator()(Geometry2 const& geometry2) const
        {
            return equals<Geometry1, Geometry2>
                   ::apply(m_geometry1, geometry2);
        }

    };

    static inline bool apply(
        Geometry1 const& geometry1,
        boost::variant<BOOST_VARIANT_ENUM_PARAMS(T)> const& geometry2
    )
    {
        return boost::apply_visitor(visitor(geometry1), geometry2);
    }
};

template <
    BOOST_VARIANT_ENUM_PARAMS(typename T1),
    BOOST_VARIANT_ENUM_PARAMS(typename T2)
>
struct equals<
    boost::variant<BOOST_VARIANT_ENUM_PARAMS(T1)>,
    boost::variant<BOOST_VARIANT_ENUM_PARAMS(T2)>
>
{
    struct visitor: static_visitor<bool>
    {
        template <typename Geometry1, typename Geometry2>
        inline bool operator()(Geometry1 const& geometry1,
                               Geometry2 const& geometry2) const
        {
            return equals<Geometry1, Geometry2>
                   ::apply(geometry1, geometry2);
        }

    };

    static inline bool apply(
        boost::variant<BOOST_VARIANT_ENUM_PARAMS(T1)> const& geometry1,
        boost::variant<BOOST_VARIANT_ENUM_PARAMS(T2)> const& geometry2
    )
    {
        return boost::apply_visitor(visitor(), geometry1, geometry2);
    }
};

} // namespace resolve_variant

/*!
\brief \brief_check{are spatially equal}
\details \details_check12{equals, is spatially equal}. Spatially equal means
    that the same point set is included. A box can therefore be spatially equal
    to a ring or a polygon, or a linestring can be spatially equal to a
    multi-linestring or a segment. This only works theoretically, not all
    combinations are implemented yet.
\ingroup equals
\tparam Geometry1 \tparam_geometry
\tparam Geometry2 \tparam_geometry
\param geometry1 \param_geometry
\param geometry2 \param_geometry
\return \return_check2{are spatially equal}

\qbk{[include reference/algorithms/equals.qbk]}

 */
template <typename Geometry1, typename Geometry2>
inline bool equals(Geometry1 const& geometry1, Geometry2 const& geometry2)
{
    return resolve_variant::equals<Geometry1, Geometry2>
                          ::apply(geometry1, geometry2);
}

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_ALGORITHMS_EQUALS_HPP


#include <Python.h>
#include <numpy/arrayobject.h>
#include <Eigen/Dense>

#include "george.h"

#if PY_MAJOR_VERSION >= 3
#define PyInt_AsLong PyLong_AsLong
#endif

using namespace george;

using george::HODLRSolver;

// Eigen is column major and numpy is row major. Barf.
typedef Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic,
                          Eigen::RowMajor> > RowMajorMap;

#define PARSE_ARRAY(o) (PyArrayObject*) PyArray_FROM_OTF(o, NPY_DOUBLE, \
        NPY_IN_ARRAY)

extern "C" {

// Type definition.
typedef struct {

    PyObject_HEAD
    kernels::Kernel* kernel;
    HODLRSolver<kernels::Kernel>* solver;

} _george_object;

// Allocation/deallocation.
static void _george_dealloc (_george_object* self);
static PyObject* _george_new (PyTypeObject* type, PyObject* args, PyObject* kwds);
static int _george_init (_george_object* self, PyObject* args, PyObject* kwds);

// Methods.
static PyObject* _george_compute (_george_object* self, PyObject* args);
static PyObject* _george_lnlikelihood (_george_object* self, PyObject* args);
static PyObject* _george_computed (_george_object* self, PyObject* args);
static PyObject* _george_predict (_george_object* self, PyObject* args);
static PyObject* _george_get_matrix (_george_object* self, PyObject* args);

// Python 3.

#if PY_MAJOR_VERSION >= 3

PyObject *PyInit__george (void);
static int _george_traverse(PyObject* m, visitproc visit, void* arg);
static int _george_clear(PyObject* m);

#else

void init_george (void);

#endif

}

static void _george_dealloc (_george_object* self)
{
    if (self->kernel != NULL) delete self->kernel;
    if (self->solver != NULL) delete self->solver;
    Py_TYPE(self)->tp_free((PyObject*)self);
}

static PyObject *_george_new (PyTypeObject* type, PyObject* args, PyObject* kwds)
{
    _george_object* self = (_george_object*)type->tp_alloc(type, 0);
    self->kernel = NULL;
    self->solver = NULL;
    return (PyObject*)self;
}

kernels::Kernel* parse_kernel (PyObject* kernel)
{
    // Check the kernel.
    if (!PyObject_HasAttrString(kernel, "is_kernel")) {
        PyErr_SetString(PyExc_TypeError, "Invalid kernel");
        return NULL;
    }

    // Deal with operators.
    PyObject* is_kernel_obj = PyObject_GetAttrString(kernel, "is_kernel");
    if (is_kernel_obj == Py_False) {
        Py_DECREF(is_kernel_obj);

        // Get the Python kernel objects and the type of operator.
        PyObject* k1 = PyObject_GetAttrString(kernel, "k1"),
                * k2 = PyObject_GetAttrString(kernel, "k2"),
                * otype_obj = PyObject_GetAttrString(kernel, "operator_type");
        long otype = PyInt_AsLong(otype_obj);
        Py_XDECREF(otype_obj);

        // Validate the kernels.
        if (k1 == NULL || k2 == NULL || PyErr_Occurred() != NULL) {
            PyErr_SetString(PyExc_TypeError, "Invalid operator");
            Py_XDECREF(k1);
            Py_XDECREF(k2);
            return NULL;
        }

        // Parse the combined kernels.
        kernels::Kernel* kernel1 = parse_kernel(k1);
        kernels::Kernel* kernel2 = parse_kernel(k2);
        Py_DECREF(k1);
        Py_DECREF(k2);

        // Check that these were successfully parsed.
        if (kernel1 == NULL || kernel2 == NULL) {
            if (kernel1 != NULL) delete kernel1;
            if (kernel2 != NULL) delete kernel2;
            return NULL;
        }

        // Combine the kernels.
        if (otype == 0)
            return new kernels::Sum (kernel1, kernel2);
        else if (otype == 1)
            return new kernels::Product (kernel1, kernel2);

        // If we get here then the operator type was unknown.
        if (kernel1 != NULL) delete kernel1;
        if (kernel2 != NULL) delete kernel2;
        PyErr_SetString(PyExc_TypeError, "Unknown operator");
        return NULL;
    }

    // If we get here then the object claims to be a kernel.
    Py_DECREF(is_kernel_obj);

    // Get the kernel type.
    PyObject* ktype_obj = PyObject_GetAttrString(kernel, "kernel_type");
    long ktype = PyInt_AsLong(ktype_obj);
    Py_XDECREF(ktype_obj);
    if (PyErr_Occurred() != NULL) {
        PyErr_SetString(PyExc_TypeError, "Invalid kernel type");
        return NULL;
    }

    // Get the number of dimensions.
    PyObject* ndim_obj = PyObject_GetAttrString(kernel, "ndim");
    long ndim = PyInt_AsLong(ndim_obj);
    Py_XDECREF(ndim_obj);
    if (PyErr_Occurred() != NULL) {
        PyErr_SetString(PyExc_TypeError,
                        "Couldn't parse number of kernel dimensions");
        return NULL;
    }

    // Get the parameter vector.
    PyObject* pars_obj = PyObject_GetAttrString(kernel, "pars");
    PyArrayObject* pars_array = PARSE_ARRAY(pars_obj);
    Py_XDECREF(pars_obj);
    if (pars_array == NULL) {
        PyErr_SetString(PyExc_ValueError, "Invalid parameter vector");
        Py_XDECREF(pars_array);
        return NULL;
    }

    // Get the number of parameters and a pointer to the parameter data.
    int npars = PyArray_DIM(pars_array, 0);
    double* pars = (double*)PyArray_DATA(pars_array);

    // Build the kernel.
    kernels::Kernel* k = NULL;

#define GEORGE_DEFINE_KERNEL(N,NAME) { \
    if (npars != N) { \
        PyErr_SetString(PyExc_ValueError, "The NAME requires N parameters."); \
        Py_DECREF(pars_array); \
        return NULL; \
    } \
    k = new kernels::NAME(ndim, pars); \
}

    if (ktype == 0)      GEORGE_DEFINE_KERNEL (1, ConstantKernel)
    else if (ktype == 1) GEORGE_DEFINE_KERNEL (0, DotProductKernel)
    else if (ktype == 2) GEORGE_DEFINE_KERNEL ((ndim*ndim+ndim)/2, ExpKernel)
    else if (ktype == 3) GEORGE_DEFINE_KERNEL ((ndim*ndim+ndim)/2, RBFKernel)
    else if (ktype == 4) GEORGE_DEFINE_KERNEL (1, CosineKernel)
    else if (ktype == 5) GEORGE_DEFINE_KERNEL (2, ExpSine2Kernel)
    else if (ktype == 6) GEORGE_DEFINE_KERNEL ((ndim*ndim+ndim)/2, Matern32Kernel)
    else if (ktype == 7) GEORGE_DEFINE_KERNEL ((ndim*ndim+ndim)/2, Matern52Kernel)
    else if (ktype == 8) GEORGE_DEFINE_KERNEL (1, WhiteKernel)
    else if (ktype == 9) GEORGE_DEFINE_KERNEL (1+(ndim*ndim+ndim)/2, RationalQuadraticKernel)
    else PyErr_SetString(PyExc_TypeError, "Unknown kernel");

#undef GEORGE_DEFINE_KERNEL

    Py_DECREF(pars_array);
    return k;
}

static int _george_init(_george_object* self, PyObject* args, PyObject* kwds)
{
    int nleaf;
    double tol;
    PyObject* kernel = NULL;
    if (!PyArg_ParseTuple(args, "Oid", &kernel, &nleaf, &tol))
        return -1;

    // Parse the kernel.
    self->kernel = parse_kernel (kernel);
    if (self->kernel == NULL) return -2;

    // Set up the solver.
    self->solver = new HODLRSolver<kernels::Kernel> (self->kernel, nleaf, tol);

    return 0;
}

static PyObject* _george_compute (_george_object* self, PyObject* args)
{
    int seed;
    PyObject* x_obj, * yerr_obj;

    // Parse the input arguments.
    if (!PyArg_ParseTuple(args, "OOi", &x_obj, &yerr_obj, &seed))
        return NULL;

    // Decode the numpy arrays.
    PyArrayObject * x_array = PARSE_ARRAY(x_obj),
                  * yerr_array = PARSE_ARRAY(yerr_obj);
    if (x_array == NULL || yerr_array == NULL) {
        Py_XDECREF(x_array);
        Py_XDECREF(yerr_array);
        PyErr_SetString(PyExc_ValueError,
            "Failed to parse input objects as numpy arrays");
        return NULL;
    }

    // Get the dimensions.
    int nsamples = (int)PyArray_DIM(x_array, 0),
        ndim = (int)PyArray_DIM(x_array, 1);
    if ((int)PyArray_NDIM(x_array) != 2 ||
            (int)PyArray_DIM(yerr_array, 0) != nsamples) {
        Py_DECREF(x_array);
        Py_DECREF(yerr_array);
        PyErr_SetString(PyExc_ValueError, "Dimension mismatch");
        return NULL;
    }

    // Access the data.
    double * x = (double*)PyArray_DATA(x_array),
           * yerr = (double*)PyArray_DATA(yerr_array);

    // Map to vectors.
    MatrixXd x_vec = RowMajorMap(x, nsamples, ndim);
    VectorXd yerr_vec = VectorXd::Map(yerr, nsamples);

    // Pre-compute the factorization.
    int info = self->solver->compute (x_vec, yerr_vec, seed);

    // Clean up.
    Py_DECREF(x_array);
    Py_DECREF(yerr_array);

    // Check success.
    if (info != george::SOLVER_OK) {
        PyErr_SetString(PyExc_RuntimeError, "Failed to compute model");
        return NULL;
    }

    Py_INCREF(Py_None);
    return Py_None;
}

static PyObject* _george_lnlikelihood (_george_object* self, PyObject* args)
{
    PyObject *y_obj;
    if (!PyArg_ParseTuple(args, "O", &y_obj)) return NULL;

    if (!self->solver->get_computed()) {
        PyErr_SetString(PyExc_RuntimeError,
            "You need to compute the model first");
        return NULL;
    }

    PyArrayObject *y_array = PARSE_ARRAY(y_obj);
    if (y_array == NULL) {
        Py_XDECREF(y_array);
        PyErr_SetString(PyExc_ValueError,
            "Failed to parse input object as a numpy array");
        return NULL;
    }

    // Get the data.
    int nsamples = (int)PyArray_DIM(y_array, 0);
    double* y = (double*)PyArray_DATA(y_array);
    VectorXd y_vec = VectorXd::Map(y, nsamples);

    // Compute the likelihood.
    double lnlike = self->solver->log_likelihood(y_vec);
    Py_DECREF(y_array);

    // Check success.
    if (self->solver->get_status() != george::SOLVER_OK) {
        PyErr_SetString(PyExc_RuntimeError, "Failed to compute likelihood");
        return NULL;
    }

    return Py_BuildValue("d", lnlike);
}

static PyObject* _george_predict (_george_object* self, PyObject* args)
{
    PyObject* y_obj, * x_obj;
    if (!PyArg_ParseTuple(args, "OO", &y_obj, &x_obj)) return NULL;

    if (!self->solver->get_computed()) {
        PyErr_SetString(PyExc_RuntimeError,
            "You need to compute the model first");
        return NULL;
    }

    PyArrayObject *y_array = PARSE_ARRAY(y_obj),
                  *x_array = PARSE_ARRAY(x_obj);
    if (y_array == NULL || x_array == NULL) {
        Py_XDECREF(y_array);
        Py_XDECREF(x_array);
        PyErr_SetString(PyExc_ValueError,
            "Failed to parse input objects as numpy arrays");
        return NULL;
    }

    // Get the dimensions.
    int nsamples = (int)PyArray_DIM(y_array, 0),
        ntest = (int)PyArray_DIM(x_array, 0),
        ndim = (int)PyArray_DIM(x_array, 1);
    if ((int)PyArray_NDIM(x_array) != 2 ||
            self->solver->get_dimension() != nsamples) {
        Py_DECREF(x_array);
        Py_DECREF(y_array);
        PyErr_SetString(PyExc_ValueError, "Dimension mismatch");
        return NULL;
    }

    // Access the data.
    double* y = (double*) PyArray_DATA(y_array),
          * x = (double*) PyArray_DATA(x_array);

    // Compute the mean.
    VectorXd y_vec = VectorXd::Map(y, nsamples),
             mu_vec;
    MatrixXd x_vec = RowMajorMap(x, ntest, ndim);
    MatrixXd cov_mat;
    self->solver->predict(y_vec, x_vec, mu_vec, cov_mat);

    // Clean up.
    Py_DECREF(y_array);
    Py_DECREF(x_array);

    // Check success.
    if (self->solver->get_status() != george::SOLVER_OK) {
        PyErr_SetString(PyExc_RuntimeError, "Failed to compute model");
        return NULL;
    }

    // Allocate the output arrays.
    npy_intp dim[] = {ntest}, dim2[] = {ntest, ntest};
    PyArrayObject* mu_array = (PyArrayObject*)PyArray_SimpleNew(1, dim, NPY_DOUBLE),
                 * cov_array = (PyArrayObject*)PyArray_SimpleNew(2, dim2, NPY_DOUBLE);
    if (mu_array == NULL || cov_array == NULL) {
        Py_XDECREF(mu_array);
        Py_XDECREF(cov_array);
        return NULL;
    }

    // Copy over the result.
    double *mu = (double*)PyArray_DATA(mu_array),
           *cov = (double*)PyArray_DATA(cov_array);
    for (int i = 0; i < ntest; ++i) {
        mu[i] = mu_vec[i];
        for (int j = 0; j < ntest; ++j) cov[i*ntest+j] = cov_mat(i, j);
    }

    // Build the result.
    PyObject *ret = Py_BuildValue("OO", mu_array, cov_array);
    Py_DECREF(mu_array);
    Py_DECREF(cov_array);

    if (ret == NULL) {
        PyErr_SetString(PyExc_RuntimeError, "Couldn't build output tuple");
        Py_XDECREF(ret);
        return NULL;
    }

    return ret;
}

static PyObject* _george_get_matrix (_george_object* self, PyObject* args)
{
    PyObject* t_obj;
    if (!PyArg_ParseTuple(args, "O", &t_obj)) return NULL;

    PyArrayObject* t_array = PARSE_ARRAY(t_obj);
    if (t_array == NULL) {
        Py_XDECREF(t_array);
        PyErr_SetString(PyExc_ValueError,
            "Failed to parse input object as numpy array");
        return NULL;
    }

    // Check the shape of the input array.
    if ((int)PyArray_NDIM(t_array) != 2) {
        Py_DECREF(t_array);
        PyErr_SetString(PyExc_ValueError, "Dimension mismatch");
        return NULL;
    }

    // Get the dimensions of the problem.
    int n = (int)PyArray_DIM(t_array, 0),
        ndim = (int)PyArray_DIM(t_array, 1);

    // Access the data.
    double* t = (double*) PyArray_DATA(t_array);
    MatrixXd tm = RowMajorMap(t, n, ndim);

    // Allocate the output arrays.
    npy_intp dim[] = {n, n};
    PyArrayObject* out_array = (PyArrayObject*)PyArray_SimpleNew(2, dim, NPY_DOUBLE);
    if (out_array == NULL) {
        Py_DECREF(t_array);
        Py_XDECREF(out_array);
        return NULL;
    }

    // Copy over the result.
    double* matrix = (double*)PyArray_DATA(out_array), value;
    kernels::Kernel* kernel = self->kernel;
    for (int i = 0; i < n; ++i) {
        matrix[i*n+i] = kernel->evaluate(tm.row(i), tm.row(i));
        for (int j = 0; j < i; ++j) {
            value = kernel->evaluate(tm.row(i), tm.row(j));
            matrix[i*n+j] = value;
            matrix[j*n+i] = value;
        }
    }

    Py_DECREF(t_array);
    return (PyObject*)out_array;
}

static PyObject* _george_computed(_george_object* self, PyObject* args)
{
    if (self->solver->get_computed()) Py_RETURN_TRUE;
    Py_RETURN_FALSE;
}

static PyMethodDef _george_methods[] = {
    {"compute",
     (PyCFunction)_george_compute,
     METH_VARARGS,
     "Fit the GP."},
    {"lnlikelihood",
     (PyCFunction)_george_lnlikelihood,
     METH_VARARGS,
     "Get the marginalized ln likelihood of some values."
    },
    {"predict",
     (PyCFunction)_george_predict,
     METH_VARARGS,
     "Predict the mean function"
    },
    {"get_matrix",
     (PyCFunction)_george_get_matrix,
     METH_VARARGS,
     "Get the covariance matrix for a set of times."
    },
    {"computed",
     (PyCFunction)_george_computed,
     METH_NOARGS,
     "Has the GP been computed?"
    },
    {NULL}  /* Sentinel */
};

static char _george_doc[] = "This is the ``_george`` object. "
                            "There is some black magic.";
static PyTypeObject _george_type = {
    PyVarObject_HEAD_INIT(NULL, 0)
    "_george._george",           /*tp_name*/
    sizeof(_george_object),      /*tp_basicsize*/
    0,                           /*tp_itemsize*/
    (destructor)_george_dealloc, /*tp_dealloc*/
    0,                           /*tp_print*/
    0,                           /*tp_getattr*/
    0,                           /*tp_setattr*/
    0,                           /*tp_compare*/
    0,                           /*tp_repr*/
    0,                           /*tp_as_number*/
    0,                           /*tp_as_sequence*/
    0,                           /*tp_as_mapping*/
    0,                           /*tp_hash */
    0,                           /*tp_call*/
    0,                           /*tp_str*/
    0,                           /*tp_getattro*/
    0,                           /*tp_setattro*/
    0,                           /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
                                 /*tp_flags*/
    _george_doc,                 /* tp_doc */
    0,                           /* tp_traverse */
    0,                           /* tp_clear */
    0,                           /* tp_richcompare */
    0,                           /* tp_weaklistoffset */
    0,                           /* tp_iter */
    0,                           /* tp_iternext */
    _george_methods,             /* tp_methods */
    0,                           /* tp_members */
    0,                           /* tp_getset */
    0,                           /* tp_base */
    0,                           /* tp_dict */
    0,                           /* tp_descr_get */
    0,                           /* tp_descr_set */
    0,                           /* tp_dictoffset */
    (initproc)_george_init,      /* tp_init */
    0,                           /* tp_alloc */
    _george_new,                 /* tp_new */
};

//
// Initialize the module.
//

static char module_doc[] = "GP Module";

#if PY_MAJOR_VERSION >= 3

struct george_state {
    PyObject *error;
};
#define GETSTATE(m) ((struct george_state*)PyModule_GetState(m))

static int _george_traverse(PyObject* m, visitproc visit, void* arg) {
    Py_VISIT(GETSTATE(m)->error);
    return 0;
}

static int _george_clear(PyObject* m) {
    Py_CLEAR(GETSTATE(m)->error);
    return 0;
}

extern "C" {
static struct PyModuleDef moduledef = {
    PyModuleDef_HEAD_INIT,
    "_george",
    module_doc,
    sizeof(struct george_state),
    _george_methods,
    NULL,
    _george_traverse,
    _george_clear,
    NULL
};
}

#define INITERROR return NULL

PyObject *PyInit__george (void)

#else
#define INITERROR return

void init_george (void)

#endif
{
    PyObject *m;
    if (PyType_Ready(&_george_type) < 0) INITERROR;

#if PY_MAJOR_VERSION >= 3
    m = PyModule_Create(&moduledef);
#else
    m = Py_InitModule3("_george", _george_methods, module_doc);
#endif

    if (m == NULL) INITERROR;

    // Add the _george type.
    Py_INCREF(&_george_type);
    PyModule_AddObject(m, "_george", (PyObject *)&_george_type);

    // Numpy.
    import_array();

#if PY_MAJOR_VERSION >= 3
    return m;
#endif
}


#pragma once
#include <Eigen/Core>
#include <Eigen/Sparse>
#include <algorithm>
#include <atomic>
#include <cstddef>
#include <cstdint>
#include <future>
#include <iostream>
#include <random>
#include <stdexcept>
#include <thread>
#include <tuple>
#include <vector>

#include "argcheck.hpp"

namespace irspack {
namespace sparse_util {

template <typename Real>
using CSRMatrix = Eigen::SparseMatrix<Real, Eigen::RowMajor>;

template <typename Real>
using RowMajorMatrix =
    Eigen::Matrix<Real, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;

template <typename Real>
using ColMajorMatrix =
    Eigen::Matrix<Real, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>;

template <typename Real>
using DenseVector = Eigen::Matrix<Real, Eigen::Dynamic, 1>;

template <typename Real>
using DenseColVector = Eigen::Matrix<Real, 1, Eigen::Dynamic, Eigen::RowMajor>;

template <typename Real>
using CSCMatrix = Eigen::SparseMatrix<Real, Eigen::ColMajor>;

template <typename Real>
inline RowMajorMatrix<Real>
parallel_sparse_product(const CSRMatrix<Real> &left,
                        const CSCMatrix<Real> &right, const size_t n_threads) {
  RowMajorMatrix<Real> result(left.rows(), right.cols());
  result.array() = 0;
  check_arg(n_threads > 0, "n_thraed must be > 0");
  const int64_t n_row = left.rows();
  std::atomic<int64_t> cursor(0);
  std::vector<std::thread> workers;
  for (int i = 0; i < static_cast<int>(n_threads); i++) {
    workers.emplace_back([&left, &right, &cursor, n_row, &result]() {
      const int64_t chunk_size = 16;
      while (true) {
        auto current_position = cursor.fetch_add(chunk_size);
        if (current_position >= n_row) {
          break;
        }
        auto block_size =
            std::min(current_position + chunk_size, n_row) - current_position;
        result.middleRows(current_position, block_size) +=
            left.middleRows(current_position, block_size) * right;
      }
    });
  }
  for (auto &worker : workers) {
    worker.join();
  }
  return result;
}

template <typename Real, class Derived, typename Integer = int64_t>
struct SplitFunction {
  template <typename... Args>
  static std::pair<CSRMatrix<Real>, CSRMatrix<Real>>
  split_imple(const CSRMatrix<Real> &X, std::int64_t random_seed,
              Args... args) {
    Derived::check_args(args...);
    using Triplet = Eigen::Triplet<Integer>;
    std::mt19937 random_state(random_seed);
    std::vector<Integer> col_buffer;
    std::vector<Real> data_buffer;
    std::vector<uint64_t> index_;
    std::vector<Triplet> train_data, test_data;
    for (int row = 0; row < X.outerSize(); ++row) {
      col_buffer.clear(); // does not change capacity
      data_buffer.clear();
      index_.clear();
      Integer cnt = 0;
      for (typename CSRMatrix<Real>::InnerIterator it(X, row); it; ++it) {
        index_.push_back(cnt);
        col_buffer.push_back(it.col());
        data_buffer.push_back(it.value());
        cnt += 1;
      }
      std::shuffle(index_.begin(), index_.end(), random_state);
      size_t n_test = Derived::get_n_test(cnt, args...);
      for (size_t i = 0; i < n_test; i++) {
        test_data.emplace_back(row, col_buffer[index_[i]],
                               data_buffer[index_[i]]);
      }
      for (size_t i = n_test; i < col_buffer.size(); i++) {
        train_data.emplace_back(row, col_buffer[index_[i]],
                                data_buffer[index_[i]]);
      }
    }
    CSRMatrix<Real> X_train(X.rows(), X.cols()), X_test(X.rows(), X.cols());
    auto dupfunction = [](const Integer &a, const Integer &b) { return a + b; };
    X_train.setFromTriplets(train_data.begin(), train_data.end(), dupfunction);
    X_test.setFromTriplets(test_data.begin(), test_data.end(), dupfunction);
    X_train.makeCompressed();
    X_test.makeCompressed();
    return {X_train, X_test};
  }
};

template <typename Real, typename Integer = int64_t>
struct SplitByRatioFunction
    : SplitFunction<Real, SplitByRatioFunction<Real, Integer>, Integer> {
  using Base =
      SplitFunction<Real, SplitByRatioFunction<Real, Integer>, Integer>;
  static void check_args(double test_ratio, bool n_test_ceil) {
    check_arg(((test_ratio <= 1.0 && (test_ratio >= 0.0))),
              "test_ratio must be within [0.0, 1.0]");
  }

  static size_t get_n_test(size_t nnz_row, double test_ratio,
                           bool n_test_ceil) {
    if (n_test_ceil) {
      return std::ceil(nnz_row * test_ratio);
    } else {
      return std::floor(nnz_row * test_ratio);
    }
  };

  static std::pair<CSRMatrix<Real>, CSRMatrix<Real>>
  split(const CSRMatrix<Real> &X, std::int64_t random_seed, Real heldout_ratio,
        bool n_test_ceil) {
    return Base::split_imple(X, random_seed, heldout_ratio, n_test_ceil);
  }
};

template <typename Real, typename Integer = int64_t>
struct SplitFixedN : SplitFunction<Real, SplitFixedN<Real, Integer>, Integer> {
  using Base = SplitFunction<Real, SplitFixedN<Real, Integer>, Integer>;
  static void check_args(size_t n_held_out) {}

  static size_t get_n_test(size_t nnz_row, size_t n_held_out) {
    return std::min(nnz_row, n_held_out);
  };

  static std::pair<CSRMatrix<Real>, CSRMatrix<Real>>
  split(const CSRMatrix<Real> &X, std::int64_t random_seed, size_t n_held_out) {
    return Base::split_imple(X, random_seed, n_held_out);
  }
};

template <typename Real>
CSRMatrix<Real> okapi_BM_25_weight(const CSRMatrix<Real> &X, Real k1, Real b) {
  CSRMatrix<Real> result(X);
  using itertype = typename CSRMatrix<Real>::InnerIterator;
  const int N = X.rows();
  result.makeCompressed();
  DenseVector<Real> idf(X.cols());
  DenseVector<Real> doc_length(N);
  idf.array() = 0;
  doc_length.array() = 0;

  for (int i = 0; i < N; i++) {
    for (itertype iter(X, i); iter; ++iter) {
      idf(iter.col()) += 1;
      doc_length(i) += iter.value();
    }
  }
  Real avgdl = doc_length.sum() / N;
  idf.array() =
      (N / (idf.array() + static_cast<Real>(1)) + static_cast<Real>(1)).log();
  for (int i = 0; i < N; i++) {
    Real regularizer = k1 * (1 - b + b * doc_length(i) / avgdl);
    for (itertype iter(result, i); iter; ++iter) {
      iter.valueRef() = idf(iter.col()) * (iter.valueRef() * (k1 + 1)) /
                        (iter.valueRef() + regularizer);
    }
  }
  return result;
}

template <typename Real>
CSRMatrix<Real> tf_idf_weight(const CSRMatrix<Real> &X, bool smooth) {
  CSRMatrix<Real> result(X);
  using itertype = typename CSRMatrix<Real>::InnerIterator;
  const int N = X.rows();
  result.makeCompressed();
  DenseVector<Real> idf(X.cols());
  idf.array() = 0;

  for (int i = 0; i < N; i++) {
    for (itertype iter(X, i); iter; ++iter) {
      idf(iter.col()) += 1;
    }
  }
  idf.array() = (N / (idf.array() + static_cast<Real>(smooth))).log();
  for (int i = 0; i < N; i++) {
    for (itertype iter(result, i); iter; ++iter) {
      iter.valueRef() *= idf(iter.col());
    }
  }
  return result;
}

template <typename Real>
CSRMatrix<Real> remove_diagonal(const CSRMatrix<Real> &X) {
  check_arg(X.rows() == X.cols(), "X must be square");
  CSRMatrix<Real> result(X);
  using itertype = typename CSRMatrix<Real>::InnerIterator;
  const int N = X.rows();
  result.makeCompressed();
  for (int i = 0; i < N; i++) {
    for (itertype iter(result, i); iter; ++iter) {
      if (i == iter.col()) {
        iter.valueRef() = static_cast<Real>(0);
      }
    }
  }
  return result;
}

template <typename Real, bool positive_only = false,
          int block_size = Eigen::internal::packet_traits<Real>::size>
inline CSCMatrix<Real> SLIM(const CSRMatrix<Real> &X, size_t n_threads,
                            size_t n_iter, Real l2_coeff, Real l1_coeff,
                            Real tol, int64_t top_k) {
  check_arg(n_threads > 0, "n_threads must be > 0.");
  check_arg(n_iter > 0, "n_iter must be > 0.");
  check_arg(l2_coeff >= 0, "l2_coeff must be > 0.");
  check_arg(l1_coeff >= 0, "l1_coeff must be > 0.");
  using MatrixType =
      Eigen::Matrix<Real, block_size, Eigen::Dynamic, Eigen::ColMajor>;
  using VectorType = Eigen::Matrix<Real, block_size, 1>;

  CSCMatrix<Real> X_csc(X);
  X_csc.makeCompressed();
  using TripletType = Eigen::Triplet<Real>;
  using CSCIter = typename CSCMatrix<Real>::InnerIterator;
  using RealAndIndex = std::pair<Real, int64_t>;
  std::vector<std::future<std::vector<TripletType>>> workers;
  std::atomic<int64_t> cursor(0);
  for (size_t th = 0; th < n_threads; th++) {
    workers.emplace_back(std::async(std::launch::async, [&cursor, &X_csc,
                                                         l2_coeff, l1_coeff,
                                                         n_iter, tol, top_k] {
      const int64_t F = X_csc.cols();
      std::mt19937 gen(0);
      std::vector<int64_t> indices(F);
      for (int64_t i = 0; i < F; i++) {
        indices[i] = i;
      }
      MatrixType remnants(block_size, X_csc.rows());
      MatrixType coeffs(block_size, F);
      VectorType coeff_temp(block_size);
      VectorType diff(block_size);
      VectorType linear(block_size);
      VectorType linear_plus(block_size);
      VectorType linear_minus(block_size);

      std::vector<RealAndIndex> argsort_buffer;
      if (top_k >= 0) {
        argsort_buffer.resize(F);
      }

      std::vector<TripletType> local_results;
      while (true) {
        int64_t current_cursor = cursor.fetch_add(block_size);
        if (current_cursor >= F) {
          break;
        }

        int64_t block_begin = current_cursor;
        int64_t block_end = std::min(block_begin + block_size, F);
        int64_t valid_block_size = block_end - block_begin;
        remnants.setZero();
        coeffs.setZero();

        for (int64_t f_cursor = block_begin; f_cursor < block_end; f_cursor++) {
          const int64_t internal_col_position = f_cursor - block_begin;
          for (CSCIter iter(X_csc, f_cursor); iter; ++iter) {
            remnants(internal_col_position, iter.row()) = -iter.value();
          }
        }

        for (size_t cd_iteration = 0; cd_iteration < n_iter; cd_iteration++) {
          std::shuffle(indices.begin(), indices.end(), gen);
          Real delta = 0;
          for (int64_t feature_index = 0; feature_index < F; feature_index++) {
            int64_t shuffled_feature_index = indices[feature_index];
            coeff_temp = coeffs.col(shuffled_feature_index);
            diff = coeff_temp;
            linear.setZero();
            Real x2_sum = 0.0;
            for (CSCIter nnz_iter(X_csc, shuffled_feature_index); nnz_iter;
                 ++nnz_iter) {
              Real x = nnz_iter.value();
              x2_sum += x * x;
              linear.noalias() += x * remnants.col(nnz_iter.row());
            }
            linear.noalias() -= x2_sum * coeff_temp;

            Real quadratic = x2_sum + l2_coeff;
            linear_plus.array() = (-linear.array() - l1_coeff) / quadratic;
            if (!positive_only) {
              linear_minus.array() = (-linear.array() + l1_coeff) / quadratic;
            }

            Real *ptr_location =
                coeffs.data() + shuffled_feature_index * block_size;
            Real *lp_ptr = linear_plus.data();
            Real *lm_ptr = linear_minus.data();

            for (int64_t inner_cursor_position = 0;
                 inner_cursor_position < block_size; inner_cursor_position++) {
              Real lplus = *(lp_ptr++);
              Real lminus = *(lm_ptr++);
              int64_t original_cursor_position =
                  inner_cursor_position + block_begin;
              if (original_cursor_position == shuffled_feature_index) {
                *(ptr_location++) = 0.0;
                continue;
              }
              if (positive_only) {
                if (lplus > 0) {
                  *(ptr_location++) = lplus;
                } else {
                  *(ptr_location++) = static_cast<Real>(0.0);
                }

              } else {
                if (lplus > 0) {
                  *(ptr_location++) = lplus;
                } else {
                  if (lminus < 0) {
                    *(ptr_location++) = lminus;
                  } else {
                    *(ptr_location++) = static_cast<Real>(0.0);
                  }
                }
              } // allow nagative block
            }
            coeff_temp.noalias() =
                coeffs.col(shuffled_feature_index) - coeff_temp;

            if (!coeff_temp.isZero()) {
              for (CSCIter nnz_iter(X_csc, shuffled_feature_index); nnz_iter;
                   ++nnz_iter) {
                const int64_t row = nnz_iter.row();
                remnants.col(row).noalias() += nnz_iter.valueRef() * coeff_temp;
              }
              delta = std::max(delta, coeff_temp.cwiseAbs().array().maxCoeff());
            }
          }
          if (delta < tol) {
            break;
          }
        }

        if (top_k < 0) {
          for (int64_t f = 0; f < F; f++) {
            for (int64_t inner_cursor_position = 0;
                 inner_cursor_position < valid_block_size;
                 inner_cursor_position++) {
              int64_t original_location = inner_cursor_position + block_begin;
              Real c = coeffs(inner_cursor_position, f);
              if (c != 0.0) {
                local_results.emplace_back(f, original_location, c);
              }
            }
          }
        } else {
          for (int64_t inner_cursor_position = 0;
               inner_cursor_position < valid_block_size;
               inner_cursor_position++) {
            int64_t original_location = inner_cursor_position + block_begin;
            auto iter = argsort_buffer.begin();
            int64_t nnz = 0;
            for (int64_t f = 0; f < F; f++) {
              Real c = coeffs(inner_cursor_position, f);
              if (c != 0.0) {
                iter->first = c;
                iter->second = f;
                iter++;
                nnz++;
              }
            }
            int64_t n_taken_coeffs = nnz;
            if (nnz > top_k) {
              std::sort(argsort_buffer.begin(), argsort_buffer.end(),
                        [](RealAndIndex &val1, RealAndIndex &val2) {
                          return val1.first > val2.first;
                        });
              n_taken_coeffs = top_k;
            }
            for (int64_t i = 0; i < n_taken_coeffs; i++) {
              local_results.emplace_back(argsort_buffer[i].second,
                                         original_location,
                                         argsort_buffer[i].first);
            }
          }
        }
      }
      return local_results;
    }));
  }
  std::vector<TripletType> nnzs;
  for (auto &fres : workers) {
    auto result = fres.get();
    for (const auto &w : result) {
      nnzs.emplace_back(w);
    }
  }

  CSCMatrix<Real> result(X.cols(), X.cols());
  result.setFromTriplets(nnzs.begin(), nnzs.end());
  result.makeCompressed();
  return result;
}

template <typename Real>
inline std::vector<std::vector<std::pair<int64_t, float>>>
retrieve_recommend_from_score(
    const RowMajorMatrix<Real> &score,
    const std::vector<std::vector<int64_t>> &allowed_indices,
    const size_t cutoff, size_t n_threads) {
  using score_and_index = std::pair<int64_t, float>;
  check_arg(n_threads > 0, "n_threads must not be 0.");
  check_arg(
      (score.rows() == static_cast<int64_t>(allowed_indices.size())) ||
          allowed_indices.empty(),
      "allowed_indices, if not empty, must have a size equal to X.rows()");
  std::vector<std::vector<score_and_index>> result(score.rows());
  std::vector<std::future<void>> workers;
  std::atomic<size_t> cursor(0);
  const size_t n_users = static_cast<size_t>(score.rows());
  for (size_t thread = 0; thread < std::min(n_threads, n_users); thread++) {
    workers.emplace_back(
        std::async([&score, cutoff, &allowed_indices, &cursor, &result]() {
          const int64_t n_rows = score.rows();
          const int64_t n_items = score.cols();
          std::vector<score_and_index> index_holder;
          index_holder.reserve(n_items);

          while (true) {
            int64_t current = cursor.fetch_add(1);
            if (current >= n_rows) {
              break;
            }

            std::vector<score_and_index> inserted;
            const Real *score_ptr = score.data() + n_items * current;

            index_holder.clear();
            if (!allowed_indices.empty()) {
              for (auto item_index : allowed_indices.at(current)) {
                if ((item_index < n_items) && (item_index >= 0)) {
                  index_holder.emplace_back(item_index, score_ptr[item_index]);
                }
              }
            } else {
              for (int64_t i = 0; i < n_items; i++) {
                index_holder.emplace_back(i, *(score_ptr++));
              }
            }
            std::partial_sort(
                index_holder.begin(),
                index_holder.begin() + std::min(cutoff, index_holder.size()),
                index_holder.end(), [](score_and_index i1, score_and_index i2) {
                  return i1.second > i2.second;
                });

            size_t items_recommended = 0;
            for (auto item_index : index_holder) {
              if (items_recommended >= cutoff) {
                break;
              }

              if (item_index.second == -std::numeric_limits<Real>::infinity()) {
                break;
              }
              result[current].emplace_back(item_index);
              items_recommended++;
            }
          }
        }));
  }
  workers.clear();
  return result;
}

} // namespace sparse_util
} // namespace irspack