#include <NTL/mat_ZZ_p.h>
#include <NTL/vec_ZZVec.h>
#include <NTL/vec_long.h>
#include <NTL/BasicThreadPool.h>

// FIXME: only needed if we use multi-modular MM
#include <NTL/MatPrime.h>
#include <NTL/mat_lzz_p.h>

NTL_START_IMPL

// ******************** Matrix Multiplication ************************

#ifdef NTL_HAVE_LL_TYPE
#define NTL_USE_MM_MATMUL (1)
#else
#define NTL_USE_MM_MATMUL (0)
#endif

#define PAR_THRESH (40000.0)

// *********************** Plain Matrix Multiplication ***************

void plain_mul_aux(mat_ZZ_p& X, const mat_ZZ_p& A, const mat_ZZ_p& B)
{
   long n = A.NumRows();
   long l = A.NumCols();
   long m = B.NumCols();

   if (l != B.NumRows())
      LogicError("matrix mul: dimension mismatch");

   X.SetDims(n, m);

   ZZ_pContext context;
   context.save();

   long sz = ZZ_p::ModulusSize();
   bool seq = (double(n)*double(l)*double(m)*double(sz)*double(sz) < PAR_THRESH);

   NTL_GEXEC_RANGE(seq, m, first, last)
   NTL_IMPORT(n)
   NTL_IMPORT(l)
   NTL_IMPORT(m)

   context.restore();

   long i, j, k;
   ZZ acc, tmp;

   vec_ZZ_p B_col;
   B_col.SetLength(l);

   for (j = first; j < last; j++) {
      for (k = 0; k < l; k++) B_col[k] = B[k][j];

      for (i = 0; i < n; i++) {
         clear(acc);
         for (k = 0; k < l; k++) {
            mul(tmp, rep(A[i][k]), rep(B_col[k]));
            add(acc, acc, tmp);
         }
         conv(X[i][j], acc);
      }
   }

   NTL_GEXEC_RANGE_END
}

void plain_mul(mat_ZZ_p& X, const mat_ZZ_p& A, const mat_ZZ_p& B)
{
   if (&X == &A || &X == &B) {
      mat_ZZ_p tmp;
      plain_mul_aux(tmp, A, B);
      X = tmp;
   }
   else
      plain_mul_aux(X, A, B);
}

// X = A*transpose(B)

void plain_mul_transpose_aux(mat_ZZ_p& X, const mat_ZZ_p& A, const mat_ZZ_p& B)
{
   long n = A.NumRows();
   long l = A.NumCols();
   long m = B.NumRows();

   if (l != B.NumCols())
      LogicError("matrix mul: dimension mismatch");

   X.SetDims(n, m);

   ZZ_pContext context;
   context.save();

   long sz = ZZ_p::ModulusSize();
   bool seq = (double(n)*double(l)*double(m)*double(sz)*double(sz) < PAR_THRESH);

   NTL_GEXEC_RANGE(seq, m, first, last)
   NTL_IMPORT(n)
   NTL_IMPORT(l)
   NTL_IMPORT(m)

   context.restore();

   long i, j, k;
   ZZ acc, tmp;

   for (j = first; j < last; j++) {
      const ZZ_p *B_col = B[j].elts();

      for (i = 0; i < n; i++) {
         clear(acc);
         for (k = 0; k < l; k++) {
            mul(tmp, rep(A[i][k]), rep(B_col[k]));
            add(acc, acc, tmp);
         }
         conv(X[i][j], acc);
      }
   }

   NTL_GEXEC_RANGE_END
}

void plain_mul_transpose(mat_ZZ_p& X, const mat_ZZ_p& A, const mat_ZZ_p& B)
{
   if (&X == &A || &X == &B) {
      mat_ZZ_p tmp;
      plain_mul_transpose_aux(tmp, A, B);
      X = tmp;
   }
   else
      plain_mul_transpose_aux(X, A, B);
}

// ***************** Multi-modular Matrix Multiplication *************

struct mat_ZZ_p_crt_rep {

   Vec< Mat<MatPrime_residue_t> > rep;

};

static
const MatPrime_crt_helper& get_MatPrime_crt_helper_info()
{
   do {
      Lazy<MatPrime_crt_helper,ZZ_pInfoT::MatPrime_crt_helper_deleter_policy>::Builder
         builder(ZZ_pInfo->MatPrime_crt_helper_info);
      if (!builder()) break;

      UniquePtr<MatPrime_crt_helper,ZZ_pInfoT::MatPrime_crt_helper_deleter_policy> p;
      p.make();
      build(*p, ZZ_pInfo->p);
      builder.move(p);
   } while (0);

   return *ZZ_pInfo->MatPrime_crt_helper_info;
}

static
void RawConvert(Mat<zz_p>& X, const Mat<MatPrime_residue_t>& A)
{
   long n = A.NumRows();
   long m = A.NumCols();

   X.SetDims(n, m);
   for (long i = 0; i < n; i++) {
      const MatPrime_residue_t *Ai = A[i].elts();
      zz_p *Xi = X[i].elts();
      for (long j = 0; j < m; j++)
         Xi[j].LoopHole() = Ai[j];
   }
}

static
void RawConvertTranspose(Mat<zz_p>& X, const Mat<MatPrime_residue_t>& A)
{
   long n = A.NumRows();
   long m = A.NumCols();

   X.SetDims(m, n);
   for (long i = 0; i < n; i++) {
      const MatPrime_residue_t *Ai = A[i].elts();
      for (long j = 0; j < m; j++)
         X[j][i] = Ai[j];
   }
}

static
void RawConvert(Mat<MatPrime_residue_t>& X, const Mat<zz_p>& A)
{
   long n = A.NumRows();
   long m = A.NumCols();

   X.SetDims(n, m);
   for (long i = 0; i < n; i++) {
      const zz_p *Ai = A[i].elts();
      MatPrime_residue_t *Xi = X[i].elts();
      for (long j = 0; j < m; j++)
         Xi[j] = rep(Ai[j]);
   }
}

#define CRT_BLK (8)

void to_mat_ZZ_p_crt_rep(mat_ZZ_p_crt_rep& X, const mat_ZZ_p& A)
{
   long n = A.NumRows();
   long m = A.NumCols();

   const MatPrime_crt_helper& H = get_MatPrime_crt_helper_info();
   long nprimes = H.GetNumPrimes();

   if (NTL_OVERFLOW(nprimes, CRT_BLK, 0))
      ResourceError("overflow"); // this is pretty academic

   X.rep.SetLength(nprimes);
   for (long k = 0; k < nprimes; k++) X.rep[k].SetDims(n, m);

   ZZ_pContext context;
   context.save();

   bool seq = (double(n)*double(m)*H.GetCost() < PAR_THRESH);

   // FIXME: right now, we just partition the rows, but if
   // #cols > #rows, we should perhaps partition the cols
   NTL_GEXEC_RANGE(seq, n, first, last)
   NTL_IMPORT(n)
   NTL_IMPORT(m)
   NTL_IMPORT(nprimes)

   context.restore();

   MatPrime_crt_helper_scratch scratch;
   Vec<MatPrime_residue_t> remainders_store;
   remainders_store.SetLength(nprimes*CRT_BLK);
   MatPrime_residue_t *remainders = remainders_store.elts();

   for (long i = first; i < last; i++) {
      const ZZ_p *a = A[i].elts();

      long jj = 0;
      for (; jj <= m-CRT_BLK; jj += CRT_BLK) {
         for (long j = 0; j < CRT_BLK; j++)
            reduce(H, rep(a[jj+j]), remainders + j*nprimes, scratch);
         for (long k = 0; k < nprimes; k++) {
            MatPrime_residue_t *x = X.rep[k][i].elts();
            for (long j = 0; j < CRT_BLK; j++)
               x[jj+j] = remainders[j*nprimes+k];
         }
      }
      if (jj < m) {
         for (long j = 0; j < m-jj; j++)
            reduce(H, rep(a[jj+j]), remainders + j*nprimes, scratch);
         for (long k = 0; k < nprimes; k++) {
            MatPrime_residue_t *x = X.rep[k][i].elts();
            for (long j = 0; j < m-jj; j++)
               x[jj+j] = remainders[j*nprimes+k];
         }
      }
   }

   NTL_GEXEC_RANGE_END
}

void from_mat_ZZ_p_crt_rep(const mat_ZZ_p_crt_rep& X, mat_ZZ_p& A)
{
   long n = X.rep[0].NumRows();
   long m = X.rep[0].NumCols();

   const MatPrime_crt_helper& H = get_MatPrime_crt_helper_info();
   long nprimes = H.GetNumPrimes();

   if (NTL_OVERFLOW(nprimes, CRT_BLK, 0))
      ResourceError("overflow"); // this is pretty academic

   A.SetDims(n, m);

   ZZ_pContext context;
   context.save();

   bool seq = (double(n)*double(m)*H.GetCost() < PAR_THRESH);

   // FIXME: right now, we just partition the rows, but if
   // #cols > #rows, we should perhaps partition the cols
   NTL_GEXEC_RANGE(seq, n, first, last)
   NTL_IMPORT(n)
   NTL_IMPORT(m)
   NTL_IMPORT(nprimes)

   context.restore();

   MatPrime_crt_helper_scratch scratch;
   Vec<MatPrime_residue_t> remainders_store;
   remainders_store.SetLength(nprimes*CRT_BLK);
   MatPrime_residue_t *remainders = remainders_store.elts();

   for (long i = first; i < last; i++) {
      ZZ_p *a = A[i].elts();

      long jj = 0;
      for (; jj <= m-CRT_BLK; jj += CRT_BLK) {
         for (long k = 0; k < nprimes; k++) {
            const MatPrime_residue_t *x = X.rep[k][i].elts();
            for (long j = 0; j < CRT_BLK; j++)
               remainders[j*nprimes+k] = x[jj+j];
         }
         for (long j = 0; j < CRT_BLK; j++)
            reconstruct(H, a[jj+j].LoopHole(), remainders + j*nprimes, scratch);
      }
      if (jj < m) {
         for (long k = 0; k < nprimes; k++) {
            const MatPrime_residue_t *x = X.rep[k][i].elts();
            for (long j = 0; j < m-jj; j++)
               remainders[j*nprimes+k] = x[jj+j];
         }
         for (long j = 0; j < m-jj; j++)
            reconstruct(H, a[jj+j].LoopHole(), remainders + j*nprimes, scratch);
      }
   }

   NTL_GEXEC_RANGE_END
}

void mul(mat_ZZ_p_crt_rep& X, const mat_ZZ_p_crt_rep& A, const mat_ZZ_p_crt_rep& B)
{
   long nprimes = A.rep.length();

   long n = A.rep[0].NumRows();
   long l = A.rep[0].NumCols();
   long m = B.rep[0].NumCols();

   X.rep.SetLength(nprimes);
   for (long k = 0; k < nprimes; k++) X.rep[k].SetDims(n, m);

   bool seq = (double(n)*double(l)*double(m)*double(nprimes) < PAR_THRESH);

   NTL_GEXEC_RANGE(seq, nprimes, first, last)
   NTL_IMPORT(n)
   NTL_IMPORT(l)
   NTL_IMPORT(m)

   zz_pPush push;

   Mat<zz_p> x, a, b;
   x.SetDims(n, m);
   a.SetDims(n, l);
   b.SetDims(l, m);

   for (long k = first; k < last; k++) {
      RestoreMatPrime(k);
      RawConvert(a, A.rep[k]);
      RawConvert(b, B.rep[k]);
      mul(x, a, b);
      RawConvert(X.rep[k], x);
   }

   NTL_GEXEC_RANGE_END
}

// X = A*transpose(B)
void mul_transpose(mat_ZZ_p_crt_rep& X, const mat_ZZ_p_crt_rep& A, const mat_ZZ_p_crt_rep& B)
{
   long nprimes = A.rep.length();

   long n = A.rep[0].NumRows();
   long l = A.rep[0].NumCols();
   long m = B.rep[0].NumRows();

   X.rep.SetLength(nprimes);
   for (long k = 0; k < nprimes; k++) X.rep[k].SetDims(n, m);

   bool seq = (double(n)*double(l)*double(m)*double(nprimes) < PAR_THRESH);

   NTL_GEXEC_RANGE(seq, nprimes, first, last)
   NTL_IMPORT(n)
   NTL_IMPORT(l)
   NTL_IMPORT(m)

   zz_pPush push;

   Mat<zz_p> x, a, b;
   x.SetDims(n, m);
   a.SetDims(n, l);
   b.SetDims(l, m);

   for (long k = first; k < last; k++) {
      RestoreMatPrime(k);
      RawConvert(a, A.rep[k]);
      RawConvertTranspose(b, B.rep[k]);
      mul(x, a, b);
      RawConvert(X.rep[k], x);
   }

   NTL_GEXEC_RANGE_END
}

void multi_modular_mul(mat_ZZ_p& X, const mat_ZZ_p& A, const mat_ZZ_p& B)
{
   long l = A.NumCols();

   if (l != B.NumRows())
      LogicError("matrix mul: dimension mismatch");

   if (l > NTL_MatPrimeLimit)
      ResourceError("matrix mul: dimension too large");

   mat_ZZ_p_crt_rep x, a, b;

   to_mat_ZZ_p_crt_rep(a, A);
   to_mat_ZZ_p_crt_rep(b, B);
   mul(x, a, b);
   from_mat_ZZ_p_crt_rep(x, X);
}

void multi_modular_mul(mat_ZZ_p& X, const mat_ZZ_p& A, const mat_ZZ_p_crt_rep& B)
{
   long l = A.NumCols();

   if (l != B.rep[0].NumRows())
      LogicError("matrix mul: dimension mismatch");

   if (l > NTL_MatPrimeLimit)
      ResourceError("matrix mul: dimension too large");

   mat_ZZ_p_crt_rep x, a;

   to_mat_ZZ_p_crt_rep(a, A);
   mul(x, a, B);
   from_mat_ZZ_p_crt_rep(x, X);
}

void multi_modular_mul_transpose(mat_ZZ_p& X, const mat_ZZ_p& A, const mat_ZZ_p_crt_rep& B)
{
   long l = A.NumCols();

   if (l != B.rep[0].NumCols())
      LogicError("matrix mul: dimension mismatch");

   if (l > NTL_MatPrimeLimit)
      ResourceError("matrix mul: dimension too large");

   mat_ZZ_p_crt_rep x, a;

   to_mat_ZZ_p_crt_rep(a, A);
   mul_transpose(x, a, B);
   from_mat_ZZ_p_crt_rep(x, X);
}

// ******************** mat_ZZ_p_opaque implementation ************

// This could (and maybe eventually will) be implemented using
// derived types, if we want run-time polymorphism...we'll see

struct mat_ZZ_p_opaque_body_crt : mat_ZZ_p_opaque_body {
   mat_ZZ_p_crt_rep body;

   mat_ZZ_p_opaque_body* clone() const
   {
      return MakeRaw<mat_ZZ_p_opaque_body_crt>(*this);
   }

   long NumRows() const
   {
      return body.rep.length() == 0 ? 0 : body.rep[0].NumRows();
   }

   long NumCols() const
   {
      return body.rep.length() == 0 ? 0 : body.rep[0].NumCols();
   }

   void mul(mat_ZZ_p& X, const mat_ZZ_p& A) const
   {
      multi_modular_mul(X, A, body);
   }

   void mul_transpose(mat_ZZ_p& X, const mat_ZZ_p& A) const
   {
      multi_modular_mul_transpose(X, A, body);
   }

};

struct mat_ZZ_p_opaque_body_plain : mat_ZZ_p_opaque_body {
   mat_ZZ_p body;

   mat_ZZ_p_opaque_body* clone() const
   {
      return MakeRaw<mat_ZZ_p_opaque_body_plain>(*this);
   }

   long NumRows() const
   {
      return body.NumRows();
   }

   long NumCols() const
   {
      return body.NumCols();
   }

   void mul(mat_ZZ_p& X, const mat_ZZ_p& A) const
   {
      plain_mul(X, A, body);
   }

   void mul_transpose(mat_ZZ_p& X, const mat_ZZ_p& A) const
   {
      plain_mul_transpose(X, A, body);
   }

};

// This is a "factory" method that makes a mat_ZZ_p_opaque_body
// from a matrix A.  The matrix A is destroyed in the process.

mat_ZZ_p_opaque_body *mat_ZZ_p_opaque_body_move(mat_ZZ_p& A)
{
   if (NTL_USE_MM_MATMUL && A.NumRows() >= 16 && A.NumCols() >= 16) {
      UniquePtr<mat_ZZ_p_opaque_body_crt> tmp;
      tmp.make();
      to_mat_ZZ_p_crt_rep(tmp->body, A);
      A.kill();
      return tmp.release();
   }
   else {
      UniquePtr<mat_ZZ_p_opaque_body_plain> tmp;
      tmp.make();
      tmp->body.move(A);
      return tmp.release();
   }
}

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

void mul(mat_ZZ_p& X, const mat_ZZ_p& A, const mat_ZZ_p& B)
{
   long n = A.NumRows();
   long l = A.NumCols();
   long m = B.NumCols();

   if (l != B.NumRows()) LogicError("matrix mul: dimension mismatch");

   if (NTL_USE_MM_MATMUL && n >= 24 && l >= 24 && m >= 24)
      multi_modular_mul(X, A, B);
   else
      plain_mul(X, A, B);
}

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

void add(mat_ZZ_p& X, const mat_ZZ_p& A, const mat_ZZ_p& B)
{
   long n = A.NumRows();
   long m = A.NumCols();

   if (B.NumRows() != n || B.NumCols() != m)
      LogicError("matrix add: dimension mismatch");

   X.SetDims(n, m);

   long i, j;
   for (i = 1; i <= n; i++)
      for (j = 1; j <= m; j++)
         add(X(i,j), A(i,j), B(i,j));
}

void sub(mat_ZZ_p& X, const mat_ZZ_p& A, const mat_ZZ_p& B)
{
   long n = A.NumRows();
   long m = A.NumCols();

   if (B.NumRows() != n || B.NumCols() != m)
      LogicError("matrix sub: dimension mismatch");

   X.SetDims(n, m);

   long i, j;
   for (i = 1; i <= n; i++)
      for (j = 1; j <= m; j++)
         sub(X(i,j), A(i,j), B(i,j));
}

void negate(mat_ZZ_p& X, const mat_ZZ_p& A)
{
   long n = A.NumRows();
   long m = A.NumCols();

   X.SetDims(n, m);

   long i, j;
   for (i = 1; i <= n; i++)
      for (j = 1; j <= m; j++)
         negate(X(i,j), A(i,j));
}

static
void mul_aux(vec_ZZ_p& x, const mat_ZZ_p& A, const vec_ZZ_p& b)
{
   long n = A.NumRows();
   long l = A.NumCols();

   if (l != b.length())
      LogicError("matrix mul: dimension mismatch");

   x.SetLength(n);

   long i, k;
   ZZ acc, tmp;

   for (i = 1; i <= n; i++) {
      clear(acc);
      for (k = 1; k <= l; k++) {
         mul(tmp, rep(A(i,k)), rep(b(k)));
         add(acc, acc, tmp);
      }
      conv(x(i), acc);
   }
}

void mul(vec_ZZ_p& x, const mat_ZZ_p& A, const vec_ZZ_p& b)
{
   if (&b == &x || A.position1(x) != -1) {
      vec_ZZ_p tmp;
      mul_aux(tmp, A, b);
      x = tmp;
   }
   else
      mul_aux(x, A, b);
}

static
void mul_aux(vec_ZZ_p& x, const vec_ZZ_p& a, const mat_ZZ_p& B)
{
   long n = B.NumRows();
   long l = B.NumCols();

   if (n != a.length())
      LogicError("matrix mul: dimension mismatch");

   x.SetLength(l);

   long i, k;
   ZZ acc, tmp;

   for (i = 1; i <= l; i++) {
      clear(acc);
      for (k = 1; k <= n; k++) {
         mul(tmp, rep(a(k)), rep(B(k,i)));
         add(acc, acc, tmp);
      }
      conv(x(i), acc);
   }
}

void mul(vec_ZZ_p& x, const vec_ZZ_p& a, const mat_ZZ_p& B)
{
   if (&a == &x) {
      vec_ZZ_p tmp;
      mul_aux(tmp, a, B);
      x = tmp;
   }
   else
      mul_aux(x, a, B);
}

void ident(mat_ZZ_p& X, long n)
{
   X.SetDims(n, n);
   long i, j;

   for (i = 1; i <= n; i++)
      for (j = 1; j <= n; j++)
         if (i == j)
            set(X(i, j));
         else
            clear(X(i, j));
}

void determinant(ZZ_p& d, const mat_ZZ_p& M_in)
{
   ZZ t1, t2;

   const ZZ& p = ZZ_p::modulus();

   long n = M_in.NumRows();

   if (M_in.NumCols() != n)
      LogicError("determinant: nonsquare matrix");

   if (n == 0) {
      set(d);
      return;
   }

   Vec<ZZVec> M;
   sqr(t1, p);
   mul(t1, t1, n);

   M.SetLength(n);
   for (long i = 0; i < n; i++) {
      M[i].SetSize(n, t1.size());
      for (long j = 0; j < n; j++)
         M[i][j] = rep(M_in[i][j]);
   }

   ZZ det;
   set(det);

   for (long k = 0; k < n; k++) {
      long pos = -1;
      for (long i = k; i < n; i++) {
         rem(t1, M[i][k], p);
         M[i][k] = t1;
         if (pos == -1 && !IsZero(t1))
            pos = i;
      }

      if (pos != -1) {
         if (k != pos) {
            swap(M[pos], M[k]);
            NegateMod(det, det, p);
         }

         MulMod(det, det, M[k][k], p);

         // make M[k, k] == -1 mod p, and make row k reduced

         InvMod(t1, M[k][k], p);
         NegateMod(t1, t1, p);
         for (long j = k+1; j < n; j++) {
            rem(t2, M[k][j], p);
            MulMod(M[k][j], t2, t1, p);
         }

         bool seq =
            double(n-(k+1))*(n-(k+1))*double(p.size())*double(p.size()) < PAR_THRESH;
         NTL_GEXEC_RANGE(seq, n-(k+1), first, last)
         NTL_IMPORT(n)
         NTL_IMPORT(k)

         ZZ t1, t2;

         for (long ii = first; ii < last; ii++) {
            long i = ii + k+1;

            // M[i] = M[i] + M[k]*M[i,k]

            t1 = M[i][k];   // this is already reduced

            ZZ *x = M[i].elts() + (k+1);
            ZZ *y = M[k].elts() + (k+1);

            for (long j = k+1; j < n; j++, x++, y++) {
               // *x = *x + (*y)*t1

               mul(t2, *y, t1);
               add(*x, *x, t2);
            }
         }

         NTL_GEXEC_RANGE_END

      }
      else {
         clear(d);
         return;
      }
   }

   conv(d, det);
}

long IsIdent(const mat_ZZ_p& A, long n)
{
   if (A.NumRows() != n || A.NumCols() != n)
      return 0;

   long i, j;

   for (i = 1; i <= n; i++)
      for (j = 1; j <= n; j++)
         if (i != j) {
            if (!IsZero(A(i, j))) return 0;
         }
         else {
            if (!IsOne(A(i, j))) return 0;
         }

   return 1;
}

void transpose(mat_ZZ_p& X, const mat_ZZ_p& A)
{
   long n = A.NumRows();
   long m = A.NumCols();

   long i, j;

   if (&X == & A) {
      if (n == m)
         for (i = 1; i <= n; i++)
            for (j = i+1; j <= n; j++)
               swap(X(i, j), X(j, i));
      else {
         mat_ZZ_p tmp;
         tmp.SetDims(m, n);
         for (i = 1; i <= n; i++)
            for (j = 1; j <= m; j++)
               tmp(j, i) = A(i, j);
         X.kill();
         X = tmp;
      }
   }
   else {
      X.SetDims(m, n);
      for (i = 1; i <= n; i++)
         for (j = 1; j <= m; j++)
            X(j, i) = A(i, j);
   }
}

static
void solve_impl(ZZ_p& d, vec_ZZ_p& X, const mat_ZZ_p& A, const vec_ZZ_p& b, bool trans)

{
   long n = A.NumRows();
   if (A.NumCols() != n)
      LogicError("solve: nonsquare matrix");

   if (b.length() != n)
      LogicError("solve: dimension mismatch");

   if (n == 0) {
      set(d);
      X.SetLength(0);
      return;
   }

   ZZ t1, t2;

   const ZZ& p = ZZ_p::modulus();

   Vec<ZZVec> M;
   sqr(t1, p);
   mul(t1, t1, n);

   M.SetLength(n);

   for (long i = 0; i < n; i++) {
      M[i].SetSize(n+1, t1.size());

      if (trans)
         for (long j = 0; j < n; j++) M[i][j] = rep(A[j][i]);
      else
         for (long j = 0; j < n; j++) M[i][j] = rep(A[i][j]);

      M[i][n] = rep(b[i]);
   }

   ZZ det;
   set(det);

   for (long k = 0; k < n; k++) {
      long pos = -1;
      for (long i = k; i < n; i++) {
         rem(t1, M[i][k], p);
         M[i][k] = t1;
         if (pos == -1 && !IsZero(t1)) {
            pos = i;
         }
      }

      if (pos != -1) {
         if (k != pos) {
            swap(M[pos], M[k]);
            NegateMod(det, det, p);
         }

         MulMod(det, det, M[k][k], p);

         // make M[k, k] == -1 mod p, and make row k reduced

         InvMod(t1, M[k][k], p);
         NegateMod(t1, t1, p);
         for (long j = k+1; j <= n; j++) {
            rem(t2, M[k][j], p);
            MulMod(M[k][j], t2, t1, p);
         }

         bool seq =
            double(n-(k+1))*(n-(k+1))*double(p.size())*double(p.size()) < PAR_THRESH;
         NTL_GEXEC_RANGE(seq, n-(k+1), first, last)
         NTL_IMPORT(n)
         NTL_IMPORT(k)

         ZZ t1, t2;

         for (long ii = first; ii < last; ii++) {
            long i = ii + k+1;

            // M[i] = M[i] + M[k]*M[i,k]

            t1 = M[i][k];   // this is already reduced

            ZZ *x = M[i].elts() + (k+1);
            ZZ *y = M[k].elts() + (k+1);

            for (long j = k+1; j <= n; j++, x++, y++) {
               // *x = *x + (*y)*t1

               mul(t2, *y, t1);
               add(*x, *x, t2);
            }
         }

         NTL_GEXEC_RANGE_END

      }
      else {
         clear(d);
         return;
      }
   }

   X.SetLength(n);
   for (long i = n-1; i >= 0; i--) {
      clear(t1);
      for (long j = i+1; j < n; j++) {
         mul(t2, rep(X[j]), M[i][j]);
         add(t1, t1, t2);
      }
      sub(t1, t1, M[i][n]);
      conv(X[i], t1);
   }

   conv(d, det);
}

void solve(ZZ_p& d, vec_ZZ_p& x, const mat_ZZ_p& A, const vec_ZZ_p& b)
{
   solve_impl(d, x, A, b, true);
}

void solve(ZZ_p& d, const mat_ZZ_p& A, vec_ZZ_p& x,  const vec_ZZ_p& b)
{
   solve_impl(d, x, A, b, false);
}

void inv(ZZ_p& d, mat_ZZ_p& X, const mat_ZZ_p& A)
{
   long n = A.NumRows();

   if (A.NumCols() != n)
      LogicError("inv: nonsquare matrix");

   if (n == 0) {
      set(d);
      X.SetDims(0, 0);
      return;
   }

   const ZZ& p = ZZ_p::modulus();

   ZZ t1, t2;
   ZZ pivot;
   ZZ pivot_inv;

   Vec<ZZVec> M;
   // scratch space

   sqr(t1, p);
   mul(t1, t1, n);

   M.SetLength(n);
   for (long i = 0; i < n; i++) {
      M[i].SetSize(n, t1.size());
      for (long j = 0; j < n; j++)
         M[i][j] = rep(A[i][j]);
   }

   ZZ det;
   det = 1;

   Vec<long> P;
   P.SetLength(n);
   for (long k = 0; k < n; k++) P[k] = k;
   // records swap operations

   bool seq =
      double(n)*double(n)*double(p.size())*double(p.size()) < PAR_THRESH;

   bool pivoting = false;

   for (long k = 0; k < n; k++) {

      long pos = -1;

      for (long i = k; i < n; i++) {
         rem(pivot, M[i][k], p);
         if (pivot != 0) {
            InvMod(pivot_inv, pivot, p);
            pos = i;
            break;
         }
      }

      if (pos != -1) {
         if (k != pos) {
            swap(M[pos], M[k]);
            NegateMod(det, det, p);
            P[k] = pos;
            pivoting = true;
         }

         MulMod(det, det, pivot, p);

         {
            // multiply row k by pivot_inv
            ZZ *y = &M[k][0];
            for (long j = 0; j < n; j++) {
               rem(t2, y[j], p);
               MulMod(y[j], t2, pivot_inv, p);
            }
            y[k] = pivot_inv;
         }

         NTL_GEXEC_RANGE(seq, n, first, last)
         NTL_IMPORT(n)
         NTL_IMPORT(k)

         ZZ *y = &M[k][0];
         ZZ t1, t2;

         for (long i = first; i < last; i++) {
            if (i == k) continue; // skip row k

            ZZ *x = &M[i][0];
            rem(t1, x[k], p);
            NegateMod(t1, t1, p);
            x[k] = 0;
            if (t1 == 0) continue;

            // add t1 * row k to row i
            for (long j = 0; j < n; j++) {
               mul(t2, y[j], t1);
               add(x[j], x[j], t2);
            }
         }
         NTL_GEXEC_RANGE_END
      }
      else {
         clear(d);
         return;
      }
   }

   if (pivoting) {
      // pivot colums, using reverse swap sequence

      for (long i = 0; i < n; i++) {
         ZZ *x = &M[i][0];

         for (long k = n-1; k >= 0; k--) {
            long pos = P[k];
            if (pos != k) swap(x[pos], x[k]);
         }
      }
   }

   X.SetDims(n, n);
   for (long i = 0; i < n; i++)
      for (long j = 0; j < n; j++)
         conv(X[i][j], M[i][j]);

   conv(d, det);
}

long gauss(mat_ZZ_p& M_in, long w)
{
   ZZ t1, t2;
   ZZ piv;

   long n = M_in.NumRows();
   long m = M_in.NumCols();

   if (w < 0 || w > m)
      LogicError("gauss: bad args");

   const ZZ& p = ZZ_p::modulus();

   Vec<ZZVec> M;
   sqr(t1, p);
   mul(t1, t1, n);

   M.SetLength(n);
   for (long i = 0; i < n; i++) {
      M[i].SetSize(m, t1.size());
      for (long j = 0; j < m; j++) {
         M[i][j] = rep(M_in[i][j]);
      }
   }

   long l = 0;
   for (long k = 0; k < w && l < n; k++) {

      long pos = -1;
      for (long i = l; i < n; i++) {
         rem(t1, M[i][k], p);
         M[i][k] = t1;
         if (pos == -1 && !IsZero(t1)) {
            pos = i;
         }
      }

      if (pos != -1) {
         swap(M[pos], M[l]);

         InvMod(piv, M[l][k], p);
         NegateMod(piv, piv, p);

         for (long j = k+1; j < m; j++) {
            rem(M[l][j], M[l][j], p);
         }

         bool seq =
            double(n-(l+1))*double(m-(k+1))*double(p.size())*double(p.size()) < PAR_THRESH;

         NTL_GEXEC_RANGE(seq, n-(l+1), first, last)
         NTL_IMPORT(m)
         NTL_IMPORT(k)
         NTL_IMPORT(l)

         ZZ t1, t2;

         for (long ii = first; ii < last; ii++) {
            long i = ii + l+1;

            // M[i] = M[i] + M[l]*M[i,k]*piv

            MulMod(t1, M[i][k], piv, p);

            clear(M[i][k]);

            ZZ *x = M[i].elts() + (k+1);
            ZZ *y = M[l].elts() + (k+1);

            for (long j = k+1; j < m; j++, x++, y++) {
               // *x = *x + (*y)*t1

               mul(t2, *y, t1);
               add(t2, t2, *x);
               *x = t2;
            }
         }

         NTL_GEXEC_RANGE_END

         l++;
      }
   }

   for (long i = 0; i < n; i++)
      for (long j = 0; j < m; j++)
         conv(M_in[i][j], M[i][j]);

   return l;
}

long gauss(mat_ZZ_p& M)
{
   return gauss(M, M.NumCols());
}

void image(mat_ZZ_p& X, const mat_ZZ_p& A)
{
   mat_ZZ_p M;
   M = A;
   long r = gauss(M);
   M.SetDims(r, M.NumCols());
   X = M;
}

void kernel(mat_ZZ_p& X, const mat_ZZ_p& A)
{
   long m = A.NumRows();
   long n = A.NumCols();

   const ZZ& p = ZZ_p::modulus();

   mat_ZZ_p M;

   transpose(M, A);
   long r = gauss(M);

   if (r == 0) {
      ident(X, m);
      return;
   }

   X.SetDims(m-r, m);

   if (m-r == 0 || m == 0) return;

   Vec<long> D;
   D.SetLength(m);
   for (long j = 0; j < m; j++) D[j] = -1;

   Vec<ZZ_p> inverses;
   inverses.SetLength(m);

   for (long i = 0, j = -1; i < r; i++) {
      do {
         j++;
      } while (IsZero(M[i][j]));

      D[j] = i;
      inv(inverses[j], M[i][j]);
   }

   bool seq =
      double(m-r)*double(r)*double(r)*double(p.size())*double(p.size()) < PAR_THRESH;

   NTL_GEXEC_RANGE(seq, m-r, first, last)
   NTL_IMPORT(m)
   NTL_IMPORT(r)

   ZZ t1, t2;
   ZZ_p T3;

   for (long k = first; k < last; k++) {
      vec_ZZ_p& v = X[k];
      long pos = 0;
      for (long j = m-1; j >= 0; j--) {
         if (D[j] == -1) {
            if (pos == k)
               set(v[j]);
            else
               clear(v[j]);
            pos++;
         }
         else {
            long i = D[j];

            clear(t1);

            for (long s = j+1; s < m; s++) {
               mul(t2, rep(v[s]), rep(M[i][s]));
               add(t1, t1, t2);
            }

            conv(T3, t1);
            mul(T3, T3, inverses[j]);
            negate(v[j], T3);
         }
      }
   }

   NTL_GEXEC_RANGE_END
}

void mul(mat_ZZ_p& X, const mat_ZZ_p& A, const ZZ_p& b_in)
{
   NTL_ZZ_pRegister(b);
   b = b_in;
   long n = A.NumRows();
   long m = A.NumCols();

   X.SetDims(n, m);

   long i, j;
   for (i = 0; i < n; i++)
      for (j = 0; j < m; j++)
         mul(X[i][j], A[i][j], b);
}

void mul(mat_ZZ_p& X, const mat_ZZ_p& A, long b_in)
{
   NTL_ZZ_pRegister(b);
   b = b_in;
   long n = A.NumRows();
   long m = A.NumCols();

   X.SetDims(n, m);

   long i, j;
   for (i = 0; i < n; i++)
      for (j = 0; j < m; j++)
         mul(X[i][j], A[i][j], b);
}

void diag(mat_ZZ_p& X, long n, const ZZ_p& d_in)
{
   ZZ_p d = d_in;
   X.SetDims(n, n);
   long i, j;

   for (i = 1; i <= n; i++)
      for (j = 1; j <= n; j++)
         if (i == j)
            X(i, j) = d;
         else
            clear(X(i, j));
}

long IsDiag(const mat_ZZ_p& A, long n, const ZZ_p& d)
{
   if (A.NumRows() != n || A.NumCols() != n)
      return 0;

   long i, j;

   for (i = 1; i <= n; i++)
      for (j = 1; j <= n; j++)
         if (i != j) {
            if (!IsZero(A(i, j))) return 0;
         }
         else {
            if (A(i, j) != d) return 0;
         }

   return 1;
}

long IsZero(const mat_ZZ_p& a)
{
   long n = a.NumRows();
   long i;

   for (i = 0; i < n; i++)
      if (!IsZero(a[i]))
         return 0;

   return 1;
}

void clear(mat_ZZ_p& x)
{
   long n = x.NumRows();
   long i;
   for (i = 0; i < n; i++)
      clear(x[i]);
}

mat_ZZ_p operator+(const mat_ZZ_p& a, const mat_ZZ_p& b)
{
   mat_ZZ_p res;
   add(res, a, b);
   NTL_OPT_RETURN(mat_ZZ_p, res);
}

mat_ZZ_p operator*(const mat_ZZ_p& a, const mat_ZZ_p& b)
{
   mat_ZZ_p res;
   mul(res, a, b);
   NTL_OPT_RETURN(mat_ZZ_p, res);
}

mat_ZZ_p operator-(const mat_ZZ_p& a, const mat_ZZ_p& b)
{
   mat_ZZ_p res;
   sub(res, a, b);
   NTL_OPT_RETURN(mat_ZZ_p, res);
}

mat_ZZ_p operator-(const mat_ZZ_p& a)
{
   mat_ZZ_p res;
   negate(res, a);
   NTL_OPT_RETURN(mat_ZZ_p, res);
}

vec_ZZ_p operator*(const mat_ZZ_p& a, const vec_ZZ_p& b)
{
   vec_ZZ_p res;
   mul_aux(res, a, b);
   NTL_OPT_RETURN(vec_ZZ_p, res);
}

vec_ZZ_p operator*(const vec_ZZ_p& a, const mat_ZZ_p& b)
{
   vec_ZZ_p res;
   mul_aux(res, a, b);
   NTL_OPT_RETURN(vec_ZZ_p, res);
}

void inv(mat_ZZ_p& X, const mat_ZZ_p& A)
{
   ZZ_p d;
   inv(d, X, A);
   if (d == 0) ArithmeticError("inv: non-invertible matrix");
}

void power(mat_ZZ_p& X, const mat_ZZ_p& A, const ZZ& e)
{
   if (A.NumRows() != A.NumCols()) LogicError("power: non-square matrix");

   if (e == 0) {
      ident(X, A.NumRows());
      return;
   }

   mat_ZZ_p T1, T2;
   long i, k;

   k = NumBits(e);
   T1 = A;

   for (i = k-2; i >= 0; i--) {
      sqr(T2, T1);
      if (bit(e, i))
         mul(T1, T2, A);
      else
         T1 = T2;
   }

   if (e < 0)
      inv(X, T1);
   else
      X = T1;
}

void random(mat_ZZ_p& x, long n, long m)
{
   x.SetDims(n, m);
   for (long i = 0; i < n; i++) random(x[i], m);
}

NTL_END_IMPL


/* ----------------------------------------------------------------------- *//**
 *
 * @file metric.cpp
 *
 * @brief Metric operations
 *
 *//* ----------------------------------------------------------------------- */

#include <dbconnector/dbconnector.hpp>
#include <limits>
#include <string>
#include <vector>
#include <set>
#include <algorithm>
#include <sstream>
#include <cmath>
#include <boost/algorithm/string.hpp>

#include "metric.hpp"

using std::string;

namespace madlib {

using namespace dbconnector::postgres;

namespace modules {

namespace linalg {

namespace {

template <class TupleType>
struct ReverseLexicographicComparator {
    /**
     * @brief Return true if the first argument is less than the second
     */
    bool operator()(const TupleType& inTuple1, const TupleType& inTuple2) {
        // This could be a real reverse lexicographic comparator in C++11, but
        // lacking variadic template arguments, we simply pretend here that
        // all tuples contain only 2 elements.
        return std::get<1>(inTuple1) < std::get<1>(inTuple2) ||
            (std::get<1>(inTuple1) == std::get<1>(inTuple2) &&
                std::get<0>(inTuple1) < std::get<0>(inTuple2));
    }
};

} // anonymous namespace

/**
 * @brief Compute the k columns of a matrix that are closest to a vector
 *
 * @tparam DistanceFunction Type of the distance function. This can be a
 *     function pointer, or a class with method <tt>operator()</tt> (like
 *     \c FunctionHandle)
 * @tparam NumClosestColumns The number of closest columns to compute. We assume
 *     that \c outIndicesAndDistances is of size \c NumClosestColumns.
 *
 * @param inMatrix Matrix \f$ M \f$
 * @param inVector Vector \f$ \vec x \f$
 * @param[out] outClosestColumns A list of \c NumClosestColumns pairs
 *     \f$ (i, d) \f$ sorted in ascending order of \f$ d \f$, where $i$ is a
 *     0-based column index in \f$ M \f$ and
 *     \f$ d \f$ is the distance (using \c inMetric) between \f$ M_i \f$ and
 *     \f$ x \f$.
 */
template <class DistanceFunction, class RandomAccessIterator>
void
closestColumnsAndDistances(
    const MappedMatrix& inMatrix,
    const MappedColumnVector& inVector,
    DistanceFunction& inMetric,
    RandomAccessIterator ioFirst,
    RandomAccessIterator ioLast)
{

    ReverseLexicographicComparator<
        typename std::iterator_traits<RandomAccessIterator>::value_type>
            comparator;

    std::fill(ioFirst, ioLast,
        std::make_tuple(0, std::numeric_limits<double>::infinity()));
    for (Index i = 0; i < inMatrix.cols(); ++i) {
        double currentDist;
        currentDist
            = AnyType_cast<double>(
                    inMetric(MappedColumnVector(inMatrix.col(i)), inVector)
                    );

        // outIndicesAndDistances is a heap, so the first element is maximal
        if (currentDist < std::get<1>(*ioFirst)) {
            // Unfortunately, the STL does not have a decrease-key function,
            // so we are wasting a bit of performance here
            std::pop_heap(ioFirst, ioLast, comparator);
            *(ioLast - 1) = std::make_tuple(i, currentDist);
            std::push_heap(ioFirst, ioLast, comparator);
        }
    }
    std::sort_heap(ioFirst, ioLast, comparator);
}

template <class RandomAccessIterator>
void
closestColumnsAndDistancesUDF(
    const MappedMatrix& inMatrix,
    const MappedColumnVector& inVector,
    RandomAccessIterator ioFirst,
    RandomAccessIterator ioLast,
    Oid oid)
{

    ReverseLexicographicComparator<
        typename std::iterator_traits<RandomAccessIterator>::value_type>
            comparator;

    std::fill(ioFirst, ioLast,
        std::make_tuple(0, std::numeric_limits<double>::infinity()));
    for (Index i = 0; i < inMatrix.cols(); ++i) {
        double currentDist;
        currentDist = static_cast<double>(DatumGetFloat8(OidFunctionCall2(
                        oid,
                        PointerGetDatum(VectorToNativeArray(inMatrix.col(i))),
                        PointerGetDatum(VectorToNativeArray(inVector))
                        )));

        // outIndicesAndDistances is a heap, so the first element is maximal
        if (currentDist < std::get<1>(*ioFirst)) {
            // Unfortunately, the STL does not have a decrease-key function,
            // so we are wasting a bit of performance here
            std::pop_heap(ioFirst, ioLast, comparator);
            *(ioLast - 1) = std::make_tuple(i, currentDist);
            std::push_heap(ioFirst, ioLast, comparator);
        }
    }
    std::sort_heap(ioFirst, ioLast, comparator);
}

double
distPNorm(const MappedColumnVector& inX, const MappedColumnVector& inY, double p) {
    if (inX.size() != inY.size()) {
        throw std::runtime_error("Found input arrays of "
                "different lengths unexpectedly.");
    }
    if (p <= 0 || std::isnan(p)) {
        throw std::runtime_error("Expect input p to be positive.");
    }

    if (!std::isfinite(p)) {
        return (inX - inY).lpNorm<Eigen::Infinity>();
    } else {
        double res = 0.0;
        for (int i = 0; i < inX.size(); i++) {
            res += std::pow(std::abs(inX(i) - inY(i)), p);
        }
        return std::pow(res, 1./p);
    }
}

double
distNorm1(const MappedColumnVector& inX, const MappedColumnVector& inY) {
    if (inX.size() != inY.size()) {
        throw std::runtime_error("Found input arrays of "
                "different lengths unexpectedly.");
    }

    return (inX - inY).lpNorm<1>();
}

double
distNorm2(const MappedColumnVector& inX, const MappedColumnVector& inY) {
    if (inX.size() != inY.size()) {
        throw std::runtime_error("Found input arrays of "
                "different lengths unexpectedly.");
    }

    return (inX - inY).norm();
}

double
cosineSimilarity(const MappedColumnVector& inX, const MappedColumnVector& inY) {
    if (inX.size() != inY.size()) {
        throw std::runtime_error("Found input arrays of "
                "different lengths unexpectedly.");
    }

	double xnorm = inX.norm(), ynorm = inY.norm();
	if (xnorm < std::numeric_limits<double>::denorm_min()
		|| ynorm < std::numeric_limits<double>::denorm_min()) {
        return -1;
    }

    return inX.dot(inY) / (xnorm * ynorm);
}

double
squaredDistNorm2(const MappedColumnVector& inX, const MappedColumnVector& inY) {
    if (inX.size() != inY.size()) {
        throw std::runtime_error("Found input arrays of "
                "different lengths unexpectedly.");
    }

    return (inX - inY).squaredNorm();
}

double
distAngle(const MappedColumnVector& inX, const MappedColumnVector& inY) {
    if (inX.size() != inY.size()) {
        throw std::runtime_error("Found input arrays of "
                "different lengths unexpectedly.");
    }

	// Deal with the undefined case where one of the norm is zero
	// Angle is not defined. Just return \pi.
	double xnorm = inX.norm(), ynorm = inY.norm();
	if (xnorm < std::numeric_limits<double>::denorm_min()
		|| ynorm < std::numeric_limits<double>::denorm_min())
		return std::acos(-1);

    double cosine = dot(inX, inY) / (xnorm * ynorm);
    if (cosine > 1)
        cosine = 1;
    else if (cosine < -1)
        cosine = -1;
    return std::acos(cosine);
}

double
distTanimoto(const MappedColumnVector& inX, const MappedColumnVector& inY) {
    if (inX.size() != inY.size()) {
        throw std::runtime_error("Found input arrays of "
                "different lengths unexpectedly.");
    }

    // Note that this is not a metric in general!
    double dotProduct = dot(inX, inY);
    double tanimoto = inX.squaredNorm() + inY.squaredNorm();
    return (tanimoto - 2 * dotProduct) / (tanimoto - dotProduct);
}

double
distJaccard(const ArrayHandle<text*>& inX, const ArrayHandle<text*>& inY) {
    if (inX.size() == 0 && inY.size() == 0)
        return 0.0;  // both empty are treated as zero distance
    if (inX.size() == 0 || inY.size() == 0)
        return 1.0;   // one of the sets being empty is treated as max distance

    std::set<string> x_set;
    for (size_t i = 0; i < inX.size(); i++){
        x_set.insert(std::string(VARDATA_ANY(inX[i]),
                                 VARSIZE_ANY(inX[i]) - VARHDRSZ));
    }

    size_t n_intersection = 0;
    size_t n_union = x_set.size();
    std::set<string> y_set;
    for (size_t i = 0; i < inY.size(); i++){
        string y_elem = std::string(VARDATA_ANY(inY[i]),
                                    VARSIZE_ANY(inY[i]) - VARHDRSZ);
        if (y_set.count(y_elem) == 0){
            // for sets, count returns 1 if an element with that value
            // exists in the container, and zero otherwise.

            if (x_set.count(y_elem) == 1){
                // element present in both sets (already counted for union)
                n_intersection++;
            }
            else{
                // element only in inY, not yet counted for union
                n_union++;
            }
        }
        y_set.insert(y_elem);
    }
    return 1.0 - static_cast<double>(n_intersection) / static_cast<double>(n_union);
}

/**
 * @brief Compute the k columns of a matrix that are closest to a vector
 *
 * For performance, we cheat here: For the following four distance functions, we
 * take a special shortcut.
 * FIXME: FunctionHandle should be tuned so that this shortcut no longer
 * impacts performance by more than, say, ~10%.
 */

std::string dist_fn_name(string s)
{
    std::istringstream ss(s);
    std::string token, fname;
    if (std::getline(ss, token, '.')) fname = token; // suppose there is no schema name
    if (std::getline(ss, token, '.')) fname = token; // previous part is schema name
    return fname;
}

template <class RandomAccessIterator>
inline
void
closestColumnsAndDistancesShortcut(
    const MappedMatrix& inMatrix,
    const MappedColumnVector& inVector,
    FunctionHandle &inDist,
    std::string fname,
    RandomAccessIterator ioFirst,
    RandomAccessIterator ioLast) {

    // Sorted in the order of expected use
    if (fname.compare("squared_dist_norm2") == 0)
        closestColumnsAndDistances(inMatrix, inVector, squaredDistNorm2,
            ioFirst, ioLast);
    else if (fname.compare("dist_norm2") == 0)
        closestColumnsAndDistances(inMatrix, inVector, distNorm2,
            ioFirst, ioLast);
    else if (fname.compare("dist_norm1") == 0)
        closestColumnsAndDistances(inMatrix, inVector, distNorm1,
            ioFirst, ioLast);
    else if (fname.compare("dist_angle") == 0)
        closestColumnsAndDistances(inMatrix, inVector, distAngle,
            ioFirst, ioLast);
    else if (fname.compare("dist_tanimoto") == 0) {
        closestColumnsAndDistances(inMatrix, inVector, distTanimoto,
            ioFirst, ioLast);
    } else {
        closestColumnsAndDistancesUDF(inMatrix, inVector, ioFirst,
                ioLast, inDist.funcID());
    }
}

/**
 * @brief Compute the minimum distance between a vector and any column of a
 *     matrix
 *
 * This function calls a user-supplied function, for which it does not do
 * garbage collection. It is therefore meant to be called only constantly many
 * times before control is returned to the backend.
 */
AnyType
closest_column::run(AnyType& args) {
    //if (true) throw std::runtime_error("Begin cc run\n");
    try{
        MappedMatrix M = args[0].getAs<MappedMatrix>();
        MappedColumnVector x = args[1].getAs<MappedColumnVector>();
        FunctionHandle dist = args[2].getAs<FunctionHandle>()
            .unsetFunctionCallOptions(FunctionHandle::GarbageCollectionAfterCall);
        string dist_fname = args[3].getAs<char *>();
        std::string fname = dist_fn_name(dist_fname);
        std::tuple<Index, double> result;
        closestColumnsAndDistancesShortcut(M, x, dist, fname, &result, &result + 1);

        AnyType tuple;
        return tuple
            << static_cast<int32_t>(std::get<0>(result))
            << std::get<1>(result);
    }catch (const ArrayWithNullException &e) {
        return Null();
    }
}

AnyType
closest_column_hawq::run(AnyType& args) {
    MappedMatrix M = args[0].getAs<MappedMatrix>();
    MappedColumnVector x = args[1].getAs<MappedColumnVector>();
    string distance_metric_str = args[2].getAs<char *>();
    boost::trim(distance_metric_str);

    double (*distance_metric)(const MappedColumnVector&, const MappedColumnVector&);

    // we hard-code comparision and selection of the distance function since
    // we are currently limited in not being able to access the catalog
    // in a function executed at the segments. This is a limitation in HAWQ
    // and will probably be eliminated in a future HAWQ release
    if ((distance_metric_str.compare("squared_dist_norm2") == 0) ||
            (distance_metric_str.compare("madlib.squared_dist_norm2") == 0)){
        distance_metric = squaredDistNorm2;
    } else if ((distance_metric_str.compare("dist_norm2") == 0) ||
            (distance_metric_str.compare("madlib.dist_norm2") == 0)){
        distance_metric = distNorm2;
    } else if ((distance_metric_str.compare("dist_norm1") == 0) ||
            (distance_metric_str.compare("madlib.dist_norm1") == 0)){
        distance_metric = distNorm1;
    } else if ((distance_metric_str.compare("dist_angle") == 0) ||
            (distance_metric_str.compare("madlib.dist_angle") == 0)){
        distance_metric = distAngle;
    } else if ((distance_metric_str.compare("dist_tanimoto") == 0) ||
            (distance_metric_str.compare("madlib.dist_tanimoto") == 0)){
        distance_metric = distTanimoto;
    } else{
        string errorMessage = string("Invalid distance metric provided: ") + \
                                distance_metric_str + \
                                string(". Currently only madlib provided distance functions are supported.");
        throw std::invalid_argument(errorMessage);
    }

    std::tuple<Index, double> result;
    closestColumnsAndDistances(M, x, distance_metric, &result, &result + 1);

    AnyType tuple;
    return tuple
        << static_cast<int32_t>(std::get<0>(result))
        << std::get<1>(result);
}

/**
 * @brief Compute the minimum distance between a vector and any column of a
 *     matrix
 *
 * This function calls a user-supplied function, for which it does not do
 * garbage collection. It is therefore meant to be called only constantly many
 * times before control is returned to the backend.
 */
AnyType
closest_columns::run(AnyType& args) {
    MappedMatrix M = args[0].getAs<MappedMatrix>();
    MappedColumnVector x = args[1].getAs<MappedColumnVector>();
    uint32_t num = args[2].getAs<uint32_t>();
    FunctionHandle dist = args[3].getAs<FunctionHandle>()
        .unsetFunctionCallOptions(FunctionHandle::GarbageCollectionAfterCall);
    string dist_fname = args[4].getAs<char *>();

    std::string fname = dist_fn_name(dist_fname);

    std::vector<std::tuple<Index, double> > result(num);
    closestColumnsAndDistancesShortcut(M, x, dist, fname, result.begin(),
        result.end());

    MutableArrayHandle<int32_t> indices = allocateArray<int32_t,
        dbal::FunctionContext, dbal::DoNotZero, dbal::ThrowBadAlloc>(num);
    MutableArrayHandle<double> distances = allocateArray<double,
        dbal::FunctionContext, dbal::DoNotZero, dbal::ThrowBadAlloc>(num);
    for (uint32_t i = 0; i < num; ++i)
        std::tie(indices[i], distances[i]) = result[i];

    AnyType tuple;
    return tuple << indices << distances;
}

AnyType
closest_columns_hawq::run(AnyType& args) {
    MappedMatrix M = args[0].getAs<MappedMatrix>();
    MappedColumnVector x = args[1].getAs<MappedColumnVector>();
    uint32_t num = args[2].getAs<uint32_t>();
    string distance_metric_str = args[3].getAs<char *>();
    boost::trim(distance_metric_str);

    if (0 == num) {
        throw std::invalid_argument("the parameter number should be a positive integer");
    }

    double (*distance_metric)(const MappedColumnVector&, const MappedColumnVector&);

    if ((distance_metric_str.compare("squared_dist_norm2") == 0) ||
            (distance_metric_str.compare("madlib.squared_dist_norm2") == 0)){
        distance_metric = squaredDistNorm2;
    } else if ((distance_metric_str.compare("dist_norm2") == 0) ||
            (distance_metric_str.compare("madlib.dist_norm2") == 0)){
        distance_metric = distNorm2;
    } else if ((distance_metric_str.compare("dist_norm1") == 0) ||
            (distance_metric_str.compare("madlib.dist_norm1") == 0)){
        distance_metric = distNorm1;
    } else if ((distance_metric_str.compare("dist_angle") == 0) ||
            (distance_metric_str.compare("madlib.dist_angle") == 0)){
        distance_metric = distAngle;
    } else if ((distance_metric_str.compare("dist_tanimoto") == 0) ||
            (distance_metric_str.compare("madlib.dist_tanimoto") == 0)){
        distance_metric = distTanimoto;
    } else{
        string errorMessage = string("Invalid distance metric provided: ") + \
                                distance_metric_str + \
                                string(". Currently only madlib provided distance functions are supported.");
        throw std::invalid_argument(errorMessage);
    }

    std::vector<std::tuple<Index, double> > result(num);
    closestColumnsAndDistances(M, x, distance_metric, result.begin(), result.end());

    MutableArrayHandle<int32_t> indices = allocateArray<int32_t,
        dbal::FunctionContext, dbal::DoNotZero, dbal::ThrowBadAlloc>(num);
    MutableArrayHandle<double> distances = allocateArray<double,
        dbal::FunctionContext, dbal::DoNotZero, dbal::ThrowBadAlloc>(num);
    for (uint32_t i = 0; i < num; ++i)
        std::tie(indices[i], distances[i]) = result[i];

    AnyType tuple;
    return tuple << indices << distances;
}

AnyType
norm1::run(AnyType& args) {
    return static_cast<double>(args[0].getAs<MappedColumnVector>().lpNorm<1>());
}

AnyType
norm2::run(AnyType& args) {
    return static_cast<double>(args[0].getAs<MappedColumnVector>().norm());
}

AnyType
dist_inf_norm::run(AnyType& args) {
    return distPNorm(
        args[0].getAs<MappedColumnVector>(),
        args[1].getAs<MappedColumnVector>(),
        std::numeric_limits<double>::infinity()
    );
}

AnyType
dist_pnorm::run(AnyType& args) {
    return distPNorm(
        args[0].getAs<MappedColumnVector>(),
        args[1].getAs<MappedColumnVector>(),
        args[2].getAs<double>()
    );
}

AnyType
dist_norm1::run(AnyType& args) {
    return distNorm1(
        args[0].getAs<MappedColumnVector>(),
        args[1].getAs<MappedColumnVector>()
    );
}

AnyType
dist_norm2::run(AnyType& args) {
    // FIXME: it would be nice to declare this as a template function (so it
    // works for dense and sparse vectors), and the C++ AL takes care of the
    // rest...
    return distNorm2(
        args[0].getAs<MappedColumnVector>(),
        args[1].getAs<MappedColumnVector>()
    );
}

AnyType
cosine_similarity::run(AnyType& args) {
    return cosineSimilarity(
        args[0].getAs<MappedColumnVector>(),
        args[1].getAs<MappedColumnVector>()
    );
}

AnyType
squared_dist_norm2::run(AnyType& args) {
    return squaredDistNorm2(
        args[0].getAs<MappedColumnVector>(),
        args[1].getAs<MappedColumnVector>()
    );
}

AnyType
dist_angle::run(AnyType& args) {
    return distAngle(
        args[0].getAs<MappedColumnVector>(),
        args[1].getAs<MappedColumnVector>()
    );
}

AnyType
dist_tanimoto::run(AnyType& args) {
    return distTanimoto(
        args[0].getAs<MappedColumnVector>(),
        args[1].getAs<MappedColumnVector>()
    );
}

AnyType
dist_jaccard::run(AnyType& args) {
    return distJaccard(
        args[0].getAs<ArrayHandle<text*> >(),
        args[1].getAs<ArrayHandle<text*> >()
    );
}

} // namespace linalg

} // namespace modules

} // namespace regress


#include "panoramic_stitching.h"

#include <unordered_map>
#include <algorithm>
#include <iostream>
#include <stdlib.h>
#include <cmath>
#include <unordered_set>
#include <unsupported/Eigen/NonLinearOptimization>
#include <Eigen/src/Core/util/DisableStupidWarnings.h>

namespace panSti {
panoramicStitching::panoramicStitching(int base, const std::vector<std::string>& name,
                                       std::vector<std::pair<int, int>>& image_pairs, featureMethod method): m_image_pairs(image_pairs),
    m_base_index(base)	{
    for(int i = 0; i < name.size(); i++) {
        m_images.push_back(cv::imread(name[i]));
        if (!m_images.back().data) {
            std::cout << "cannot read image" << std::endl;
            return;
        }
        cv::Mat gray_image;
        cv::cvtColor(m_images.back(), gray_image, CV_RGB2GRAY);
        m_gray_images.push_back(gray_image);
    }
    if (method == featureMethod::SIFT) {
        bool result = siftMatchExtract();
        if (!result) {
            std::cout << "sift extract fails" << std::endl;
			return;
        }
    } else {
		std::cout << "At present only SIFT provided" << std::endl;
	}
}
cv::Mat panoramicStitching::stitch_image() {
    calculateHomography();
    //-- step 1: calculate all the homography with the base
    std::unordered_map<int, int> hash_map; // record the
    int pair_index = 0;
    std::vector<std::pair<std::pair<int, int>, int>> unsolved_pairs;
    for (std::pair<int, int>& image_pair : m_image_pairs) {
        if (image_pair.first == m_base_index) {
            if (hash_map.find(image_pair.second) != hash_map.end()) {
                pair_index++;
                continue;
            }
            hash_map.insert({image_pair.second, pair_index});
            pair_index++;
        } else if (image_pair.second == m_base_index) {
            if (hash_map.find(image_pair.first) != hash_map.end()) {
                pair_index++;
                continue;
            }
            hash_map.insert({image_pair.first, pair_index});
            // invert the homography
            m_homographies[pair_index] = m_homographies[pair_index].inv();
            image_pair.second = image_pair.first;
            image_pair.first = m_base_index;
            pair_index++;
        } else {
            if (hash_map.find(image_pair.first) != hash_map.end()) {
                if (hash_map.find(image_pair.second) != hash_map.end()) {
                    pair_index++;
                    continue;
                }
                m_homographies[pair_index] = m_homographies[hash_map.at(image_pair.first)] *
					m_homographies[pair_index];
                //image_pair.second = image_pair.first;
                image_pair.first = m_base_index;
                hash_map.insert({image_pair.second, pair_index});
                pair_index++;
            } else if (hash_map.find(image_pair.second) != hash_map.end()) {
                if (hash_map.find(image_pair.first) != hash_map.end()) {
                    pair_index++;
                    continue;
                }
                m_homographies[pair_index] = m_homographies[hash_map.at(image_pair.second)] *
					m_homographies[pair_index].inv();
				image_pair.second = image_pair.first;
                image_pair.first = m_base_index;
                hash_map.insert({image_pair.second, pair_index});
                pair_index++;
            } else {
                // two images are all not in the
                unsolved_pairs.push_back({image_pair, pair_index});
                pair_index++;
            }
        }
    }
    int max_iter = 3;
    while (!unsolved_pairs.empty() && max_iter >= 0) {
        for (auto i = unsolved_pairs.begin(); i != unsolved_pairs.end();) {
            if (hash_map.find(i->first.first) != hash_map.end() ||
                    hash_map.find(i->first.second) != hash_map.end()) {
                if (hash_map.find(i->first.first) != hash_map.end() &&
                        hash_map.find(i->first.second) == hash_map.end()) {
                    m_homographies[i->second] = m_homographies[i->second] *
                                                m_homographies[hash_map.at(i->first.first)];
                    hash_map.insert({i->first.second, i->second});
                } else if (hash_map.find(i->first.first) == hash_map.end() &&
                           hash_map.find(i->first.second) != hash_map.end()) {
                    m_homographies[i->second] = m_homographies[i->second].inv() *
                                                m_homographies[hash_map.at(i->first.second)];
                    hash_map.insert({i->first.first, i->second});
                }
                unsolved_pairs.erase(i++); // erase this element and let i plus one
            } else {
                i++;
            }
        }
        max_iter--;
    }
    if (max_iter < 0) {
        std::cout << "there is no enough link between images" << std::endl;
        cv::Mat tmp;
        return tmp;
    }
    //-- step 2: calculate four edge points of the source image to find the size of the image
    // first of all, calculate the size of the big image
    int left_most = 0;
    int right_most = m_gray_images[m_base_index].cols - 1;
    int down_most = m_gray_images[m_base_index].rows - 1;
    int up_most = 0;
    // the interval of every_image before adding the left_most and up_most
    std::unordered_map<int, std::pair<int, int>> x_interval;
    std::unordered_map<int, std::pair<int, int>> y_interval;
    x_interval.insert({m_base_index, {left_most, right_most}});
    y_interval.insert({m_base_index, {up_most, down_most}});
    for (auto i = hash_map.begin(); i != hash_map.end(); i++) {
        int image_index = i->first;
        int homo_index = i->second;
        m_homographies[homo_index] = m_homographies[homo_index] /
                                     m_homographies[homo_index].at<double>(2, 2);
        int x[4] = {0, 0,
                    m_gray_images[image_index].cols - 1, m_gray_images[image_index]. cols - 1
                   };
        int y[4] = {0, m_gray_images[image_index].rows - 1,
                    0, m_gray_images[image_index]. rows - 1
                   };
        int this_left = INT_MAX;
        int this_right = INT_MIN;
        int this_up = INT_MAX;
        int this_down = INT_MIN;
        for (int j = 0; j < 4; j++) {
            cv::Mat coor = m_homographies[homo_index].col(0) * x[j] +
                           m_homographies[homo_index].col(1) * y[j] +
                           m_homographies[homo_index].col(2);
            double x = coor.at<double>(0, 0) / coor.at<double>(2, 0);
            double y = coor.at<double>(1, 0) / coor.at<double>(2, 0);
            left_most = std::min(left_most, int(std::floor(x)));
            right_most = std::max(right_most, int(std::ceil(x)));
            up_most = std::min(up_most, int(std::floor(y)));
            down_most = std::max(down_most, int(std::ceil(y)));
            this_left = std::min(this_left, int(std::floor(x)));
            this_right = std::max(this_right, int(std::ceil(x)));
            this_up = std::min(this_up, int(std::floor(y)));
            this_down = std::max(this_down, int(std::ceil(y)));
        }
        x_interval.insert({image_index, {this_left, this_right}});
        y_interval.insert({image_index, {this_up, this_down}});
    }
    // all images should be translated with (fabs(left_most), fabs(up_most))
    // and then get the warped_images
    cv::Mat base_homo = (cv::Mat_<double>(3, 3) <<
                         1, 0, std::abs(left_most),
                         0, 1, std::abs(up_most),
                         0, 0, 1);
    std::vector<cv::Mat> warped_images;
    int new_width = right_most - left_most + 1;
    int new_height = down_most - up_most + 1;
    for (int i = 0; i < m_gray_images.size(); i++) {
        if (i != m_base_index) {
            int homo_index = hash_map.at(i);
            m_homographies[homo_index] = base_homo * m_homographies[homo_index];
            cv::Mat warped_image;
            cv::warpPerspective(m_images[i], warped_image,
                                m_homographies[homo_index], cv::Size2i(new_width, new_height));
            warped_images.push_back(warped_image);
        } else {
            cv::Mat warped_image;
            cv::warpPerspective(m_images[m_base_index], warped_image,
                                base_homo, cv::Size2i(new_width, new_height));
            warped_images.push_back(warped_image);
        }
    }
    // secondly, combine the corresponding warpped image
    // for every pixel, set the pixel-value of image whose center is nearest the the pixel
    // calculate every center coordinate first of all
    std::unordered_map<int, std::pair<int, int>> center_coor; // the coordinate is (y, x)
    center_coor.insert({ m_base_index, {m_gray_images[m_base_index].rows / 2 + std::abs(up_most),
                                        m_gray_images[m_base_index].cols / 2 + std::abs(left_most)
                                       }
                       });
    for (int i = 0; i < m_images.size(); i++) {
        if (i == m_base_index)
            continue;
        cv::Mat tmp_H = m_homographies[hash_map.at(i)];
        int mid_col = m_gray_images[i].cols >> 1;
        int mid_row = m_gray_images[i].rows >> 1;
        int mid_x = tmp_H.at<double>(0, 0) * mid_col + tmp_H.at<double>(0, 1) * mid_row + tmp_H.at<double>(0, 2);
        int mid_y = tmp_H.at<double>(1, 0) * mid_col + tmp_H.at<double>(1, 1) * mid_row + tmp_H.at<double>(1, 2);
        double ww = tmp_H.at<double>(2, 0) * mid_col + tmp_H.at<double>(2, 1) * mid_row + tmp_H.at<double>(2, 2);
        mid_x /= ww;
        mid_y /= ww;
        center_coor.insert({ i, {mid_y, mid_x} });
    }
    // set pixel value for new image
    cv::Mat result = warped_images[m_base_index];
    for (int i = 0; i < new_height; i++) {
        for (int j = 0; j < new_width; j++) {
            int nearest_value[3] = { -1, -1, -1 };
            int nearest_dst = INT_MAX;
            for (int k = 0; k < m_images.size(); k++) {
                if (j < x_interval.at(k).second + std::abs(left_most) && j > x_interval.at(k).first + std::abs(left_most) &&
                        i < y_interval.at(k).second + std::abs(up_most)&& i > y_interval.at(k).first + std::abs(up_most)) {
                    if (warped_images[k].at<cv::Vec3b>(i, j)[0] == 0 &&
                            warped_images[k].at<cv::Vec3b>(i, j)[1] == 0 &&
                            warped_images[k].at<cv::Vec3b>(i, j)[2] == 0)
                        continue;
                    else {
                        int dist = std::fabs(i - center_coor.at(k).first) + std::fabs(j - center_coor.at(k).second);
                        if (dist < nearest_dst) {
                            nearest_dst = dist;
                            nearest_value[0] = warped_images[k].at<cv::Vec3b>(i, j)[0];
                            nearest_value[1] = warped_images[k].at<cv::Vec3b>(i, j)[1];
                            nearest_value[2] = warped_images[k].at<cv::Vec3b>(i, j)[2];
                        }
                    }
                }
            }
            if (nearest_value[0] == -1)
                continue;
            else {
                result.at<cv::Vec3b>(i, j)[0] = nearest_value[0];
                result.at<cv::Vec3b>(i, j)[1] = nearest_value[1];
                result.at<cv::Vec3b>(i, j)[2] = nearest_value[2];
            }
        }
    }
    return result;
}

bool panoramicStitching::siftMatchExtract() {
    cv::SiftFeatureDetector detector;
    cv::SiftDescriptorExtractor extractor;
    std::vector<cv::Mat> descriptor_vec;
    for (const cv::Mat& image : m_gray_images) {
        if (!image.data) {
            std::cout << "image data wrong" << std::endl;
            return false;
        }
        //-- Step 1: detect the feature points with the Sift Detector
        std::vector<cv::KeyPoint> keypoints;
        detector.detect(image, keypoints);
        m_keyPoints.push_back(keypoints);
        //-- Step 2: calculate the descriptor
        cv::Mat descriptor;
        extractor.compute(image, keypoints, descriptor);
        descriptor_vec.push_back(descriptor);
    }
    double good_match_threshold = 0.02;
    for (const std::pair<int, int>& pair_index : m_image_pairs) {
        //-- Step 3: Match the descriptor vectors using FLANN matcher
        cv::Mat l_descriptor = descriptor_vec[pair_index.first];
        cv::Mat r_descriptor = descriptor_vec[pair_index.second];
        cv::FlannBasedMatcher matcher;
        std::vector<cv::DMatch> matches;
        matcher.match(l_descriptor, r_descriptor, matches);
        //-- Calculate good matches whose distance is less than 2*min_dist
        //-- or a small arbitary value (0.02) when min_dist is very small
        double min_dist = 100, max_dist = 0;
        for (const cv::DMatch& match : matches) {
            min_dist = std::min(min_dist, double(match.distance));
            max_dist = std::max(max_dist, double(match.distance));
        }
        std::vector<cv::DMatch> good_matches;
        for (const cv::DMatch& match : matches) {
            if (match.distance <= std::max(2 * min_dist, good_match_threshold))
                good_matches.push_back(match);
        }
        m_matches.push_back(good_matches);
    }
    return true;
}

bool panoramicStitching::calculateHomography() {
    const double confidence = 0.95;
    const int maxIter = 2000;
    const double ransac_threshold = 3;
    bool result = false;
    for (int index = 0; index < m_matches.size(); index++) {
        int left_index = m_image_pairs[index].first;
        int right_index = m_image_pairs[index].second;
        int inlier_number = 0;
        cv::Mat inlier_mask = runRansac(m_keyPoints[left_index], m_keyPoints[right_index], m_matches[index],
                                        ransac_threshold, confidence, maxIter, result, inlier_number);
        if (!result) {
            std::cout << "ransac fails" << std::endl;
            return false;
        } else {
            std::cout << "ransac succeeds" << std::endl;
        }
        cv::Mat left_points(2, inlier_number, CV_64F);
        cv::Mat right_points(2, inlier_number, CV_64F);
        double *p2lx = left_points.ptr<double>(0);
        double *p2ly = left_points.ptr<double>(1);
        double *p2rx = right_points.ptr<double>(0);
        double *p2ry = right_points.ptr<double>(1);
        int *p2mask = inlier_mask.ptr<int>(0);
        int tmp_index = 0; // the tmp index of the point matrix
        for (int i = 0; i < inlier_mask.cols; i++) {
            if (*(p2mask + i) == 1) {
                *(p2lx + tmp_index) = m_keyPoints[left_index][m_matches[index][i].queryIdx].pt.x;
                *(p2ly + tmp_index) = m_keyPoints[left_index][m_matches[index][i].queryIdx].pt.y;
                *(p2rx + tmp_index) = m_keyPoints[right_index][m_matches[index][i].trainIdx].pt.x;
                *(p2ry + tmp_index) = m_keyPoints[right_index][m_matches[index][i].trainIdx].pt.y;
                tmp_index++;
            }
        }
        // use the inlier mask to get the initial homography with DLT
        cv::Mat homo = runKernel(left_points, right_points);
        // use the iteration to refine the homography
        homo = refine(left_points, right_points, homo);
        m_homographies.push_back(homo);
    }
    return true;
}

cv::Mat panoramicStitching::runRansac(const std::vector<cv::KeyPoint>& l_point,
                                      const std::vector<cv::KeyPoint>& r_point,
                                      const std::vector<cv::DMatch>& match,
                                      double reproj_threshold, double confidence,
                                      int maxIter, bool & succeed, int &max_inlier_number) {
    const int least_point = 4; // the least number points needed to calculate the homography
    const int n_points = match.size();
    cv::Mat l_pointsM(3, n_points, CV_64F);
    cv::Mat r_pointsM(3, n_points, CV_64F);
    // get the coordinate matrix for all points
    for (int i = 0 ; i < n_points; i++) {
        l_pointsM.at<double>(0, i) = l_point[match[i].queryIdx].pt.x;
        l_pointsM.at<double>(1, i) = l_point[match[i].queryIdx].pt.y;
        l_pointsM.at<double>(2, i) = 1;
        r_pointsM.at<double>(0, i) = r_point[match[i].trainIdx].pt.x;
        r_pointsM.at<double>(1, i) = r_point[match[i].trainIdx].pt.y;
        r_pointsM.at<double>(2, i) = 1;
    }
    int iter_times = 0;
    max_inlier_number = 0;
    cv::Mat most_inlier_mask;
    while (iter_times <= maxIter) {
        std::unordered_set<int> hash_set; // use hash_set to ensure no duplicate points in the example points
        cv::Mat l_pointSample(3, least_point, CV_64F);
        cv::Mat r_pointSample(3, least_point, CV_64F);
        int exist_points = 0;
        //-- step 1:get 4 points randomly
        while (exist_points < least_point) {
            int next_index = rand() % n_points;
            if (hash_set.find(next_index) != hash_set.end())
                continue;
            hash_set.insert(next_index);
            l_pointsM.col(next_index).copyTo(l_pointSample.col(exist_points));
            r_pointsM.col(next_index).copyTo(r_pointSample.col(exist_points));
            exist_points++;
            if (exist_points >= 3 &&
                    (isColinear(l_pointSample, exist_points) ||
                     isColinear(r_pointSample, exist_points)))
                exist_points--;
        }
        //-- step 2: use 4 points to get one homography
        cv::Mat model = runKernel(l_pointSample, r_pointSample);
        //-- step 3: find the inlines
        int tmp_inlier_number = 0;
        cv::Mat tmp_mask = findInliers(l_pointsM, r_pointsM, model,
                                       reproj_threshold, tmp_inlier_number);
        //-- change the iteration times according to the confidence and the inlier number
        if (tmp_inlier_number > max_inlier_number) {
            most_inlier_mask = tmp_mask;
            max_inlier_number = tmp_inlier_number;
            maxIter = std::min(maxIter,
                               updateMaxIter(max_inlier_number, n_points, confidence));
        }
        iter_times++;
    }
    if (max_inlier_number > least_point) {
        succeed = true;
        return most_inlier_mask;
    } else {
        succeed = false;
        return most_inlier_mask;
    }
}

// In order to increase the speed, we use the knowledge that everytime we will
// only add one point, and the previous point is guaranted to be not colinear
bool panoramicStitching::isColinear(cv::Mat points, int number) {
    if (number < 3) return false;
    // choose every possible 2 points
    int k = number - 1;
    for (int i = 0; i < number - 2; i++) {
        for (int j = i + 1; j < number - 1; j++) {
            double x1 = points.at<double>(0, i);
            double y1 = points.at<double>(1, i);
            double x2 = points.at<double>(0, j);
            double y2 = points.at<double>(1, j);
            double x3 = points.at<double>(0, k);
            double y3 = points.at<double>(1, k);
            double dx1 = x1 - x2;
            double dy1 = y1 - y2;
            double dx2 = x1 - x3;
            double dy2 = y1 - y3;
            // DBL_EPSILON: is used to compare double precision with 0
            // the equation is derived from the distance from one point to one line
            if (fabs(dx1 * dy2 - dx2 * dy1) <
                    DBL_EPSILON * (fabs(dx1) + fabs(dy1) + fabs(dx2) + fabs(dy2)))
                return true;
        }
    }
    return false;
}

cv::Mat panoramicStitching::runKernel(cv::Mat l_points, cv::Mat r_points) {
    //-- Step 1: Normalize, center is 0, and average distance to center is sqrt(2)
    // the detailed algorithm explain can be seen in Multiview Geometry in Computer Vision P109
    int data_number = l_points.cols;
    double *p2lx = l_points.ptr<double>(0);
    double *p2ly = l_points.ptr<double>(1);
    double *p2rx = r_points.ptr<double>(0);
    double *p2ry = r_points.ptr<double>(1);
    double l_center[2] = {0, 0};
    double r_center[2] = {0, 0};
    for (int i = 0; i < data_number; i++) {
        l_center[0] += *(p2lx + i);
        l_center[1] += *(p2ly + i);
        r_center[0] += *(p2rx + i);
        r_center[1] += *(p2ry + i);
    }
    l_center[0] /= data_number;
    l_center[1] /= data_number;
    r_center[0] /= data_number;
    r_center[1] /= data_number;
    double sum_l[2] = {0, 0};
    double sum_r[2] = {0, 0};
    // use the sum of absolute value to approximate the distance between points and center
    for (int i = 0; i < data_number; i++) {
        sum_l[0] += fabs(*(p2lx + i) - l_center[0]);
        sum_l[1] += fabs(*(p2ly + i) - l_center[1]);
        sum_r[0] += fabs(*(p2rx + i) - r_center[0]);
        sum_r[1] += fabs(*(p2ry + i) - r_center[1]);
    }
    if (fabs(sum_l[0]) < DBL_EPSILON || fabs(sum_l[1]) < DBL_EPSILON ||
            fabs(sum_r[0]) < DBL_EPSILON || fabs(sum_r[1]) < DBL_EPSILON) {
        std::cout << "the inlier points have something wrong" << std::endl;
        return cv::Mat(0, 0, CV_64F);
    }
    sum_l[0] = data_number / sum_l[0];
    sum_l[1] = data_number / sum_l[1];
    sum_r[0] = data_number / sum_r[0];
    sum_r[1] = data_number / sum_r[1];
    // final homography is T'^(-1) * H_normalize * T
    // Be careful that the Homography is used to transfer the right to the left
    cv::Mat T1_inv = (cv::Mat_<double>(3, 3) <<
                      1 / sum_l[0], 0, l_center[0],
                      0, 1 / sum_l[1], l_center[1],
                      0, 0, 1);
    cv::Mat T2 = (cv::Mat_<double>(3, 3) <<
                  sum_r[0], 0, -sum_r[0] * r_center[0],
                  0, sum_r[1], -sum_r[1] * r_center[1],
                  0, 0, 1);
    //--Step 2: use the normalized coordinate to calculate DLT
    cv::Mat L(cv::Mat::zeros(2 * data_number, 9, CV_64F));
    for (int i = 0; i < data_number; i++) {
        double lx = (*(p2lx + i) - l_center[0]) * sum_l[0];
        double ly = (*(p2ly + i) - l_center[1]) * sum_l[1];
        double rx = (*(p2rx + i) - r_center[0]) * sum_r[0];
        double ry = (*(p2ry + i) - r_center[1]) * sum_r[1];
        double *p2up = L.ptr<double>(i * 2);
        double *p2down = L.ptr<double>(i * 2 + 1);
        *(p2up + 3) = -rx;
        *(p2up + 4) = -ry;
        *(p2up + 5) = -1;
        *(p2up + 6) = ly * rx;
        *(p2up + 7) = ly * rx;
        *(p2up + 8) = ly;
        *p2down = rx;
        *(p2down + 1) = ry;
        *(p2down + 2) = 1;
        *(p2down + 6) = -lx * rx;
        *(p2down + 7) = -lx * ry;
        *(p2down + 8) = -lx;
    }
    cv::Mat LtL = L.t() * L;
    cv::Mat eigenValue, eigenVector;
    cv::eigen(LtL, eigenValue, eigenVector);
    cv::Mat result = eigenVector.row(8);
    // we need to clone the result here because the reshape needs that matrix data is continuous
    result = T1_inv * result.clone().reshape(0, 3) * T2;
    result = result / result.at<double>(2, 2);
    // return the 3*3 H matrix
    return result;
}

cv::Mat panoramicStitching::findInliers(cv::Mat l_points, cv::Mat r_points,
                                        cv::Mat homography,
                                        const double threshold, int& inlier_number) {
    int n = l_points.cols; // number of points
    cv::Mat result_mask(1, n, CV_32S);
    // calculate the geometry error
    double h[9];
    for (int i = 0; i < 9; i++) {
        h[i] = homography.at<double>(i / 3, i - 3 * (i / 3));
    }
    double* p2rx = r_points.ptr<double>(0);
    double* p2ry = r_points.ptr<double>(1);
    double* p2rw = r_points.ptr<double>(2);
    double* p2lx = l_points.ptr<double>(0);
    double* p2ly = l_points.ptr<double>(1);
    double* p2lw = l_points.ptr<double>(2);
    double threshold_square = threshold * threshold;
    int *p2result = result_mask.ptr<int>(0);
    int count = 0;
    for (int i = 0; i < n; i++) {
        double ww = 1.0 / (h[6] * (*(p2rx + i)) + h[7] * (*(p2ry + i)) + h[8] * (*(p2rw + i)));
        double xx = ww * (h[0] * (*(p2rx + i)) + h[1] * (*(p2ry + i)) + h[2] * (*(p2rw + i)))
                    - (*(p2lx + i)) / (*(p2lw + i));
        double yy = ww * (h[3] * (*(p2rx + i)) + h[4] * (*(p2ry + i)) + h[5] * (*(p2rw + i)))
                    - (*(p2ly + i)) / (*(p2lw + i));
        if ((xx * xx + yy * yy) <= threshold_square) {
            *(p2result + i) = 1;
            count++;
        } else {
            *(p2result + i) = 0;
        }
    }
    inlier_number = count;
    return result_mask;
}

int panoramicStitching::updateMaxIter(int inlier_number, int total_number, double confidence) {
    int model_points = 4; // we need 4 points to calculate one homography
    double e_p = 1 - double(inlier_number) / total_number;
    if (e_p > 1) {
        std::cout << "the error rate is greater than 1" << std::endl;
        return -1;
    }
    double num = std::max(1.0 - confidence, DBL_MIN); // avoid inf
    double denom = 1.0 - pow(1.0 - e_p, model_points);
    if (denom < DBL_MIN) return 0; // all points are inliers
    num = std::log(num) / std::log(denom);
	// prevent the overflow
	if (num > 10000)
		return 10000;
    return std::round(num);
}

cv::Mat panoramicStitching::refine(cv::Mat left_points, cv::Mat right_points, cv::Mat ini_homo) {
    Eigen::VectorXd x(8);
    double h8 = ini_homo.at<double>(2, 2);
    for (int i = 0; i <= 7; i++) {
        x[i] = ini_homo.at<double>(i / 3, i - 3 * (i / 3)) / h8;
    }
    //Eigen::VectorXd x_previous = x.replicate(1, 1);
    geometricError functor(left_points, right_points);
    Eigen::NumericalDiff<geometricError> num_diff(functor);
    Eigen::LevenbergMarquardt<Eigen::NumericalDiff<geometricError>> lm(num_diff);
    int ret = lm.minimize(x);
    //std::cout << x - x_previous << std::endl;
    for (int i = 0; i <= 7; i++) {
        ini_homo.at<double>(i / 3, i - 3 * (i / 3)) = x[i];
    }
    ini_homo.at<double>(2, 2) = 1;
    return ini_homo;
}

geometricError::geometricError(cv::Mat left_points, cv::Mat right_points):
    Functor<double>(8, 2 * left_points.cols),
    left_points(left_points),
    right_points(right_points) {
}

int geometricError::operator()(const Eigen::VectorXd & x, Eigen::VectorXd & fvec) const {
    double* p2lx = const_cast<double*>(left_points.ptr<double>(0));
    double* p2ly = const_cast<double*>(left_points.ptr<double>(1));
    double* p2rx = const_cast<double*>(right_points.ptr<double>(0));
    double* p2ry = const_cast<double*>(right_points.ptr<double>(1));
    for (int i = 0; i < values() / 2; i++) {
        double lx = *(p2lx + i);
        double ly = *(p2ly + i);
        double rx = *(p2rx + i);
        double ry = *(p2ry + i);
        double ww = x[6] * rx + x[7] * ry + 1;
        ww = std::fabs(ww) > DBL_EPSILON ? 1. / ww : 0;
        double lx_est = ww * (x[0] * rx + x[1] * ry + x[2]);
        double ly_est = ww * (x[3] * rx + x[4] * ry + x[5]);
        fvec[i * 2] = lx - lx_est;
        fvec[i * 2 + 1] = ly - ly_est;
    }
    return 0;
}

int geometricError::df(const Eigen::VectorXd & x, Eigen::MatrixXd & fjac) const {
    double* p2lx = const_cast<double*>(left_points.ptr<double>(0));
    double* p2ly = const_cast<double*>(left_points.ptr<double>(1));
    double* p2rx = const_cast<double*>(right_points.ptr<double>(0));
    double* p2ry = const_cast<double*>(right_points.ptr<double>(1));
    for (int i = 0; i < values() >> 1; i++) {
        double rx = *(p2rx + i);
        double ry = *(p2ry + i);
        double ww = x[6] * rx + x[7] * ry + 1;
        ww = std::fabs(ww) > DBL_EPSILON ? 1. / ww : 0;
        double lx_est = ww * (x[0] * rx + x[1] * ry + x[2]);
        double ly_est = ww * (x[3] * rx + x[4] * ry + x[3]);
        fjac(i * 2, 0) = rx * ww;
        fjac(i * 2, 1) = ry * ww;
        fjac(i * 2, 2) = ww;
        fjac(i * 2, 3) = 0;
        fjac(i * 2, 4) = 0;
        fjac(i * 2, 5) = 0;
        fjac(i * 2, 6) = -rx * lx_est * ww;
        fjac(i * 2, 7) = -ry * lx_est * ww;
        fjac(i * 2 + 1, 0) = 0;
        fjac(i * 2 + 1, 1) = 0;
        fjac(i * 2 + 1, 2) = 0;
        fjac(i * 2 + 1, 3) = rx * ww;
        fjac(i * 2 + 1, 4) = ry * ww;
        fjac(i * 2 + 1, 5) = ww;
        fjac(i * 2 + 1, 6) = -rx * ly_est * ww;
        fjac(i * 2 + 1, 7) = -ry * ly_est * ww;
    }
    return 0;
}

}


/*
   Copyright 2020 The Silkworm Authors

   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.
*/

#include "snark.hpp"

#include <algorithm>
#include <boost/endian/conversion.hpp>
#include <cassert>
#include <cstring>
#include <libff/algebra/curves/alt_bn128/alt_bn128_pp.hpp>
#include <libff/common/profiling.hpp>

namespace silkworm::snark {

void init_libff() noexcept {
    // magic static
    [[maybe_unused]] static bool initialized = []() noexcept {
        libff::inhibit_profiling_info = true;
        libff::inhibit_profiling_counters = true;
        libff::alt_bn128_pp::init_public_params();
        return true;
    }();
}

Scalar to_scalar(ByteView be) noexcept {
    mpz_t m;
    mpz_init(m);
    mpz_import(m, be.size(), /*order=*/1, /*size=*/1, /*endian=*/0, /*nails=*/0, &be[0]);
    Scalar out{m};
    mpz_clear(m);
    return out;
}

// Notation warning: Yellow Paper's p is the same libff's q.
// Returns x < p (YP notation).
static bool valid_element_of_fp(const Scalar& x) noexcept {
    return mpn_cmp(x.data, libff::alt_bn128_modulus_q.data, libff::alt_bn128_q_limbs) < 0;
}

std::optional<libff::alt_bn128_G1> decode_g1_element(ByteView bytes64_be) noexcept {
    assert(bytes64_be.size() == 64);

    Scalar x{to_scalar(bytes64_be.substr(0, 32))};
    if (!valid_element_of_fp(x)) {
        return {};
    }

    Scalar y{to_scalar(bytes64_be.substr(32, 32))};
    if (!valid_element_of_fp(y)) {
        return {};
    }

    if (x.is_zero() && y.is_zero()) {
        return libff::alt_bn128_G1::zero();
    }

    libff::alt_bn128_G1 point{x, y, libff::alt_bn128_Fq::one()};
    if (!point.is_well_formed()) {
        return {};
    }
    return point;
}

static std::optional<libff::alt_bn128_Fq2> decode_fp2_element(ByteView bytes64_be) noexcept {
    assert(bytes64_be.size() == 64);

    // big-endian encoding
    Scalar c0{to_scalar(bytes64_be.substr(32, 32))};
    Scalar c1{to_scalar(bytes64_be.substr(0, 32))};

    if (!valid_element_of_fp(c0) || !valid_element_of_fp(c1)) {
        return {};
    }

    return libff::alt_bn128_Fq2{c0, c1};
}

std::optional<libff::alt_bn128_G2> decode_g2_element(ByteView bytes128_be) noexcept {
    assert(bytes128_be.size() == 128);

    std::optional<libff::alt_bn128_Fq2> x{decode_fp2_element(bytes128_be.substr(0, 64))};
    if (!x) {
        return {};
    }

    std::optional<libff::alt_bn128_Fq2> y{decode_fp2_element(bytes128_be.substr(64, 64))};
    if (!y) {
        return {};
    }

    if (x->is_zero() && y->is_zero()) {
        return libff::alt_bn128_G2::zero();
    }

    libff::alt_bn128_G2 point{*x, *y, libff::alt_bn128_Fq2::one()};
    if (!point.is_well_formed()) {
        return {};
    }

    if (!(libff::alt_bn128_G2::order() * point).is_zero()) {
        // wrong order, doesn't belong to the subgroup G2
        return {};
    }

    return point;
}

Bytes encode_g1_element(libff::alt_bn128_G1 p) noexcept {
    Bytes out(64, '\0');
    if (p.is_zero()) {
        return out;
    }

    p.to_affine_coordinates();

    auto x{p.X.as_bigint()};
    auto y{p.Y.as_bigint()};

    // Here we convert little-endian data to big-endian output
    static_assert(boost::endian::order::native == boost::endian::order::little);
    static_assert(sizeof(x.data) == 32);

    std::memcpy(&out[0], y.data, 32);
    std::memcpy(&out[32], x.data, 32);

    std::reverse(out.begin(), out.end());
    return out;
}
}  // namespace silkworm::snark


/*
 * MIT License
 *
 * Copyright (c) 2020 Robert Grupp
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

#include "xregImgSimMetric2DPatchNCCOCL.h"

// ITK pollutes the global namespace with a macro, and causes
// a compile failure with vienna cl
#ifdef vcl_size_t
#undef vcl_size_t
#endif
#ifdef vcl_ptrdiff_t
#undef vcl_ptrdiff_t
#endif

#include <boost/compute/utility/source.hpp>

#include <viennacl/matrix.hpp>
#include <viennacl/vector.hpp>
#include <viennacl/linalg/prod.hpp>

#include "xregAssert.h"
#include "xregITKOpenCVUtils.h"

namespace
{

const char* kPATCH_NCC_OPENCL_SRC = BOOST_COMPUTE_STRINGIZE_SOURCE(

float2 xregComputePatchStats(__global const float* patch,
                             const ulong patch_stride,
                             const uint patch_num_cols,
                             const uint patch_num_rows)
{
  __global const float* patch_cur_row = patch;

  const uint num_pix_in_patch = patch_num_cols * patch_num_rows;

  float2 mean_and_std_dev = (float2) (0, 0);

  for (uint r = 0; r < patch_num_rows; ++r, patch_cur_row += patch_stride)
  {
    for (uint c = 0; c < patch_num_cols; ++c)
    {
      mean_and_std_dev.x += patch_cur_row[c];
    }
  }

  mean_and_std_dev.x /= num_pix_in_patch;

  patch_cur_row = patch;

  float tmp_diff = 0;

  for (uint r = 0; r < patch_num_rows; ++r, patch_cur_row += patch_stride)
  {
    for (uint c = 0; c < patch_num_cols; ++c)
    {
      tmp_diff = mean_and_std_dev.x - patch_cur_row[c];
      mean_and_std_dev.y += tmp_diff * tmp_diff;
    }
  }

  mean_and_std_dev.y = max(sqrt(mean_and_std_dev.y / (num_pix_in_patch - 1)), 1.0e-6f);

  return mean_and_std_dev;
}

__kernel void FixedImagePatchStats(__global const float* fixed_img,
                                   const ulong  img_num_cols,
                                   const ulong num_patches,
                                   __global float2* fixed_img_patch_means_and_std_devs,
                                   __global const uint4* patch_start_stop_infos)
{
  const uint patch_idx = get_global_id(0);

  if (patch_idx < num_patches)
  {
    const uint4 cur_patch_info = patch_start_stop_infos[patch_idx];

    const uint patch_start_row = cur_patch_info.x;
    const uint patch_start_col = cur_patch_info.y;
    const uint patch_stop_row  = cur_patch_info.z;
    const uint patch_stop_col  = cur_patch_info.w;

    fixed_img_patch_means_and_std_devs[patch_idx] = xregComputePatchStats(
                                                     fixed_img + (img_num_cols * patch_start_row) + patch_start_col,
                                                     img_num_cols,
                                                     patch_stop_row - patch_start_row + 1,
                                                     patch_stop_col - patch_start_col + 1);
  }
}

__kernel void ProcFixedImagePatches(__global const float* fixed_img,
                                    __global float* proc_fixed_img_patches,
                                    const ulong  img_num_cols,
                                    const ulong  proc_patch_stride,
                                    const ulong num_patches,
                                    __global const float2* fixed_img_patch_means_and_std_devs,
                                    __global const uint4* patch_start_stop_infos)
{
  const uint patch_idx = get_global_id(0);

  if (patch_idx < num_patches)
  {
    const uint4 cur_patch_info = patch_start_stop_infos[patch_idx];

    const uint patch_start_row = cur_patch_info.x;
    const uint patch_start_col = cur_patch_info.y;
    const uint patch_stop_row  = cur_patch_info.z;
    const uint patch_stop_col  = cur_patch_info.w;

    const uint nr = patch_stop_row - patch_start_row + 1;
    const uint nc = patch_stop_col - patch_start_col + 1;

    __global const float* cur_fixed_row = fixed_img + (patch_start_row * img_num_cols) + patch_start_col;

    __global float* cur_proc_row = proc_fixed_img_patches + (patch_idx * proc_patch_stride);

    const float2 mean_and_std_dev = fixed_img_patch_means_and_std_devs[patch_idx];

    const float s = mean_and_std_dev.y * nr * nc;

    for (uint r = 0; r < nr; ++r, cur_proc_row += nc, cur_fixed_row += img_num_cols)
    {
      for (uint c = 0; c < nc; ++c)
      {
        cur_proc_row[c] = (cur_fixed_row[c] - mean_and_std_dev.x) / s;
      }
    }
  }
}

__kernel void ProcMovImagePatches(__global const  float* proc_fixed_img_patches,
                                  const ulong proc_fixed_patch_stride,
                                  const ulong num_patches,
                                  const uint num_imgs,
                                  const uint mov_proj_off,
                                  const ulong img_num_cols,
                                  const ulong img_num_rows,
                                  __global const uint4* patch_start_stop_infos,
                                  __global const float* mov_imgs,
                                  __global float* patch_nccs,
                                  __global const ulong* global_patch_idx_lut)
{
  const uint patch_idx_inc = get_global_size(0);
  const uint img_idx       = get_global_id(1);

  if (img_idx < num_imgs)
  {
    __global const float* cur_mov_img  = mov_imgs + ((mov_proj_off + img_idx) * img_num_rows * img_num_cols);
    __global       float* cur_img_nccs = patch_nccs + (img_idx * num_patches);

    for (uint local_patch_idx = get_global_id(0); local_patch_idx < num_patches; local_patch_idx += patch_idx_inc)
    {
      const ulong global_patch_idx = global_patch_idx_lut[local_patch_idx];

      __global const float* cur_proc_fixed_patch = proc_fixed_img_patches + (global_patch_idx * proc_fixed_patch_stride);

      const uint4 cur_patch_info = patch_start_stop_infos[global_patch_idx];

      const uint patch_start_row = cur_patch_info.x;
      const uint patch_start_col = cur_patch_info.y;
      const uint patch_stop_row  = cur_patch_info.z;
      const uint patch_stop_col  = cur_patch_info.w;

      const uint nr = patch_stop_row - patch_start_row + 1;
      const uint nc = patch_stop_col - patch_start_col + 1;

      __global const float* cur_mov_patch_row = cur_mov_img + (img_num_cols * patch_start_row) + patch_start_col;

      const float2 mov_patch_mean_std_dev = xregComputePatchStats(cur_mov_patch_row, img_num_cols, nr, nc);

      float ncc = 0;

      for (uint r = 0; r < nr; ++r, cur_mov_patch_row += img_num_cols, cur_proc_fixed_patch += nc)
      {
        for (uint c = 0; c < nc; ++c)
        {
          ncc += cur_proc_fixed_patch[c] * (cur_mov_patch_row[c] - mov_patch_mean_std_dev.x);
        }
      }

      ncc /= mov_patch_mean_std_dev.y;

      cur_img_nccs[local_patch_idx] = 1 - ncc;
    }
  }
}

);

}  // un-named

xreg::ImgSimMetric2DPatchNCCOCL::ImgSimMetric2DPatchNCCOCL(const boost::compute::device& dev)
  : ImgSimMetric2DOCL(dev)
{ }

xreg::ImgSimMetric2DPatchNCCOCL::ImgSimMetric2DPatchNCCOCL(const boost::compute::context& ctx,
                                                           const boost::compute::command_queue& queue)
  : ImgSimMetric2DOCL(ctx, queue)
{ }

void xreg::ImgSimMetric2DPatchNCCOCL::allocate_resources()
{
  namespace bc = boost::compute;

  xregASSERT(this->patch_radius_ > 0);

  // unsupported options when running on the GPU:
  xregASSERT(!this->use_mask_for_patch_stats_);

  auto itk_size = this->fixed_img_->GetLargestPossibleRegion().GetSize();

  img_num_rows_ = itk_size[1];
  img_num_cols_ = itk_size[0];

  // create the initial, full, set of patches

  // the mask is only used here to compute initial weights on the CPU, so we do not require
  // any other previous processing done to it
  const bool use_mask = this->mask_;

	cv::Mat ocv_mask;
  if (this->mask_)
  {
    ocv_mask = ShallowCopyItkToOpenCV(this->mask_.GetPointer());
  }
  this->setup_patches(img_num_rows_, img_num_cols_,
                      use_mask ? &ocv_mask : nullptr,
                      this->num_mov_imgs_);

  const size_type num_patches = this->patch_infos_.size();

  // compute the patch bounds on the gpu device
  {
    patch_start_stops_dev_.reset(new DevBufUInt4(this->ctx_));

    std::vector<bc::uint4_> patch_start_stops_host(num_patches);

    fixed_img_proc_patches_max_len_ = 0;

    for (size_type patch_idx = 0; patch_idx < num_patches; ++patch_idx)
    {
      const auto& cur_patch_info = this->patch_infos_[patch_idx];

      auto& cur_start_stop = patch_start_stops_host[patch_idx];

      cur_start_stop[0] = cur_patch_info.start_row;
      cur_start_stop[1] = cur_patch_info.start_col;
      cur_start_stop[2] = cur_patch_info.stop_row;
      cur_start_stop[3] = cur_patch_info.stop_col;

      fixed_img_proc_patches_max_len_ = std::max(fixed_img_proc_patches_max_len_,
                                             (cur_patch_info.stop_row - cur_patch_info.start_row + 1) *
                                             (cur_patch_info.stop_col - cur_patch_info.start_col + 1));
    }

    patch_start_stops_dev_->assign(patch_start_stops_host.begin(), patch_start_stops_host.end(), this->queue_);
  }

  // compile, and create custom kernels
  bc::program prog = bc::program::create_with_source(kPATCH_NCC_OPENCL_SRC, this->ctx_);

  try
  {
    prog.build();
  }
  catch (bc::opencl_error &)
  {
    std::cerr << "OpenCL Kernel Compile Error (ImgSimMetric2DPatchNCCOCL):\n"
              << prog.build_log() << std::endl;
    throw;
  }

  fixed_img_stats_krnl_        = prog.create_kernel("FixedImagePatchStats");
  fixed_img_proc_patches_krnl_ = prog.create_kernel("ProcFixedImagePatches");
  proc_mov_img_patches_krnl_   = prog.create_kernel("ProcMovImagePatches");

  fixed_img_stats_proc_done_ = false;

  // this will result in the process_mask method being called, which may trigger some weight recomputation,
  // and also fixed image statistics/pre-processing, which is why we have previously setup patches
  ImgSimMetric2DOCL::allocate_resources();
}

void xreg::ImgSimMetric2DPatchNCCOCL::compute()
{
  namespace bc  = boost::compute;
  namespace vcl = viennacl;

  this->pre_compute();

  // this is the current number patches to be used, e.g. if a random subset of 10 patches should be used
  // this number is 10
  const size_type num_patches = this->num_patches();

  const bool use_mask = this->mask_;

  cv::Mat ocv_mask;
  if (this->mask_)
  {
    ocv_mask = ShallowCopyItkToOpenCV(this->mask_.GetPointer());
  }

  // This uses some state to determine if the weights actually need to be recomputed
  this->compute_weights(this->mask_ ? &ocv_mask : nullptr);

  if (!this->do_not_update_patch_inds_to_use_)
  {
    // this is where random patches are sampled
    this->patch_inds_to_use_ = this->patch_indices_to_use();
  }

  // copy the patch indices and weights to use onto the GPU
  patch_inds_to_use_host_.clear();

  for (const auto& p : this->patch_inds_to_use_)
  {
    patch_inds_to_use_host_.push_back(p);
  }

  bc::copy(patch_inds_to_use_host_.begin(), patch_inds_to_use_host_.end(),
           patch_inds_to_use_dev_->begin(), this->queue_);

  wgts_to_use_host_.clear();

  Scalar tot_wgt = 0;
  if (this->weight_patch_sims_in_combine_)
  {
    for (const auto& p : this->patch_inds_to_use_)
    {
      const auto& w = this->patch_infos_[p].weight;

      tot_wgt += w;

      wgts_to_use_host_.push_back(w);
    }
  }
  else
  {
    wgts_to_use_host_.assign(num_patches, 1);
    tot_wgt = num_patches;
  }

  if (this->compute_mean_of_patch_sims_ || this->weight_patch_sims_in_combine_)
  {
    for (auto& w : wgts_to_use_host_)
    {
      w /= tot_wgt;
    }
  }

  bc::copy(wgts_to_use_host_.begin(), wgts_to_use_host_.end(),
           wgts_to_use_dev_->begin(), this->queue_);

  // TODO: it would be really nice to have to some state that avoids transferring the
  //       patches used and/or weights when they are unchanged.

  proc_mov_img_patches_krnl_.set_arg(0, *proc_fixed_img_patches_dev_);
  proc_mov_img_patches_krnl_.set_arg(1, bc::ulong_(fixed_img_proc_patches_max_len_));
  proc_mov_img_patches_krnl_.set_arg(2, bc::ulong_(num_patches));
  proc_mov_img_patches_krnl_.set_arg(3, bc::uint_(this->num_mov_imgs_));
  proc_mov_img_patches_krnl_.set_arg(4, bc::uint_(this->proj_off_));
  proc_mov_img_patches_krnl_.set_arg(5, bc::ulong_(img_num_cols_));
  proc_mov_img_patches_krnl_.set_arg(6, bc::ulong_(img_num_rows_));
  proc_mov_img_patches_krnl_.set_arg(7, *patch_start_stops_dev_);
  proc_mov_img_patches_krnl_.set_arg(8, *this->mov_imgs_buf_);
  proc_mov_img_patches_krnl_.set_arg(9, *patch_nccs_dev_);
  proc_mov_img_patches_krnl_.set_arg(10, *patch_inds_to_use_dev_);

  std::array<std::size_t,2> global_size = { num_patches, this->num_mov_imgs_ };
  this->queue_.enqueue_nd_range_kernel(proc_mov_img_patches_krnl_, 2, nullptr, global_size.data(), nullptr).wait();

  // compute weighted sums, averages, whichever

  vcl::matrix<float> patch_nccs_mat(patch_nccs_dev_->get_buffer().get(),
                                    this->num_mov_imgs_, num_patches);

  vcl::vector<float> wgts_vec(wgts_to_use_dev_->get_buffer().get(), num_patches);

  vcl::vector<float> sims_vec(sim_vals_dev_->get_buffer().get(), this->num_mov_imgs_);

  vcl::linalg::prod_impl(patch_nccs_mat, wgts_vec, sims_vec);

  bc::copy(sim_vals_dev_->begin(), sim_vals_dev_->end(), this->sim_vals_.begin(), this->queue_);
}

void xreg::ImgSimMetric2DPatchNCCOCL::process_mask()
{
  namespace bc = boost::compute;

  // NOTE: this call is commented out, because we do not need the mask represented as a float image
  // and moved to the GPU for this sim metric
  //ImgSimMetric2DOCL::process_mask();

  const size_type num_patches = this->patch_infos_.size();

  if (!fixed_img_stats_proc_done_)
  {
    // mean and std. devs of fixed image patches
    fixed_img_patch_stats_dev_.reset(new DevBufFloat2(this->ctx_));
    fixed_img_patch_stats_dev_->resize(num_patches, this->queue_);

    fixed_img_stats_krnl_.set_arg(0, *this->fixed_img_ocl_buf_);
    fixed_img_stats_krnl_.set_arg(1, bc::ulong_(img_num_cols_));
    fixed_img_stats_krnl_.set_arg(2, bc::ulong_(num_patches));
    fixed_img_stats_krnl_.set_arg(3, *fixed_img_patch_stats_dev_);
    fixed_img_stats_krnl_.set_arg(4, *patch_start_stops_dev_);

    std::size_t global_size = num_patches;

    this->queue_.enqueue_nd_range_kernel(fixed_img_stats_krnl_, 1, nullptr, &global_size, nullptr).wait();

    // pre-process the fixed image patches using the mean and std. devs.
    proc_fixed_img_patches_dev_.reset(new DevBuf(this->ctx_));
    proc_fixed_img_patches_dev_->resize(num_patches * fixed_img_proc_patches_max_len_, this->queue_);

    fixed_img_proc_patches_krnl_.set_arg(0, *this->fixed_img_ocl_buf_);
    fixed_img_proc_patches_krnl_.set_arg(1, *proc_fixed_img_patches_dev_);
    fixed_img_proc_patches_krnl_.set_arg(2, bc::ulong_(img_num_cols_));
    fixed_img_proc_patches_krnl_.set_arg(3, bc::ulong_(fixed_img_proc_patches_max_len_));
    fixed_img_proc_patches_krnl_.set_arg(4, bc::ulong_(num_patches));
    fixed_img_proc_patches_krnl_.set_arg(5, *fixed_img_patch_stats_dev_);
    fixed_img_proc_patches_krnl_.set_arg(6, *patch_start_stops_dev_);

    this->queue_.enqueue_nd_range_kernel(fixed_img_proc_patches_krnl_, 1, nullptr, &global_size, nullptr).wait();

    // allocate maximum capacity buffers for storing which patches to use
    // and the patch weights

    patch_inds_to_use_host_.reserve(num_patches);
    patch_inds_to_use_dev_.reset(new DevBufULong(this->ctx_));
    patch_inds_to_use_dev_->resize(num_patches, this->queue_);

    wgts_to_use_host_.reserve(num_patches);
    wgts_to_use_dev_.reset(new DevBuf(this->ctx_));
    wgts_to_use_dev_->resize(num_patches, this->queue_);

    patch_nccs_dev_.reset(new DevBuf(this->ctx_));
    patch_nccs_dev_->resize(num_patches * this->num_mov_imgs_, this->queue_);

    sim_vals_dev_.reset(new DevBuf(this->ctx_));
    sim_vals_dev_->resize(this->num_mov_imgs_, this->queue_);

    fixed_img_stats_proc_done_ = true;
  }

  // TODO: check about using the fixed image patch std. devs as weights

  const bool use_mask = this->mask_;

  cv::Mat ocv_mask;
  if (this->mask_)
  {
    ocv_mask = ShallowCopyItkToOpenCV(this->mask_.GetPointer());
  }

  // the mask has changed, we need to make sure the weights are recomputed
  this->need_to_recompute_weights_ = true;
  this->compute_weights(use_mask ? &ocv_mask : nullptr);
}