#include <NTL/mat_lzz_pE.h>
#include <NTL/BasicThreadPool.h>

NTL_START_IMPL

// ==========================================

#define PAR_THRESH (40000.0)

static double
zz_pE_SizeInWords()
{
   return deg(zz_pE::modulus());
}

static
void mul_aux(Mat<zz_pE>& X, const Mat<zz_pE>& A, const Mat<zz_pE>& 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 zz_p_context;
   zz_p_context.save();
   zz_pEContext zz_pE_context;
   zz_pE_context.save();
   double sz = zz_pE_SizeInWords();

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

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

   zz_p_context.restore();
   zz_pE_context.restore();

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

   Vec<zz_pE> 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 mul(mat_zz_pE& X, const mat_zz_pE& A, const mat_zz_pE& B)
{
   if (&X == &A || &X == &B) {
      mat_zz_pE tmp;
      mul_aux(tmp, A, B);
      X = tmp;
   }
   else
      mul_aux(X, A, B);
}

void inv(zz_pE& d, Mat<zz_pE>& X, const Mat<zz_pE>& 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_pXModulus& G = zz_pE::modulus();

   zz_pX t1, t2;
   zz_pX pivot;
   zz_pX pivot_inv;

   Vec< Vec<zz_pX> > M;
   // scratch space

   M.SetLength(n);
   for (long i = 0; i < n; i++) {
      M[i].SetLength(n);
      for (long j = 0; j < n; j++) {
         M[i][j].SetMaxLength(2*deg(G)-1);
         M[i][j] = rep(A[i][j]);
      }
   }

   zz_pX det;
   det = 1;

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

   zz_pContext zz_p_context;
   zz_p_context.save();
   double sz = zz_pE_SizeInWords();

   bool seq = double(n)*double(n)*sz*sz < 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], G);
         if (pivot != 0) {
            InvMod(pivot_inv, pivot, G);
            pos = i;
            break;
         }
      }

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

         MulMod(det, det, pivot, G);

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

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

         zz_p_context.restore();

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

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

            zz_pX *x = &M[i][0];
            rem(t1, x[k], G);
            negate(t1, t1);
            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_pX *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);
}

static
void solve_impl(zz_pE& d, Vec<zz_pE>& X,
                const Mat<zz_pE>& A, const Vec<zz_pE>& 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_pX t1, t2;

   const zz_pXModulus& G = zz_pE::modulus();

   Vec< Vec<zz_pX> > M;

   M.SetLength(n);

   for (long i = 0; i < n; i++) {
      M[i].SetLength(n+1);
      for (long j = 0; j < n; j++) M[i][j].SetMaxLength(2*deg(G)-1);

      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_pX det;
   set(det);

   zz_pContext zz_p_context;
   zz_p_context.save();
   double sz = zz_pE_SizeInWords();

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

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

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

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

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

         bool seq =
            double(n-(k+1))*(n-(k+1))*sz*sz < PAR_THRESH;

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

         zz_p_context.restore();

         zz_pX 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_pX *x = M[i].elts() + (k+1);
            zz_pX *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_pE& d, Vec<zz_pE>& x,
               const Mat<zz_pE>& A, const Vec<zz_pE>& b)
{
   solve_impl(d, x, A, b, true);
}

void solve(zz_pE& d, const Mat<zz_pE>& A,
               Vec<zz_pE>& x, const Vec<zz_pE>& b)
{
   solve_impl(d, x, A, b, false);
}

long gauss(Mat<zz_pE>& M_in, long w)
{
   zz_pX t1, t2;
   zz_pX piv;

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

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

   const zz_pXModulus& G = zz_pE::modulus();

   Vec< Vec<zz_pX> > M;

   M.SetLength(n);
   for (long i = 0; i < n; i++) {
      M[i].SetLength(m);
      for (long j = 0; j < m; j++) {
         M[i][j].SetLength(2*deg(G)-1);
         M[i][j] = rep(M_in[i][j]);
      }
   }

   zz_pContext zz_p_context;
   zz_p_context.save();
   double sz = zz_pE_SizeInWords();

   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], G);
         M[i][k] = t1;
         if (pos == -1 && !IsZero(t1)) {
            pos = i;
         }
      }

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

         InvMod(piv, M[l][k], G);
         negate(piv, piv);

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

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

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

         zz_p_context.restore();

         zz_pX 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, G);

            clear(M[i][k]);

            zz_pX *x = M[i].elts() + (k+1);
            zz_pX *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_pE>& M)
{
   return gauss(M, M.NumCols());
}

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

void kernel(Mat<zz_pE>& X, const Mat<zz_pE>& A)
{
   long m = A.NumRows();
   long n = A.NumCols();

   const zz_pXModulus& G = zz_pE::modulus();

   Mat<zz_pE> 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_pE> 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]);
   }

   zz_pEContext zz_pE_context;
   zz_pE_context.save();
   zz_pContext zz_p_context;
   zz_p_context.save();
   double sz = zz_pE_SizeInWords();

   bool seq =
      double(m-r)*double(r)*double(r)*sz*sz < PAR_THRESH;

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

   zz_p_context.restore();
   zz_pE_context.restore();

   zz_pX t1, t2;
   zz_pE T3;

   for (long k = first; k < last; k++) {
      Vec<zz_pE>& 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 determinant(zz_pE& d, const Mat<zz_pE>& M_in)
{
   zz_pX t1, t2;

   const zz_pXModulus& G = zz_pE::modulus();

   long n = M_in.NumRows();

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

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

   Vec< Vec<zz_pX> > M;

   M.SetLength(n);
   for (long i = 0; i < n; i++) {
      M[i].SetLength(n);
      for (long j = 0; j < n; j++) {
         M[i][j].SetMaxLength(2*deg(G)-1);
         M[i][j] = rep(M_in[i][j]);
      }
   }

   zz_pX det;
   set(det);

   zz_pContext zz_p_context;
   zz_p_context.save();
   double sz = zz_pE_SizeInWords();

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

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

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

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

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

         bool seq =
            double(n-(k+1))*(n-(k+1))*sz*sz < PAR_THRESH;

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

         zz_p_context.restore();

         zz_pX 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_pX *x = M[i].elts() + (k+1);
            zz_pX *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);
}

// ==========================================

void add(mat_zz_pE& X, const mat_zz_pE& A, const mat_zz_pE& 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_pE& X, const mat_zz_pE& A, const mat_zz_pE& 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_pE& X, const mat_zz_pE& 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_pE& x, const mat_zz_pE& A, const vec_zz_pE& 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_pX 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_pE& x, const mat_zz_pE& A, const vec_zz_pE& b)
{
   if (&b == &x || A.alias(x)) {
      vec_zz_pE tmp;
      mul_aux(tmp, A, b);
      x = tmp;
   }
   else
      mul_aux(x, A, b);
}

static
void mul_aux(vec_zz_pE& x, const vec_zz_pE& a, const mat_zz_pE& 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_pX 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_pE& x, const vec_zz_pE& a, const mat_zz_pE& B)
{
   if (&a == &x) {
      vec_zz_pE tmp;
      mul_aux(tmp, a, B);
      x = tmp;
   }
   else
      mul_aux(x, a, B);

}

void ident(mat_zz_pE& 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));
}

long IsIdent(const mat_zz_pE& 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_pE& X, const mat_zz_pE& 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_pE 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);
   }
}

void mul(mat_zz_pE& X, const mat_zz_pE& A, const zz_pE& b_in)
{
   zz_pE 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_pE& X, const mat_zz_pE& 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_pE& X, const mat_zz_pE& 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_pE& X, long n, const zz_pE& d_in)
{
   zz_pE 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_pE& A, long n, const zz_pE& 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_pE& 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_pE& x)
{
   long n = x.NumRows();
   long i;
   for (i = 0; i < n; i++)
      clear(x[i]);
}

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

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

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

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

vec_zz_pE operator*(const mat_zz_pE& a, const vec_zz_pE& b)
{
   vec_zz_pE res;
   mul_aux(res, a, b);
   NTL_OPT_RETURN(vec_zz_pE, res);
}

vec_zz_pE operator*(const vec_zz_pE& a, const mat_zz_pE& b)
{
   vec_zz_pE res;
   mul_aux(res, a, b);
   NTL_OPT_RETURN(vec_zz_pE, res);
}

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

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

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

   mat_zz_pE 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_pE& x, long n, long m)
{
   x.SetDims(n, m);
   for (long i = 0; i < n; i++) random(x[i], m);
}

NTL_END_IMPL


//=======================================================================
// Copyright (c) 2014 Andrzej Pacuk, Piotr Wygocki
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
/**
 * @file lsh_functions.hpp
 * @brief
 * @author Andrzej Pacuk, Piotr Wygocki
 * @version 1.0
 * @date 2014-10-07
 */
#ifndef PAAL_LSH_FUNCTIONS_HPP
#define PAAL_LSH_FUNCTIONS_HPP

#include "paal/data_structures/ublas_traits.hpp"
#include "paal/utils/functors.hpp"
#include "paal/utils/type_functions.hpp"

#include <boost/range/algorithm/equal.hpp>
#include <boost/range/algorithm/generate.hpp>
#include <boost/range/algorithm/min_element.hpp>
#include <boost/range/algorithm_ext/iota.hpp>
#include <boost/range/counting_range.hpp>
#include <boost/range/empty.hpp>
#include <boost/range/iterator.hpp>
#include <boost/range/numeric.hpp>
#include <boost/range/size.hpp>
#include <boost/range/adaptor/transformed.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/vector_sparse.hpp>
#include <boost/numeric/ublas/vector_expression.hpp>

#include <cassert>
#include <cmath>
#include <cstddef>
#include <random>
#include <type_traits>
#include <utility>

namespace paal {

namespace lsh {

/**
 * @brief simple hash function
 */
template <typename IntType = std::size_t>
class projection_hash_function {
    IntType m_chosen_position;

public:
    ///serialize
    template<class Archive>
    void serialize(Archive & ar, const unsigned int version) {
        ar & m_chosen_position;
    }

    ///default constructor
    projection_hash_function() = default;

    /**
     * @brief constructor
     *
     * @param chosen_position
     *
     */
    projection_hash_function(IntType chosen_position) :
            m_chosen_position(chosen_position) {
    }

    ///operator==
    bool operator==(projection_hash_function const & other) const {
        return m_chosen_position == other.m_chosen_position;
    }

    /**
     * @brief operator()
     *
     * @tparam Range
     * @param range
     *
     * @return
     */
    template <typename Range>
    //TODO change to decltype(auto), when it starts working
    auto operator()(Range &&range) const ->
            range_to_elem_t<decltype(range)> {
        //TODO takes to long, use boost assert
        //assert(m_chosen_position < boost::size(range));
        return range[m_chosen_position];
    }
};

/**
 * @brief Factory class for projection_hash_function
 */
template <typename RandomEngine = std::default_random_engine,
          typename IntType = std::size_t>
class random_projection_hash_function_generator {
    mutable RandomEngine m_generator;
    using distribution_t = std::uniform_int_distribution<IntType>;
    mutable distribution_t m_distribution;

    static IntType convert_to_inclusive(IntType upper_bound) {
        assert(upper_bound > 0);
        return upper_bound - 1;
    }

public:
    /**
     * @brief constructor
     *
     * @param range_size
     * @param random_engine
     */
    random_projection_hash_function_generator(IntType range_size,
            RandomEngine random_engine = RandomEngine{}) :
            m_generator(std::move(random_engine)),
            m_distribution(0, convert_to_inclusive(range_size)) {
    }

    /**
     * @brief operator()
     *
     * @return new projection_hash_function
     */
    //TODO change to auto, when it starts working
    projection_hash_function<IntType> operator()() const {
        return projection_hash_function<IntType>(m_distribution(m_generator));
    }
};

using hamming_hash_function_generator =
    random_projection_hash_function_generator<>;

/**
 */
/// hash_function for l_p distance for p in range (0,2]
template <typename FloatType = double>
class l_p_hash_function {
public:
    using r_param_t = boost::numeric::ublas::vector<FloatType>;

private:
    r_param_t m_r;
    FloatType m_b;
    FloatType m_w;

public:

    ///serialize
    template<class Archive>
    void serialize(Archive & ar, const unsigned int version) {
        ar & m_r;
        ar & m_b;
        ar & m_w;
    }

    /**
     * @brief constructor
     *
     * @param r r elements has to be generated by p-stable distribution
     * @param b b is a random value from [0, w)
     * @param w w should be much greater than expected points neighborhood
     * radius
     */
    l_p_hash_function(r_param_t r,
                      FloatType b, FloatType w) :
            m_r(std::move(r)), m_b(b), m_w(w) {
    }

    ///default constructor
    l_p_hash_function() = default;

    ///operator==
    bool operator==(l_p_hash_function const & other) const {
        return boost::equal(m_r, other.m_r) &&
               m_b == other.m_b &&
               m_w == other.m_w;
    }

    /**
     * @brief operator()
     *
     * @tparam Range
     * @param range Range modeling boost uBLAS VectorExpression concept
     *
     * @return
     */
    template <typename Range>
    auto operator()(Range &&range) const {
        //TODO use concept check with VectorExpressionConcept<Range>,
        //when it works

        //TODO takes to long, use boost assert
        //assert(boost::size(m_r) == boost::size(range));

        using boost::numeric::ublas::inner_prod;
        auto inner_product = inner_prod(m_r, range);
        return std::floor((m_b + inner_product) / m_w);
    }

};

/**
 * @brief Factory class for l_p_hash_function
 */
template <typename FloatType = double,
          typename RandomEngine = std::default_random_engine,
          typename Distribution = std::normal_distribution<FloatType>>
class l_p_hash_function_generator {
    std::size_t m_range_size;
    FloatType m_w;
    mutable RandomEngine m_generator;
    mutable Distribution m_r_distribution;
    mutable std::uniform_real_distribution<FloatType> m_b_distribution;

public:
    /**
     * @brief constructor
     *
     * @param range_size
     * @param w
     * @param random_engine
     */
    l_p_hash_function_generator(std::size_t range_size,
                                FloatType w,
                                RandomEngine random_engine = RandomEngine{}) :
            m_range_size(range_size),
            m_w(w),
            m_generator(std::move(random_engine)),
            m_b_distribution(FloatType{}, w) {
    }

    /**
     * @brief operator()
     *
     * @return new l_p_hash_function
     */
    l_p_hash_function<FloatType> operator()() const {
        typename l_p_hash_function<FloatType>::r_param_t r(m_range_size);
        boost::generate(r, [&]() {
            return m_r_distribution(m_generator);
        });

        FloatType b = m_b_distribution(m_generator);
        return l_p_hash_function<FloatType>(std::move(r), b, m_w);
    }
};

/// Cauchy distribution is 1-stable
template <typename FloatType = double,
          typename RandomEngine = std::default_random_engine>
using l_1_hash_function_generator =
    l_p_hash_function_generator<FloatType, RandomEngine,
                                std::cauchy_distribution<FloatType>>;

/// Gaussian distribution is 2-stable
template <typename FloatType = double,
          typename RandomEngine = std::default_random_engine>
using l_2_hash_function_generator =
    l_p_hash_function_generator<FloatType, RandomEngine>;

///min-wise independent permutations locality sensitive hashing (Jaccard)
class min_hash_function {
    //permutation
    std::vector<std::size_t> m_perm;

public:

    ///serialize
    template<class Archive>
    void serialize(Archive & ar, const unsigned int version) {
        ar & m_perm;
    }

    ///constructor
    template <typename RandomEngine>
    min_hash_function(std::size_t set_element_upper_bound, RandomEngine && rng)
            : m_perm(set_element_upper_bound)  {
        boost::iota(m_perm, 0);
        std::shuffle(m_perm.begin(), m_perm.end(), rng);
    }

    ///default constructor
    min_hash_function() = default;

    ///operator==
    bool operator==(min_hash_function const & other) const {
        return boost::equal(m_perm, other.m_perm);
    }

    /**
     * @brief operator()
     *
     * @tparam Range
     * @param range forward range modeling boost uBLAS sparse VectorExpression concept,
     * set is represented as set of indexes of range elements
     *
     * @return
     */
    template <typename Range>
    auto operator()(Range &&range) const {
        static_assert(data_structures::is_sparse_row<Range>::value, "vector must be sparse");
        using iter_t = typename boost::range_iterator<Range>::type;

        auto perm_function = [&](iter_t it) {
            auto set_elem = it.index();
            assert(set_elem < m_perm.size());
            return m_perm[set_elem];
        };
        auto perm_range =
                boost::counting_range(range)
                | boost::adaptors::transformed(utils::make_assignable_functor(perm_function));

        assert(!boost::empty(perm_range));
        return *boost::min_element(perm_range);
    }
};

///Factory class for min_hash_function
template <typename RandomEngine = std::default_random_engine>
class min_hash_function_generator {
    std::size_t m_range_size;
    mutable RandomEngine m_generator;

public:
    ///constructor
    min_hash_function_generator(std::size_t range_size,
                                RandomEngine random_engine = RandomEngine{}) :
            m_range_size(range_size),
            m_generator(std::move(random_engine)) { }

    ///operator()
    min_hash_function operator()() const {
        return min_hash_function(m_range_size, m_generator);
    }
};
using jaccard_hash_function_generator =
    min_hash_function_generator<>;

} //! lsh

} //! paal

#endif // PAAL_LSH_FUNCTIONS_HPP


//==================================================================================================
/*!
  @file
  @copyright 2015 NumScale SAS
  @copyright 2015 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_FUNCTION_SCALAR_STIRLING_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_FUNCTION_SCALAR_STIRLING_HPP_INCLUDED

#include <boost/config.hpp>
#include <boost/dispatch/function/overload.hpp>
#include <boost/simd/arch/common/detail/generic/stirling_kernel.hpp>
#include <boost/simd/constant/half.hpp>
#include <boost/simd/constant/inf.hpp>
#include <boost/simd/constant/one.hpp>
#include <boost/simd/constant/real.hpp>
#include <boost/simd/constant/sqrt_2pi.hpp>
#include <boost/simd/function/scalar/exp.hpp>
#include <boost/simd/function/scalar/fma.hpp>
#include <boost/simd/function/scalar/is_eqz.hpp>
#include <boost/simd/function/scalar/pow.hpp>
#include <boost/simd/function/scalar/rec.hpp>
#ifndef BOOST_SIMD_NO_INVALIDS
#include <boost/simd/function/scalar/is_nan.hpp>
#endif

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;
  BOOST_DISPATCH_OVERLOAD ( stirling_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_< bd::floating_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() (A0 a0) const BOOST_NOEXCEPT
    {
      const A0 Stirlinglargelim = Real<A0, 0x4065800000000000ULL, 0X420C28F3UL>();// 172, 35.0399895f
      const A0 Stirlingsplitlim = Real<A0, 0X4061E083BA3443D4ULL, 0X41D628F6UL>();// 143.01608, 26.77f
      #ifndef BOOST_SIMD_NO_INVALIDS
      if (is_nan(a0)) return a0;
      #endif
      if (a0 > Stirlinglargelim) return Inf<A0>();
      A0 w = rec(a0);
      w = fma(w,detail::stirling_kernel<A0>::stirling1(w), One<A0>());
      A0 y = exp(-a0);
      if(is_eqz(y)) return Inf<A0>();
      A0 z =  a0 - Half<A0>();
      if( a0 >= Stirlingsplitlim )
      { /* Avoid overflow in pow() */
        const A0 v = pow(a0,z*Half<A0>());
        y *= v;
        y *= v;
      }
      else
      {
        y *= pow( a0, z );
      }
      y *= Sqrt_2pi<A0>()*w;
      return y;
    }
  };
} } }

#endif


//=======================================================================
// Copyright 2001 Jeremy G. Siek, Andrew Lumsdaine, Lie-Quan Lee,
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
#include <deque>                // to store the vertex ordering
#include <vector>
#include <list>
#include <iostream>
#include <boost/graph/vector_as_graph.hpp>
#include <boost/graph/topological_sort.hpp>

int
main()
{
  using namespace boost;
  const char *tasks[] = {
    "pick up kids from school",
    "buy groceries (and snacks)",
    "get cash at ATM",
    "drop off kids at soccer practice",
    "cook dinner",
    "pick up kids from soccer",
    "eat dinner"
  };
  const int n_tasks = sizeof(tasks) / sizeof(char *);

  std::vector < std::list < int > > g(n_tasks);
  g[0].push_back(3);
  g[1].push_back(3);
  g[1].push_back(4);
  g[2].push_back(1);
  g[3].push_back(5);
  g[4].push_back(6);
  g[5].push_back(6);

  std::deque < int >topo_order;

  topological_sort(g, std::front_inserter(topo_order),
                   vertex_index_map(identity_property_map()));

  int n = 1;
  for (std::deque < int >::iterator i = topo_order.begin();
       i != topo_order.end(); ++i, ++n)
    std::cout << tasks[*i] << std::endl;

  return EXIT_SUCCESS;
}


//==================================================================================================
/*!
  @file
  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_ERF_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_ERF_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{

 /*!
  @ingroup group-euler
    Function object implementing erf capabilities

   Computes the error function:
   \f$\displaystyle \frac{2}{\sqrt\pi}\int_0^{x} e^{-t^2}\mbox{d}t\f$

    @par Semantic:

    For every parameter of floating type T

    @code
    T r = erf(x);
    @endcode

    @par Decorators

    std_ for floating entries provides access to @c std::erf

    @see erfc,  erfcx

  **/
  Value erf(Value const & v0);
} }
#endif

#include <boost/simd/function/scalar/erf.hpp>
#include <boost/simd/function/simd/erf.hpp>

#endif


//==============================================================================
//         Copyright 2003 - 2011 LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2011 LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef NT2_EXPONENTIAL_FUNCTIONS_SCALAR_CBRT_HPP_INCLUDED
#define NT2_EXPONENTIAL_FUNCTIONS_SCALAR_CBRT_HPP_INCLUDED

#include <nt2/exponential/functions/cbrt.hpp>
#include <boost/simd/sdk/config.hpp>
#include <nt2/include/constants/one.hpp>
#include <nt2/include/constants/real_splat.hpp>
#include <nt2/include/constants/third.hpp>
#include <nt2/include/constants/three.hpp>
#include <nt2/include/constants/two.hpp>
#include <nt2/include/functions/scalar/abs.hpp>
#include <nt2/include/functions/scalar/bitofsign.hpp>
#include <nt2/include/functions/scalar/bitwise_or.hpp>
#include <nt2/include/functions/scalar/fast_frexp.hpp>
#include <nt2/include/functions/scalar/fast_ldexp.hpp>
#include <nt2/include/functions/scalar/is_gez.hpp>
#include <nt2/include/functions/scalar/negate.hpp>
#include <nt2/include/functions/scalar/sqr.hpp>
#include <nt2/polynomials/functions/scalar/impl/horner.hpp>
#include <nt2/sdk/meta/as_integer.hpp>

#ifndef BOOST_SIMD_NO_INFINITIES
#include <nt2/include/constants/inf.hpp>
#endif

#ifndef BOOST_SIMD_NO_DENORMALS
#include <nt2/include/constants/smallestposval.hpp>
#include <nt2/include/constants/twotomnmbo_3.hpp>
#include <nt2/include/constants/twotonmb.hpp>
#endif

namespace nt2 { namespace ext
{
  BOOST_DISPATCH_IMPLEMENT  ( cbrt_, tag::cpu_
                            , (A0)
                            , (scalar_< double_<A0> >)
                            )
  {

    typedef A0 result_type;

    NT2_FUNCTOR_CALL(1)
    {
      A0 z =  nt2::abs(a0);
#ifndef BOOST_SIMD_NO_INFINITIES
      if (z == nt2::Inf<A0>() || (z == 0)) return a0;
#else
      if (z == 0) return a0;
#endif
#ifndef BOOST_SIMD_NO_DENORMALS
      A0 f = One<A0>();
      if (z < Smallestposval<A0>())
      {
        z *= Twotonmb<A0>();
        f  = Twotomnmbo_3<A0>();
      }
#endif
      const A0 CBRT2  = double_constant< A0, 0x3ff428a2f98d728bll> ();
      const A0 CBRT4  = double_constant< A0, 0x3ff965fea53d6e3dll> ();
      const A0 CBRT2I = double_constant< A0, 0x3fe965fea53d6e3dll> ();
      const A0 CBRT4I = double_constant< A0, 0x3fe428a2f98d728bll> ();
      typedef typename meta::as_integer<A0, signed>::type int_type;
      int_type e;
      A0 x = fast_frexp(z, e);
      x = horner < NT2_HORNER_COEFF_T(A0, 5,
                            (0xbfc13c93386fdff6ll,
                             0x3fe17e1fc7e59d58ll,
                             0xbfee8a4ca3ba37b8ll,
                             0x3ff23d6ee505873all,
                             0x3fd9c0c12122a4fell)
                            ) > (x);
      const bool flag = is_gez(e);
      int_type e1 =  nt2::abs(e);
      int_type rem = e1;
      e1 /= Three<int_type>();
      rem -= e1*Three<int_type>();
      e =  negate(e1, e);
      const A0 cbrt2 = flag ? CBRT2 : CBRT2I;
      const A0 cbrt4 = flag ? CBRT4 : CBRT4I;
      A0 fact = (rem == One<int_type>()) ? cbrt2: One<A0>();
      fact = (rem == Two<int_type>() ? cbrt4 : fact);
      x = fast_ldexp(x*fact, e);
      x -= (x-z/sqr(x))*Third<A0>();
      x -= (x-z/sqr(x))*Third<A0>(); //two newton passes
#ifndef BOOST_SIMD_NO_DENORMALS
      return b_or(x, bitofsign(a0))*f;
#else
      return b_or(x, bitofsign(a0));
#endif
    }
  };

  BOOST_DISPATCH_IMPLEMENT  ( cbrt_, tag::cpu_
                            , (A0)
                            , (scalar_< single_<A0> >)
                            )
  {
    typedef A0 result_type;
    NT2_FUNCTOR_CALL(1)
    {
      A0 z =  nt2::abs(a0);
#ifndef BOOST_SIMD_NO_INFINITIES
      if (z == nt2::Inf<A0>() || (z == 0)) return a0;
#else
      if  (z == 0) return a0;
#endif
#ifndef BOOST_SIMD_NO_DENORMALS
      A0 f = One<A0>();
      if (z < Smallestposval<A0>())
      {
        z *= Twotonmb<A0>();
        f = Twotomnmbo_3<A0>();
      }
#endif
      const A0 CBRT2  = single_constant< A0, 0x3fa14518> ();
      const A0 CBRT4  = single_constant< A0, 0x3fcb2ff5> ();
      const A0 CBRT2I = single_constant< A0, 0x3f4b2ff5> ();
      const A0 CBRT4I = single_constant< A0, 0x3f214518> ();
      typedef typename meta::as_integer<A0, signed>::type int_type;
      int_type e;
      A0 x = fast_frexp(z, e);
      x = horner < NT2_HORNER_COEFF_T(A0, 5,
                                      (0xbe09e49a,
                                       0x3f0bf0fe,
                                       0xbf745265,
                                       0x3f91eb77,
                                       0x3ece0609)
                                     ) > (x);
      const bool flag = is_gez(e);
      int_type e1 =  nt2::abs(e);
      int_type rem = e1;
      e1 /= Three<int_type>();
      rem -= e1*Three<int_type>();
      e =  negate(e1, e);

      const A0 cbrt2 = flag ? CBRT2 : CBRT2I;
      const A0 cbrt4 = flag ? CBRT4 : CBRT4I;
      A0 fact = (rem ==  One<int_type>()) ? cbrt2 : One<A0>();
      fact = (rem == Two<int_type>()) ? cbrt4 : fact;
      x = fast_ldexp(x*fact, e);
      x -= (x-z/sqr(x))*Third<A0>();
#ifndef BOOST_SIMD_NO_DENORMALS
      return b_or(x, bitofsign(a0))*f;
#else
      return b_or(x, bitofsign(a0));
#endif
    }
  };
} }

#endif


#include <stan/math/prim/scal.hpp>
#include <gtest/gtest.h>
#include <boost/random/mersenne_twister.hpp>
#include <boost/math/distributions.hpp>

TEST(ProbDistributionsNegBinomial, error_check) {
  boost::random::mt19937 rng;
  EXPECT_NO_THROW(stan::math::neg_binomial_rng(6, 2, rng));
  EXPECT_NO_THROW(stan::math::neg_binomial_rng(0.5, 1, rng));
  EXPECT_NO_THROW(stan::math::neg_binomial_rng(1e9, 1, rng));

  EXPECT_THROW(stan::math::neg_binomial_rng(0, -2, rng),
                 std::domain_error);
  EXPECT_THROW(stan::math::neg_binomial_rng(6, -2, rng),
                   std::domain_error);
  EXPECT_THROW(stan::math::neg_binomial_rng(-6, -0.1, rng),
                   std::domain_error);
  EXPECT_THROW(stan::math::neg_binomial_rng(
                 stan::math::positive_infinity(), 2, rng),
                 std::domain_error);
  EXPECT_THROW(stan::math::neg_binomial_rng(
                 stan::math::positive_infinity(), 6, rng),
                 std::domain_error);
  EXPECT_THROW(stan::math::neg_binomial_rng(2,
                 stan::math::positive_infinity(), rng),
                 std::domain_error);

  std::string error_msg;
  error_msg = "neg_binomial_rng: Random number that "
              "came from gamma distribution is";
  try {
    stan::math::neg_binomial_rng(1e10, 1, rng);
    FAIL() << "neg_binomial_rng should have thrown" << std::endl;
  } catch (const std::exception& e) {
    if (std::string(e.what()).find(error_msg) == std::string::npos)
      FAIL() << "Error message is different than expected" << std::endl
             << "EXPECTED: " << error_msg << std::endl
             << "FOUND: " << e.what() << std::endl;
    SUCCEED();
  }
}

void expected_bin_sizes(double *expect, const int K,
                        const int N,
                        const double alpha, const double beta) {
  double p = 0;
  for(int i = 0 ; i < K; i++)  {
    expect[i] = N * std::exp(stan::math::neg_binomial_log(i, alpha, beta));
    p += std::exp(stan::math::neg_binomial_log(i, alpha, beta));
  }
  expect[K-1] = N * (1.0 - p);
}

TEST(ProbDistributionsNegBinomial, chiSquareGoodnessFitTest) {
  boost::random::mt19937 rng;
  double p = 0.6;
  double alpha = 5;
  double beta = p / (1 - p);
  int N = 1000;
  int K = boost::math::round(2 * std::pow(N, (1-p)));
  boost::math::chi_squared mydist(K-1);

  int loc[K - 1];
  for(int i = 1; i < K; i++)
    loc[i - 1] = i - 1;

  int count = 0;
  double bin [K];
  double expect [K];

  for(int i = 0 ; i < K; i++)
    bin[i] = 0;
  expected_bin_sizes(expect, K, N, alpha, beta);

  while (count < N) {
    int a = stan::math::neg_binomial_rng(alpha, beta, rng);
    int i = 0;
    while (i < K-1 && a > loc[i])
      ++i;
    ++bin[i];
    count++;
   }

  double chi = 0;

  for(int j = 0; j < K; j++)
    chi += ((bin[j] - expect[j]) * (bin[j] - expect[j]) / expect[j]);

  EXPECT_TRUE(chi < boost::math::quantile(boost::math::complement(mydist, 1e-6)));
}

TEST(ProbDistributionsNegBinomial, chiSquareGoodnessFitTest2) {
  boost::random::mt19937 rng;
  double p = 0.8;
  double alpha = 2.4;
  double beta = p / (1 - p);
  int N = 1000;
  int K = boost::math::round(2 * std::pow(N, (1-p)));
  boost::math::chi_squared mydist(K-1);

  int loc[K - 1];
  for(int i = 1; i < K; i++)
    loc[i - 1] = i - 1;

  int count = 0;
  double bin [K];
  double expect [K];

  for(int i = 0 ; i < K; i++)
    bin[i] = 0;
  expected_bin_sizes(expect, K, N, alpha, beta);

  while (count < N) {
    int a = stan::math::neg_binomial_rng(alpha, beta, rng);
    int i = 0;
    while (i < K-1 && a > loc[i])
      ++i;
    ++bin[i];
    count++;
   }

  double chi = 0;

  for(int j = 0; j < K; j++)
    chi += ((bin[j] - expect[j]) * (bin[j] - expect[j]) / expect[j]);

  EXPECT_TRUE(chi < boost::math::quantile(boost::math::complement(mydist, 1e-6)));
}

TEST(ProbDistributionsNegBinomial, chiSquareGoodnessFitTest3) {
  boost::random::mt19937 rng;
  double p = 0.2;
  double alpha = 0.4;
  double beta = p / (1 - p);
  int N = 1000;
  int K = boost::math::round(2 * std::pow(N, (1-p)));
  boost::math::chi_squared mydist(K-1);

  int loc[K - 1];
  for(int i = 1; i < K; i++)
    loc[i - 1] = i - 1;

  int count = 0;
  double bin [K];
  double expect [K];

  for(int i = 0 ; i < K; i++)
    bin[i] = 0;
  expected_bin_sizes(expect, K, N, alpha, beta);

  while (count < N) {
    int a = stan::math::neg_binomial_rng(alpha, beta, rng);
    int i = 0;
    while (i < K-1 && a > loc[i])
      ++i;
    ++bin[i];
    count++;
   }

  double chi = 0;

  for(int j = 0; j < K; j++)
    chi += ((bin[j] - expect[j]) * (bin[j] - expect[j]) / expect[j]);

  EXPECT_TRUE(chi < boost::math::quantile(boost::math::complement(mydist, 1e-6)));
}


/*
 * testSudoku.cpp
 * @brief develop code for Sudoku CSP solver
 * @date Jan 29, 2012
 * @author Frank Dellaert
 */

#include <CppUnitLite/TestHarness.h>
#include <gtsam/inference/Symbol.h>
#include <gtsam_unstable/discrete/CSP.h>

#include <boost/assign/std/map.hpp>
using boost::assign::insert;
#include <stdarg.h>

#include <iostream>
#include <sstream>

using namespace std;
using namespace gtsam;

#define PRINT false

/// A class that encodes Sudoku's as a CSP problem
class Sudoku : public CSP {
  size_t n_;  ///< Side of Sudoku, e.g. 4 or 9

  /// Mapping from base i,j coordinates to discrete keys:
  using IJ = std::pair<size_t, size_t>;
  std::map<IJ, DiscreteKey> dkeys_;

 public:
  /// return DiscreteKey for cell(i,j)
  const DiscreteKey& dkey(size_t i, size_t j) const {
    return dkeys_.at(IJ(i, j));
  }

  /// return Key for cell(i,j)
  Key key(size_t i, size_t j) const { return dkey(i, j).first; }

  /// Constructor
  Sudoku(size_t n, ...) : n_(n) {
    // Create variables, ordering, and unary constraints
    va_list ap;
    va_start(ap, n);
    for (size_t i = 0; i < n; ++i) {
      for (size_t j = 0; j < n; ++j) {
        // create the key
        IJ ij(i, j);
        Symbol key('1' + i, j + 1);
        dkeys_[ij] = DiscreteKey(key, n);
        // get the unary constraint, if any
        int value = va_arg(ap, int);
        if (value != 0) addSingleValue(dkeys_[ij], value - 1);
      }
      // cout << endl;
    }
    va_end(ap);

    // add row constraints
    for (size_t i = 0; i < n; i++) {
      DiscreteKeys dkeys;
      for (size_t j = 0; j < n; j++) dkeys += dkey(i, j);
      addAllDiff(dkeys);
    }

    // add col constraints
    for (size_t j = 0; j < n; j++) {
      DiscreteKeys dkeys;
      for (size_t i = 0; i < n; i++) dkeys += dkey(i, j);
      addAllDiff(dkeys);
    }

    // add box constraints
    size_t N = (size_t)sqrt(double(n)), i0 = 0;
    for (size_t I = 0; I < N; I++) {
      size_t j0 = 0;
      for (size_t J = 0; J < N; J++) {
        // Box I,J
        DiscreteKeys dkeys;
        for (size_t i = i0; i < i0 + N; i++)
          for (size_t j = j0; j < j0 + N; j++) dkeys += dkey(i, j);
        addAllDiff(dkeys);
        j0 += N;
      }
      i0 += N;
    }
  }

  /// Print readable form of assignment
  void printAssignment(const DiscreteValues& assignment) const {
    for (size_t i = 0; i < n_; i++) {
      for (size_t j = 0; j < n_; j++) {
        Key k = key(i, j);
        cout << 1 + assignment.at(k) << " ";
      }
      cout << endl;
    }
  }

  /// solve and print solution
  void printSolution() const {
    auto MPE = optimize();
    printAssignment(MPE);
  }

  // Print domain
  void printDomains(const Domains& domains) {
    for (size_t i = 0; i < n_; i++) {
      for (size_t j = 0; j < n_; j++) {
        Key k = key(i, j);
        cout << domains.at(k).base1Str();
        cout << "\t";
      }  // i
      cout << endl;
    }  // j
  }
};

/* ************************************************************************* */
TEST(Sudoku, small) {
  Sudoku csp(4,           //
             1, 0, 0, 4,  //
             0, 0, 0, 0,  //
             4, 0, 2, 0,  //
             0, 1, 0, 0);

  // optimize and check
  auto solution = csp.optimize();
  DiscreteValues expected;
  insert(expected)(csp.key(0, 0), 0)(csp.key(0, 1), 1)(csp.key(0, 2), 2)(
      csp.key(0, 3), 3)(csp.key(1, 0), 2)(csp.key(1, 1), 3)(csp.key(1, 2), 0)(
      csp.key(1, 3), 1)(csp.key(2, 0), 3)(csp.key(2, 1), 2)(csp.key(2, 2), 1)(
      csp.key(2, 3), 0)(csp.key(3, 0), 1)(csp.key(3, 1), 0)(csp.key(3, 2), 3)(
      csp.key(3, 3), 2);
  EXPECT(assert_equal(expected, solution));
  // csp.printAssignment(solution);

  // Do BP (AC1)
  auto domains = csp.runArcConsistency(4, 3);
  // csp.printDomains(domains);
  Domain domain44 = domains.at(Symbol('4', 4));
  EXPECT_LONGS_EQUAL(1, domain44.nrValues());

  // Test Creation of a new, simpler CSP
  CSP new_csp = csp.partiallyApply(domains);
  // Should only be 16 new Domains
  EXPECT_LONGS_EQUAL(16, new_csp.size());

  // Check that solution
  auto new_solution = new_csp.optimize();
  // csp.printAssignment(new_solution);
  EXPECT(assert_equal(expected, new_solution));
}

/* ************************************************************************* */
TEST(Sudoku, easy) {
  Sudoku csp(9,                          //
             0, 0, 5, 0, 9, 0, 0, 0, 1,  //
             0, 0, 0, 0, 0, 2, 0, 7, 3,  //
             7, 6, 0, 0, 0, 8, 2, 0, 0,  //

             0, 1, 2, 0, 0, 9, 0, 0, 4,  //
             0, 0, 0, 2, 0, 3, 0, 0, 0,  //
             3, 0, 0, 1, 0, 0, 9, 6, 0,  //

             0, 0, 1, 9, 0, 0, 0, 5, 8,  //
             9, 7, 0, 5, 0, 0, 0, 0, 0,  //
             5, 0, 0, 0, 3, 0, 7, 0, 0);

  // csp.printSolution(); // don't do it

  // Do BP (AC1)
  auto domains = csp.runArcConsistency(9, 10);
  // csp.printDomains(domains);
  Key key99 = Symbol('9', 9);
  Domain domain99 = domains.at(key99);
  EXPECT_LONGS_EQUAL(1, domain99.nrValues());

  // Test Creation of a new, simpler CSP
  CSP new_csp = csp.partiallyApply(domains);
  // 81 new Domains, and still 26 all-diff constraints
  EXPECT_LONGS_EQUAL(81 + 26, new_csp.size());

  // csp.printSolution(); // still don't do it ! :-(
}

/* ************************************************************************* */
TEST(Sudoku, extreme) {
  Sudoku csp(9,                             //
             0, 0, 9, 7, 4, 8, 0, 0, 0, 7,  //
             0, 0, 0, 0, 0, 0, 0, 0, 0, 2,  //
             0, 1, 0, 9, 0, 0, 0, 0, 0, 7,  //
             0, 0, 0, 2, 4, 0, 0, 6, 4, 0,  //
             1, 0, 5, 9, 0, 0, 9, 8, 0, 0,  //
             0, 3, 0, 0, 0, 0, 0, 8, 0, 3,  //
             0, 2, 0, 0, 0, 0, 0, 0, 0, 0,  //
             0, 6, 0, 0, 0, 2, 7, 5, 9, 0, 0);

  // Do BP
  csp.runArcConsistency(9, 10);

#ifdef METIS
  VariableIndexOrdered index(csp);
  index.print("index");
  ofstream os("/Users/dellaert/src/hmetis-1.5-osx-i686/extreme-dual.txt");
  index.outputMetisFormat(os);
#endif

  // Do BP (AC1)
  auto domains = csp.runArcConsistency(9, 10);
  // csp.printDomains(domains);
  Key key99 = Symbol('9', 9);
  Domain domain99 = domains.at(key99);
  EXPECT_LONGS_EQUAL(2, domain99.nrValues());

  // Test Creation of a new, simpler CSP
  CSP new_csp = csp.partiallyApply(domains);
  // 81 new Domains, and still 20 all-diff constraints
  EXPECT_LONGS_EQUAL(81 + 20, new_csp.size());

  // csp.printSolution(); // still don't do it ! :-(
}

/* ************************************************************************* */
TEST(Sudoku, AJC_3star_Feb8_2012) {
  Sudoku csp(9,                          //
             9, 5, 0, 0, 0, 6, 0, 0, 0,  //
             0, 8, 4, 0, 7, 0, 0, 0, 0,  //
             6, 2, 0, 5, 0, 0, 4, 0, 0,  //

             0, 0, 0, 2, 9, 0, 6, 0, 0,  //
             0, 9, 0, 0, 0, 0, 0, 2, 0,  //
             0, 0, 2, 0, 6, 3, 0, 0, 0,  //

             0, 0, 9, 0, 0, 7, 0, 6, 8,  //
             0, 0, 0, 0, 3, 0, 2, 9, 0,  //
             0, 0, 0, 1, 0, 0, 0, 3, 7);

  // Do BP (AC1)
  auto domains = csp.runArcConsistency(9, 10);
  // csp.printDomains(domains);
  Key key99 = Symbol('9', 9);
  Domain domain99 = domains.at(key99);
  EXPECT_LONGS_EQUAL(1, domain99.nrValues());

  // Test Creation of a new, simpler CSP
  CSP new_csp = csp.partiallyApply(domains);
  // Just the 81 new Domains
  EXPECT_LONGS_EQUAL(81, new_csp.size());

  // Check that solution
  auto solution = new_csp.optimize();
  // csp.printAssignment(solution);
  EXPECT_LONGS_EQUAL(6, solution.at(key99));
}

/* ************************************************************************* */
int main() {
  TestResult tr;
  return TestRegistry::runAllTests(tr);
}
/* ************************************************************************* */


#include "ray_intersect_triangle_mesh_brute_force.h"
#include "Ray.h"
#include "MeshTriangle.h"
#include "Object.h"
#include "AABBTree.h"
#include "warnings.h"
#include "VectorXb.h"
#include "tictoc.h"
#include "visualize_aabbtree.h"
#include <igl/read_triangle_mesh.h>
#include <Eigen/Core>
#include <string> // std::stoi
#include <iostream>
#include <iomanip> // std::setw
#include <memory> // std::shared_ptr

int main(int argc, char * argv[])
{
  /////////////////////////////////////////////////////////////////////////////
  // RAY TRIANGLE MESH INTERSECTION
  /////////////////////////////////////////////////////////////////////////////
  std::cout<<"# Ray Triangle Mesh Intersection"<<std::endl;
  // Read in a triangle mesh
  Eigen::MatrixXd V;
  Eigen::MatrixXi F;
  igl::read_triangle_mesh(argc>1?argv[1]:"../data/rubber-ducky.obj",V,F);
  std::cout<<"  |V| "<<V.rows()<<"  "<<std::endl;
  std::cout<<"  |F| "<<F.rows()<<"  "<<std::endl<<std::endl;
  // Make a bunch of random rays
  std::vector<Ray> rays;
  rays.reserve(argc>2?std::stoi(argv[2]):1000);
  // Default bounds on ray
  double min_t = 0;
  double max_t = std::numeric_limits<double>::infinity();
  for(int r = 0;r<rays.capacity();r++)
  {
    rays.emplace_back( Eigen::Vector3d::Random(), Eigen::Vector3d::Random());
  }
  std::cout<<"  Firing "<<rays.size()<<" rays..."<<std::endl<<std::endl;

  // Brute Force
  tic(); // Start the clock!
  Eigen::VectorXb bf_hit(rays.size());
  Eigen::VectorXd bf_t(rays.size());
  Eigen::VectorXi bf_I(rays.size());
  // loop over rays
  for(int r = 0;r<rays.size();r++)
  {
    bf_hit(r) =
      ray_intersect_triangle_mesh_brute_force(
      rays[r], V, F, min_t, max_t, bf_t(r), bf_I(r));
  }
  std::cout<<"  | Method      | Time in seconds |"<<std::endl;
  std::cout<<"  |:------------|----------------:|"<<std::endl;
  std::cout<<"  | brute force | " << FLOAT15 << toc() << " |"<<std::endl;

  // Build a tree
  tic();
  std::vector<std::shared_ptr<Object> > triangles;
  triangles.reserve(F.rows());
  // Create a box for each triangle
  for(int f = 0;f<F.rows();f++)
  {
    triangles.emplace_back( std::make_shared<MeshTriangle>(V,F,f) );
  }
  std::shared_ptr<AABBTree> root = std::make_shared<AABBTree>(triangles);
  std::cout<<"  | build tree  | " << FLOAT15 << toc() << " |"<<std::endl;

  // Shoot rays at tree
  tic();
  Eigen::VectorXb tree_hit(rays.size());
  Eigen::VectorXd tree_t(rays.size());
  Eigen::VectorXi tree_I(rays.size());
  {
    // loop over rays
    for(int r = 0;r<rays.size();r++)
    {
      std::shared_ptr<Object> hit_object;
      tree_hit(r) =
        root->ray_intersect(rays[r],min_t,max_t,tree_t(r),hit_object);
      if(tree_hit(r))
      {
        const std::shared_ptr<MeshTriangle> hit_triangle =
          std::static_pointer_cast<MeshTriangle>(hit_object);
        tree_I(r) = hit_triangle->f;
      }
    }
    std::cout<<"  | use tree    | " << FLOAT15 << toc() << " |"<<std::endl;
  }

  // Check that solutions match.
  for(int r = 0;r<rays.size();r++)
  {
    WARN_IF_NOT_EQUAL(bf_hit,tree_hit,r);
    if(bf_hit(r) && tree_hit(r))
    {
      WARN_IF_NOT_APPROX(bf_t,tree_t,r);
      WARN_IF_NOT_EQUAL(bf_I,tree_I,r);
    }
  }

  // Visualize the tree
  visualize_aabbtree(V,F,root);
}


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

#ifndef BOOST_MATH_SPECIAL_LAGUERRE_HPP
#define BOOST_MATH_SPECIAL_LAGUERRE_HPP

#ifdef _MSC_VER
#pragma once
#endif

#include <boost/math/special_functions/math_fwd.hpp>
#include <boost/math/tools/config.hpp>
#include <boost/math/policies/error_handling.hpp>

namespace boost{
namespace math{

// Recurrance relation for Laguerre polynomials:
template <class T1, class T2, class T3>
inline typename tools::promote_args<T1, T2, T3>::type
   laguerre_next(unsigned n, T1 x, T2 Ln, T3 Lnm1)
{
   typedef typename tools::promote_args<T1, T2, T3>::type result_type;
   return ((2 * n + 1 - result_type(x)) * result_type(Ln) - n * result_type(Lnm1)) / (n + 1);
}

namespace detail{

// Implement Laguerre polynomials via recurrance:
template <class T>
T laguerre_imp(unsigned n, T x)
{
   T p0 = 1;
   T p1 = 1 - x;

   if(n == 0)
      return p0;

   unsigned c = 1;

   while(c < n)
   {
      std::swap(p0, p1);
      p1 = laguerre_next(c, x, p0, p1);
      ++c;
   }
   return p1;
}

template <class T, class Policy>
inline typename tools::promote_args<T>::type
laguerre(unsigned n, T x, const Policy&, const mpl::true_&)
{
   typedef typename tools::promote_args<T>::type result_type;
   typedef typename policies::evaluation<result_type, Policy>::type value_type;
   return policies::checked_narrowing_cast<result_type, Policy>(detail::laguerre_imp(n, static_cast<value_type>(x)), "boost::math::laguerre<%1%>(unsigned, %1%)");
}

template <class T>
inline typename tools::promote_args<T>::type
   laguerre(unsigned n, unsigned m, T x, const mpl::false_&)
{
   return boost::math::laguerre(n, m, x, policies::policy<>());
}

} // namespace detail

template <class T>
inline typename tools::promote_args<T>::type
   laguerre(unsigned n, T x)
{
   return laguerre(n, x, policies::policy<>());
}

// Recurrence for associated polynomials:
template <class T1, class T2, class T3>
inline typename tools::promote_args<T1, T2, T3>::type
   laguerre_next(unsigned n, unsigned l, T1 x, T2 Pl, T3 Plm1)
{
   typedef typename tools::promote_args<T1, T2, T3>::type result_type;
   return ((2 * n + l + 1 - result_type(x)) * result_type(Pl) - (n + l) * result_type(Plm1)) / (n+1);
}

namespace detail{
// Laguerre Associated Polynomial:
template <class T, class Policy>
T laguerre_imp(unsigned n, unsigned m, T x, const Policy& pol)
{
   // Special cases:
   if(m == 0)
      return boost::math::laguerre(n, x, pol);

   T p0 = 1;

   if(n == 0)
      return p0;

   T p1 = m + 1 - x;

   unsigned c = 1;

   while(c < n)
   {
      std::swap(p0, p1);
      p1 = laguerre_next(c, m, x, p0, p1);
      ++c;
   }
   return p1;
}

}

template <class T, class Policy>
inline typename tools::promote_args<T>::type
   laguerre(unsigned n, unsigned m, T x, const Policy& pol)
{
   typedef typename tools::promote_args<T>::type result_type;
   typedef typename policies::evaluation<result_type, Policy>::type value_type;
   return policies::checked_narrowing_cast<result_type, Policy>(detail::laguerre_imp(n, m, static_cast<value_type>(x), pol), "boost::math::laguerre<%1%>(unsigned, unsigned, %1%)");
}

template <class T1, class T2>
inline typename laguerre_result<T1, T2>::type
   laguerre(unsigned n, T1 m, T2 x)
{
   typedef typename policies::is_policy<T2>::type tag_type;
   return detail::laguerre(n, m, x, tag_type());
}

} // namespace math
} // namespace boost

#endif // BOOST_MATH_SPECIAL_LAGUERRE_HPP


#include <anie/activation_layers/sigmoid_layer.hpp>

#include <cassert>
#include <memory>
#include <boost/compute.hpp>

namespace anie
{
	sigmoid_layer::sigmoid_layer(const anie::device& device) noexcept
		: activation_layer(device), u_(device)
	{}

	std::string_view sigmoid_layer::name() const noexcept
	{
		return "Sigmoid";
	}

	matrix sigmoid_layer::forward(const matrix& z) const
	{
		assert(device() == z.device());

		static BOOST_COMPUTE_FUNCTION(arithemtic_type, sigmoid, (const double x),
		{
			return 1. / (1. + exp(-x));
		});

		matrix result(z);
		boost::compute::transform(result.begin(), result.end(), result.begin(), sigmoid, device().queue());

		u_ = z;
		return result;
	}
	matrix sigmoid_layer::backward(const matrix& d)
	{
		assert(device() == d.device());

		static BOOST_COMPUTE_FUNCTION(arithemtic_type, dsigmoid, (const double x),
		{
			const double y = 1. / (1. + exp(-x));
			return y * (1. - y);
		});

		matrix result(u_);
		boost::compute::transform(result.begin(), result.end(), result.begin(), dsigmoid, device().queue());

		return result;
	}
}

namespace anie
{
	layer_ptr sigmoid_layer_generator::operator()(const anie::device& device) const
	{
		return std::make_shared<sigmoid_layer>(device);
	}

	layer_generator_ptr sigmoid()
	{
		return std::make_shared<sigmoid_layer_generator>();
	}
	layer_ptr sigmoid(const anie::device& device)
	{
		return std::make_shared<sigmoid_layer>(device);
	}
}


//
//  Copyright Toon Knapen, Karl Meerbergen
//
// Distributed under the Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//

#include "../../blas/test/random.hpp"

#include <boost/numeric/bindings/lapack/geqrf.hpp>
#include <boost/numeric/bindings/lapack/ormqr.hpp>
#include <boost/numeric/bindings/lapack/orgqr.hpp>
#include <boost/numeric/bindings/traits/ublas_matrix.hpp>
#include <boost/numeric/bindings/traits/ublas_vector.hpp>
#include <boost/numeric/ublas/triangular.hpp>
#include <boost/numeric/ublas/io.hpp>

#include <iostream>
#include <limits>

namespace ublas = boost::numeric::ublas;
namespace lapack = boost::numeric::bindings::lapack;

// Randomize a matrix
template <typename M>
void randomize(M& m) {
   typedef typename M::size_type  size_type ;
   typedef typename M::value_type value_type ;

   size_type size1 = m.size1() ;
   size_type size2 = m.size2() ;

   for (size_type i=0; i<size2; ++i) {
      for (size_type j=0; j<size1; ++j) {
         m(j,i) = random_value< value_type >() ;
      }
   }
} // randomize()

template <typename T>
struct transpose {
   static const char value ;
};

template <typename T>
const char transpose<T>::value = 'T';

template <typename T>
struct transpose< std::complex<T> > {
   static const char value ;
};

template <typename T>
const char transpose< std::complex<T> >::value = 'C';

template <typename M>
ublas::triangular_adaptor<const M, ublas::upper> upper_part(const M& m) {
   return ublas::triangular_adaptor<const M, ublas::upper>( m );
}

template <typename T, typename W>
int do_memory_type(int n, W workspace) {
   typedef typename boost::numeric::bindings::traits::type_traits<T>::real_type real_type ;
   typedef std::complex< real_type >                                            complex_type ;

   typedef ublas::matrix<T, ublas::column_major> matrix_type ;
   typedef ublas::vector<T>                      vector_type ;

   // Set matrix
   matrix_type a( n, n );
   vector_type tau( n );

   randomize( a );
   matrix_type a2( a );
   matrix_type a3( a );

   // Compute QR factorization.
   lapack::geqrf( a, tau, workspace ) ;

   // Apply the orthogonal transformations to a2
   lapack::ormqr( 'L', transpose<T>::value, a, tau, a2, workspace );

   // The upper triangular parts of a and a2 must be equal.
   if (norm_frobenius( upper_part( a - a2 ) )
            > std::numeric_limits<real_type>::epsilon() * 10.0 * norm_frobenius( upper_part( a ) ) ) return 255 ;

   // Generate orthogonal matrix
   lapack::orgqr( a, tau, workspace );

   // The result of lapack::ormqr and the equivalent matrix product must be equal.
   if (norm_frobenius( a2 - prod(herm(a), a3) )
            > std::numeric_limits<real_type>::epsilon() * 10.0 * norm_frobenius( a2 ) ) return 255 ;

   return 0 ;
} // do_value_type()

template <typename T>
int do_value_type() {
   const int n = 8 ;

   if (do_memory_type<T,lapack::optimal_workspace>( n, lapack::optimal_workspace() ) ) return 255 ;
   if (do_memory_type<T,lapack::minimal_workspace>( n, lapack::minimal_workspace() ) ) return 255 ;

   ublas::vector<T> work( n );
   do_memory_type<T, lapack::detail::workspace1<ublas::vector<T> > >( n, lapack::workspace(work) );
   return 0;
} // do_value_type()

int main() {
   // Run tests for different value_types
   if (do_value_type<float>()) return 255;
   if (do_value_type<double>()) return 255;
   if (do_value_type< std::complex<float> >()) return 255;
   if (do_value_type< std::complex<double> >()) return 255;

   std::cout << "Regression test succeeded\n" ;
   return 0;
}


//Link to Boost
#define BOOST_TEST_DYN_LINK

//VERY IMPORTANT - include this last
//#include <boost/test/included/unit_test.hpp>
#include <boost/test/unit_test.hpp>

#include "test.h"
#include "../BetheFunctions.h"
#include "../BethePolynomials.h"

using namespace std::literals::complex_literals;

// test suite
BOOST_FIXTURE_TEST_SUITE(BetheFactor_suite, SimpleTestFixture, * utf::label("BetheFactor"))

BOOST_AUTO_TEST_CASE(coef_test)
{
    int L = 6;
    int M = 2;
    BetheFactorLeft left(L, M);
    BetheFactorRight right(L, M);
    for (int i = 0; i < 3; i++) {
        var_t *coef1Left = left.coef1((RootType)i);
        var_t *coef1Right = right.coef1((RootType)i);
        for (int j = 0; j < 3; j++) {
            BOOST_TEST_INFO("i=" << i << ", j=" << j);
            BOOST_TEST(isZero(std::conj(coef1Left[j]) - coef1Right[j]));
        }
    }

    for (int i = 0; i < 3; i++) {
        for (int j = 0; j < 3; j++) {
            var_t *coef2Left = left.coef2((RootType)i, (RootType)j);
            var_t *coef2Right = right.coef2((RootType)i, (RootType)j);
            for (int k = 0; k < 6; k++) {
                BOOST_TEST_INFO("i=" << i << ", j=" << j << ", k=" << k);
                BOOST_TEST(isZero(std::conj(coef2Left[k]) - coef2Right[k]));
            }
        }
    }
}

BOOST_DATA_TEST_CASE(eval_test_regular, bdata::random(4, 16) ^ bdata::xrange(20), L, index)
{
    int M = random64(L/2);
    if (M <= 1) M = 2;
    Vector roots(M);
    for (int i = 0; i < M; i++) {
        roots(i) = randomComplex(0.1, 0.9);
    }

    BetheSolution sol1(roots, L, 0.0, 0.0);
    InvertableSolution sol2(roots, L);

    BetheFactorLeft left1(L, M);
    BetheFactorRight right1(L, M);

    //BethePolynomial left2(L, M, true);
    BethePolynomialLeft left2(L, M);
    //BethePolynomial right2(L, M, false);
    BethePolynomialRight right2(L, M);

    BOOST_TEST((left1.eval(sol1) - left2.eval(sol2)).norm() == 0, tt::tolerance(eps));
    Vector v1 = right1.eval(sol1);
    Vector v2 = right2.eval(sol2);
    BOOST_TEST((v1 - v2).norm() == 0, tt::tolerance(eps));
}

BOOST_DATA_TEST_CASE(eval_test_largest, bdata::random(4, 16) ^ bdata::xrange(20), L, index)
{
    int M = random64(L/2);
    if (M <= 1) M = 2;
    int indexOfLargest = random64(M - 1);
    Vector roots(M);
    roots(indexOfLargest) = randomComplex(blowup_root_cutoff/0.3, blowup_root_cutoff/0.2);

    for (int i = 0; i < M; i++) {
        if (i == indexOfLargest) continue;
        roots(i) = randomComplex(0.1, 0.9);
    }

    BetheSolution sol1(roots, L, 0.0, 0.01);
    InvertableSolution sol2(roots, L);

    BetheFactorLeft left1(L, M);
    BetheFactorRight right1(L, M);

    BethePolynomialLeft left2(L, M);
    BethePolynomialRight right2(L, M);

    BOOST_TEST((left1.eval(sol1) - left2.eval(sol2)).norm() == 0, tt::tolerance(eps));
    Vector v1 = right1.eval(sol1);
    Vector v2 = right2.eval(sol2);
//    std::cout << "indexOfLargest=" << indexOfLargest << std::endl;
//    std::cout << "v1=" << v1 << std::endl;
//    std::cout << "v2=" << v2 << std::endl;
    BOOST_TEST((v1 - v2).norm() == 0, tt::tolerance(eps));
}

BOOST_DATA_TEST_CASE(eval_test_largest_large, bdata::random(4, 16) ^ bdata::xrange(1), L, index)
{
    int M = 2;
    int indexOfLargest = 0;
    int indexOfLarge = 1;

    Vector roots(M);
    roots(indexOfLargest) = randomComplex(blowup_root_cutoff/0.3, blowup_root_cutoff/0.2);
    roots(indexOfLarge) = randomComplex(blowup_root_cutoff/0.5, blowup_root_cutoff/0.4);

    BetheSolution sol1(roots, L, 0.0, 0.01);
    InvertableSolution sol2(roots, L);

    BetheFactorLeft left1(L, M);
    BetheFactorRight right1(L, M);

    BethePolynomialLeft left2(L, M);
    BethePolynomialRight right2(L, M);

    Vector v1 = left1.eval(sol1);
    Vector v2 = left2.eval(sol2);

    var_t ratio = roots(indexOfLarge) / roots(indexOfLargest);
    v2(indexOfLargest) *= roots(indexOfLarge);
    v2(indexOfLarge) *= roots(indexOfLarge) * std::pow(ratio, L);

//    std::cout << "v1=" << v1 << std::endl;
//    std::cout << "v2=" << v2 << std::endl;

    BOOST_TEST((v1 - v2).norm() == 0, tt::tolerance(eps));
    v1 = right1.eval(sol1);
    v2 = right2.eval(sol2);
    v2(indexOfLargest) *= roots(indexOfLarge);
    v2(indexOfLarge) *= roots(indexOfLarge) * std::pow(ratio, L);
//    std::cout << "indexOfLargest=" << indexOfLargest << std::endl;
//    std::cout << "v1=" << v1 << std::endl;
//    std::cout << "v2=" << v2 << std::endl;
    BOOST_TEST((v1 - v2).norm() == 0, tt::tolerance(eps));
}

BOOST_DATA_TEST_CASE(eval_test_singular, bdata::random(2, 16) ^ bdata::xrange(20), halfL, index)
{
    int L = 2 * halfL;
    int M = 2;
    int index1 = 0;
    int index2 = 1;

    Vector roots(M);
    roots(index2) = -0.5il;
    roots(index1) = 0.5il;
    BetheSolution sol(roots, L, 0.0, 0.01);

    BetheFactorLeft left(L, M);
    BetheFactorRight right(L, M);

//    std::cout << "left.eval() = " << left.eval(sol) << std::endl;
//    std::cout << "right.eval() = " << right.eval(sol) << std::endl;

    BOOST_TEST((left.eval(sol) - right.eval(sol)).norm() == 0, tt::tolerance(eps));
}

BOOST_DATA_TEST_CASE(diff_test_regular, bdata::random(4, 16) ^ bdata::xrange(20), L, index)
{
    int M = random64(L/2);
    if (M <= 1) M = 2;
    Vector roots(M);
    for (int i = 0; i < M; i++) {
        roots(i) = randomComplex(0.1, 0.9);
    }

    BetheSolution sol1(roots, L, 0.0, 0.01);
    InvertableSolution sol2(roots, L);

    BetheFactorLeft left1(L, M);
    BetheFactorRight right1(L, M);

    BethePolynomialLeft left2(L, M);
    BethePolynomialRight right2(L, M);

    BOOST_TEST((left1.diff(sol1) - left2.diff(sol2)).norm() == 0, tt::tolerance(eps));
    Matrix v1 = right1.diff(sol1);
    Matrix v2 = right2.diff(sol2);
//    std::cout << "v1-v2=" << v1 - v2 << std::endl;
    BOOST_TEST((v1 - v2).norm() == 0, tt::tolerance(eps));
}

BOOST_DATA_TEST_CASE(diff_test_largest, bdata::random(4, 16) ^ bdata::xrange(20), L, index)
{
    int M = random64(L/2);
    if (M <= 1) M = 2;
    int indexOfLargest = random64(M - 1);
    Vector roots(M);
    roots(indexOfLargest) = randomComplex(blowup_root_cutoff/0.3, blowup_root_cutoff/0.2);

    for (int i = 0; i < M; i++) {
        if (i == indexOfLargest) continue;
        roots(i) = randomComplex(0.1, 0.9);
    }

    BetheSolution sol1(roots, L, 0.0, 0.01);
    InvertableSolution sol2(roots, L);

    BetheFactorLeft left1(L, M);
    BetheFactorRight right1(L, M);

    BethePolynomialLeft left2(L, M);
    BethePolynomialRight right2(L, M);

    Matrix w1 = left1.diff(sol1);
    Matrix w2 = left2.diff(sol2);
//    std::cout << "M=" << M << ", indexOfLargest=" << indexOfLargest << std::endl;
//    std::cout << "w1=" << w1 << std::endl;
//    std::cout << "w2=" << w2 << std::endl;
//    std::cout << "diff=" << w1-w2 << std::endl;

    BOOST_TEST((w1 - w2).norm() == 0, tt::tolerance(eps));
    Matrix v1 = right1.diff(sol1);
    Matrix v2 = right2.diff(sol2);
//    std::cout << "M=" << M << ", indexOfLargest=" << indexOfLargest << std::endl;
//    std::cout << "v1=" << v1 << std::endl;
//    std::cout << "v2=" << v2 << std::endl;
//    std::cout << "diff=" << v1-v2 << std::endl;
    BOOST_TEST((v1 - v2).norm() == 0, tt::tolerance(eps));
}

BOOST_DATA_TEST_CASE(diff_test_LeftFactor, bdata::random(4, 16) ^ bdata::xrange(20), L, index)
{
    int M = random64(L/2);
    if (M <= 1) M = 2;
    int indexOfLargest = random64(M - 1);
    Vector roots(M);
    Vector diff1(M); // infinetesimal differential of roots for regular roots
    Vector diff2(M); // infinetesimal differential of roots for BetheSolution.root();
    elem_t delta = 1e-4;

    roots(indexOfLargest) = randomComplex(blowup_root_cutoff/0.3, blowup_root_cutoff/0.2);
    diff1(indexOfLargest) = delta * (random64(3) + 1) * randomComplex(1.0, 1.1);
    diff2(indexOfLargest) = -diff1(indexOfLargest) / (roots(indexOfLargest) * roots(indexOfLargest));

    for (int i = 0; i < M; i++) {
        if (i == indexOfLargest) continue;
        diff1(i) = delta * (random64(3) + 1) * randomComplex(1.0, 1.1);
        if (i % 2 == 0) {
            // large root
            roots(i) = randomComplex(blowup_root_cutoff/0.5, blowup_root_cutoff/0.4);
            diff2(i) = diff1(i) / roots(indexOfLargest) - diff1(indexOfLargest) * roots(i) / (roots(indexOfLargest) * roots(indexOfLargest));
        } else {
            // regular root
            roots(i) = randomComplex(blowup_root_cutoff/1.5, blowup_root_cutoff/1.2);
            diff2(i) = diff1(i);
        }
    }

    Vector roots2 = roots + diff1;
    Vector roots3 = roots + diff1 * 0.5L;

    BetheSolution sol1(roots, L, 0.0, 0.01);
    BetheSolution sol2(roots2, L, 0.0, 0.01);
    BetheSolution sol3(roots3, L, 0.0, 0.01);

    BetheFactorLeft left(L, M);
    Vector res1 = left.eval(sol1);
    Vector res2 = left.eval(sol2);
    Vector resDiff = res2 - res1;

    Matrix jacobian = left.diff(sol3);
    Vector resDiff2 = jacobian * diff2;

//    std::cout << "resDiff=" << resDiff << std::endl;
//    std::cout << "resDiff2=" << resDiff2 << std::endl;
    for (int i = 0; i < M; i++) {
        var_t v1 = resDiff(i);
        var_t v2 = resDiff2(i);
        BOOST_TEST_INFO("M=" << M << ",i=" << i << ", v1=" << v1 << ", v2=" << v2);
        BOOST_TEST(std::abs(v1-v2)/std::max(std::abs(v1), std::abs(v2)) == 0, tt::tolerance(delta * 10));
    }
}

BOOST_DATA_TEST_CASE(diff_test_RightFactor, bdata::random(4, 16) ^ bdata::xrange(20), L, index)
{
    int M = random64(L/2);
    if (M <= 1) M = 2;
    int indexOfLargest = random64(M - 1);
    Vector roots(M);
    Vector diff1(M); // infinetesimal differential of roots for regular roots
    Vector diff2(M); // infinetesimal differential of roots for BetheSolution.root();
    elem_t delta = 1e-4;

    roots(indexOfLargest) = randomComplex(blowup_root_cutoff/0.3, blowup_root_cutoff/0.2);
    diff1(indexOfLargest) = delta * (random64(3) + 1) * randomComplex(1.0, 1.1);
    diff2(indexOfLargest) = -diff1(indexOfLargest) / (roots(indexOfLargest) * roots(indexOfLargest));

    for (int i = 0; i < M; i++) {
        if (i == indexOfLargest) continue;
        diff1(i) = delta * (random64(3) + 1) * randomComplex(1.0, 1.1);
        if (i % 2 == 0) {
            // large root
            roots(i) = randomComplex(blowup_root_cutoff/0.5, blowup_root_cutoff/0.4);
            diff2(i) = diff1(i) / roots(indexOfLargest) - diff1(indexOfLargest) * roots(i) / (roots(indexOfLargest) * roots(indexOfLargest));
        } else {
            // regular root
            roots(i) = randomComplex(blowup_root_cutoff/1.5, blowup_root_cutoff/1.2);
            diff2(i) = diff1(i);
        }
    }

    Vector roots2 = roots + diff1;
    Vector roots3 = roots + diff1 * 0.5L;

    BetheSolution sol1(roots, L, 0.0, 0.01);
    BetheSolution sol2(roots2, L, 0.0, 0.01);
    BetheSolution sol3(roots3, L, 0.0, 0.01);

    BetheFactorRight right(L, M);
    Vector res1 = right.eval(sol1);
    Vector res2 = right.eval(sol2);
    Vector resDiff = res2 - res1;

    Matrix jacobian = right.diff(sol3);
    Vector resDiff2 = jacobian * diff2;

//    std::cout << "resDiff=" << resDiff << std::endl;
//    std::cout << "resDiff2=" << resDiff2 << std::endl;
    for (int i = 0; i < M; i++) {
        var_t v1 = resDiff(i);
        var_t v2 = resDiff2(i);
        BOOST_TEST_INFO("M=" << M << ",i=" << i << ", v1=" << v1 << ", v2=" << v2);
        BOOST_TEST(std::abs(v1-v2)/std::max(std::abs(v1), std::abs(v2)) == 0, tt::tolerance(delta * 20));
    }
}

BOOST_AUTO_TEST_CASE(diff_test_singular)
{
    elem_t tmp = eps_may_be_singular;
    eps_may_be_singular = 0.1L;
    int L = 6;
    int M = 3;
    var_t r1(0.54236679461751705L,2.90957067279989408e-11L);
    var_t r2(0.0968423125274162171L,0.500001874059868433L);
    var_t r3(0.0968423125273995912L,-0.50000187418704665L);
    elem_t beta = 0.914203L;

    Vector roots(M);
    roots(0) = r1;
    roots(1) = r2;
    roots(2) = r3;
    BetheSolution sol(roots, L, beta, Pi/1000);
    BOOST_TEST_REQUIRE(sol.indexOfSingularRoot1() >= 1);
    BOOST_TEST_REQUIRE(sol.indexOfSingularRoot2() >= 1);

    BetheFactorLeft left(L, M);
    Vector res1 = left.eval(sol);
    Matrix jacobian = left.diff(sol);

    Vector inc(M);
    inc(0) = 0.0L;
    inc(sol.indexOfSingularRoot1()) = 0.0L;
    inc(sol.indexOfSingularRoot2()) = 0.0L;
    inc(sol.indexOfSingularRoot2()) = 0.0L;
    inc(0) = randomComplex(1e-3, 2e-3);
    inc(sol.indexOfSingularRoot1()) = randomComplex(1e-1, 2e-1);
    inc(sol.indexOfSingularRoot2()) = randomComplex(1e-3, 2e-3);

    sol.root() += inc;
    Vector res2 = left.eval(sol);
    Vector diff = res2 - res1;

    Vector expected = jacobian * inc;
//    std::cout << "res1=" << res1 << std::endl;
//    std::cout << "res2=" << res2 << std::endl;
//    std::cout << "diff=" << diff << std::endl;
//    std::cout << "expected=" << expected << std::endl;

    BOOST_TEST_INFO("inc=" << inc);
    BOOST_TEST((diff-expected).norm()/expected.norm() < 2e-2);

    eps_may_be_singular = tmp;
}

BOOST_AUTO_TEST_CASE(diff_test_singular_RightFactor)
{
    elem_t tmp = eps_may_be_singular;
    eps_may_be_singular = 0.1L;
    int L = 6;
    int M = 3;
    var_t r1(0.54236679461751705L,2.90957067279989408e-11L);
    var_t r2(0.0968423125274162171L,0.500001874059868433L);
    var_t r3(0.0968423125273995912L,-0.50000187418704665L);
    elem_t beta = 0.914203L;

    Vector roots(M);
    roots(0) = r1;
    roots(1) = r2;
    roots(2) = r3;
    BetheSolution sol(roots, L, beta, Pi/1000);
    BOOST_TEST_REQUIRE(sol.indexOfSingularRoot1() >= 1);
    BOOST_TEST_REQUIRE(sol.indexOfSingularRoot2() >= 1);

    BetheFactorRight right(L, M);
    Vector res1 = right.eval(sol);
    Matrix jacobian = right.diff(sol);

    Vector inc(M);
    inc(0) = 0.0L;
    inc(sol.indexOfSingularRoot1()) = 0.0L;
    inc(sol.indexOfSingularRoot2()) = 0.0L;
    inc(sol.indexOfSingularRoot2()) = 0.0L;
    inc(0) = randomComplex(1e-3, 2e-3);
    inc(sol.indexOfSingularRoot1()) = randomComplex(1e-1, 2e-1);
    inc(sol.indexOfSingularRoot2()) = randomComplex(1e-3, 2e-3);

    sol.root() += inc;
    Vector res2 = right.eval(sol);
    Vector diff = res2 - res1;

    Vector expected = jacobian * inc;
//    std::cout << "res1=" << res1 << std::endl;
//    std::cout << "res2=" << res2 << std::endl;
//    std::cout << "diff=" << diff << std::endl;
//    std::cout << "expected=" << expected << std::endl;

    BOOST_TEST_INFO("inc=" << inc);
    BOOST_TEST((diff-expected).norm()/expected.norm() < 2e-2);

    eps_may_be_singular = tmp;
}

BOOST_AUTO_TEST_CASE(diff_test_largest_LeftFactor)
{
    int L = 6;
    int M = 3;

    var_t r1(0.288675134570472836L, -1.98390501120640128e-11L);
    var_t r2(4833488737.1251517L, 8727942389.24525939L);
    var_t r3(-5137763027.0882383L, 8552377948.96713318L);
    elem_t beta = 1.5708L;

    Vector roots(M);
    roots(0) = r1;
    roots(1) = r2;
    roots(2) = r3;
    BetheSolution sol(roots, L, beta, Pi/1000);
    BOOST_TEST_REQUIRE(sol.indexOfBlowupRoot() >= 1);

    BetheFactorRight left(L, M);
    Vector res1 = left.eval(sol);
    Matrix jacobian = left.diff(sol);

    Vector inc = Vector::Zero(M);
    for (int i = 0; i < inc.size(); i++) {
        if (sol.type(i) == regular) {
            inc(i) = randomComplex(1e-4, 2e-4);
        } else if (sol.type(i) == largest) {
            elem_t phi = std::arg(sol.root(i));
            inc(i) = -std::abs(sol.root(i)) * random(1.0, 1.5) * std::polar(1.0L, phi);;
        } else {
            inc(i) = randomComplex(1e-4, 2e-4);
        }
    }

    sol.root() += inc;
    Vector res2 = left.eval(sol);
    Vector diff = res2 - res1;

    Vector expected = jacobian * inc;
//    std::cout << "res1=" << res1 << std::endl;
//    std::cout << "res2=" << res2 << std::endl;
//    std::cout << "diff=" << diff << std::endl;
//    std::cout << "expected=" << expected << std::endl;

    BOOST_TEST_INFO("inc=" << inc);
    BOOST_TEST((diff-expected).norm()/expected.norm() < 1e-3);
}

BOOST_AUTO_TEST_CASE(diff_test_largest_RightFactor)
{
    int L = 6;
    int M = 3;

    var_t r1(0.288675134570472836L, -1.98390501120640128e-11L);
    var_t r2(4833488737.1251517L, 8727942389.24525939L);
    var_t r3(-5137763027.0882383L, 8552377948.96713318L);
    elem_t beta = 1.5708L;

    Vector roots(M);
    roots(0) = r1;
    roots(1) = r2;
    roots(2) = r3;
    BetheSolution sol(roots, L, beta, Pi/1000);
    BOOST_TEST_REQUIRE(sol.indexOfBlowupRoot() >= 1);

    BetheFactorRight right(L, M);
    Vector res1 = right.eval(sol);
    Matrix jacobian = right.diff(sol);

    Vector inc = Vector::Zero(M);
    for (int i = 0; i < inc.size(); i++) {
        if (sol.type(i) == regular) {
            inc(i) = randomComplex(1e-4, 2e-4);
        } else if (sol.type(i) == largest) {
            elem_t phi = std::arg(sol.root(i));
            inc(i) = std::abs(sol.root(i)) * random(1.0, 1.5) * std::polar(1.0L, phi);;
        } else {
            inc(i) = randomComplex(1e-4, 2e-4);
        }
    }

    sol.root() += inc;
    Vector res2 = right.eval(sol);
    Vector diff = res2 - res1;

    Vector expected = jacobian * inc;
//    std::cout << "res1=" << res1 << std::endl;
//    std::cout << "res2=" << res2 << std::endl;
//    std::cout << "diff=" << diff << std::endl;
//    std::cout << "expected=" << expected << std::endl;

    BOOST_TEST_INFO("inc=" << inc);
    BOOST_TEST((diff-expected).norm()/expected.norm() < 1e-3);
}
/*
BOOST_AUTO_TEST_CASE(diff_test_largest2)
{
    int L = 6;
    int M = 2;
    var_t r1(-2002453610.04371309,18872186.6385834396);
    var_t r2(-2002453610.04371309,18872186.6385834396);
    elem_t beta = 0.785398;
    elem_t delta = Pi/1000;

    Vector roots(M);
    roots(0) = r1;
    roots(1) = r2;
    BetheSolution sol(roots, L, beta, delta);
    BOOST_TEST_REQUIRE(sol.indexOfBlowupRoot() >= 0);

    BetheFactorLeft left(L, M);
    BetheFactorRight right(L, M);
    Vector res1 = right.eval(sol);
    Matrix jacobian1 = left.diff(sol);
    Matrix jacobian2 = right.diff(sol);
    Matrix jacobian = jacobian1 - jacobian2 * std::polar(1.0L, 2.0L * L * beta);
    Matrix jacobianInv = jacobian.inverse();

    Vector rhs = left.eval(sol) * var_t((elem_t)0.0, 2.0 * L);

//    std::cout << "jacobian = " << jacobian << std::endl;
//    std::cout << "jacobianInv=" << jacobianInv << std::endl;
//    std::cout << "rhs=" << std::setprecision(18) << rhs << std::endl;

    Vector inc1 = jacobianInv * rhs * delta/2.0;
//    std::cout << "jacobianInv.rhs=" << std::setprecision(18) << inc1 << std::endl;
//    std::cout << "root=" << sol.root() << std::endl;
    sol.root() += inc1;
//    std::cout << "after root=" << sol.normalRoot() << std::endl;
}*/

BOOST_AUTO_TEST_SUITE_END()


// cfiles, an analysis frontend for the Chemfiles library
// Copyright (C) Guillaume Fraux and contributors -- BSD license

#include <fstream>

#include <docopt/docopt.h>
#include <fmt/format.h>
#include <fmt/ostream.h>
#include <chemfiles.hpp>
#include <Eigen/Dense>

#include "Elastic.hpp"
#include "Errors.hpp"

using namespace chemfiles;

using Matrix6 = Eigen::Matrix<double, 6, 6>;

static size_t CARTESIAN_TO_VOIGT[][2] = {
    {0, 0}, {1, 1}, {2, 2}, {1, 2}, {0, 2}, {0, 1},
};

static const std::string OPTIONS =
R"(Compute the elastic tensor of a system from the unit cell fluctuations during
a NPT simulation.

The values given here are highly dependent on having good statistic during the
simulation: both in term of having a long enough simulation time to get to the
equilibrium, and using a good barostat that does produce isobaric-isothermal
ensemble fluctuations (not just average). The theory behind this code is
described in https://dx.doi.org/10.1080/08927022.2017.1313418.

Usage:
  cfiles elastic [options] <trajectory>
  cfiles elastic (-h | --help)

Examples:
  cfiles elastic -t 328 -o elastic.dat trajectory.pdb

Options:
  -h --help                        show this help
  --format=<format>                force the input file format to be <format>
  -t <temp>, --temperature=<temp>  temperature of the simulation, in kelvin
  -o <file>, --output=<file>       write result to <file>. This default to the
                                   trajectory file name with the `.angles.dat`
                                   extension.
  --steps=<steps>                  steps to use from the input. <steps> format
                                   is <start>:<end>[:<stride>] with <start>,
                                   <end> and <stride> optional. The used steps
                                   goes from <start> to <end> (excluded) by
                                   steps of <stride>. The default values are 0
                                   for <start>, the number of steps for <end>
                                   and 1 for <stride>.
)";

static Elastic::Options parse_options(int argc, const char* argv[]) {
    auto options_str = command_header("elastic", Elastic().description());
    options_str += "Guillaume Fraux <guillaume@fraux.fr>\n\n";
    options_str += std::string(OPTIONS);
    auto args = docopt::docopt(options_str, {argv, argv + argc}, true, "");

    Elastic::Options options;
    options.trajectory = args.at("<trajectory>").asString();

    if (args["--temperature"]) {
        options.temperature = string2double(args.at("--temperature").asString());
    } else {
        throw CFilesError("missing --temperature argument");
    }

    if (args.at("--steps")) {
        options.steps = steps_range::parse(args.at("--steps").asString());
    }

    if (args["--output"]) {
        options.outfile = args.at("--output").asString();
    } else {
        options.outfile = options.trajectory + ".elastic.dat";
    }

    return options;
}

std::string Elastic::description() const {
    return "compute elastic constants from unit cell fluctuations in NPT";
}

int Elastic::run(int argc, const char* argv[]) {
    auto options = parse_options(argc, argv);
    auto cells = std::vector<Matrix3D>();

    auto trajectory = Trajectory(options.trajectory, 'r', options.format);
    for (auto step: options.steps) {
        if (step >= trajectory.nsteps()) {
            break;
        }
        auto frame = trajectory.read_step(step);
        cells.emplace_back(frame.cell().matrix());
    }

    // Use the avrerage as the reference state
    auto reference = Matrix3D::zero();
    for (auto& cell: cells) {
        reference += cell.invert();
    }
    reference /= cells.size();
    auto reference_t = reference.transpose();

    auto epsilons = std::vector<Matrix3D>();
    epsilons.reserve(cells.size());
    for (auto& cell: cells) {
        epsilons.emplace_back(0.5 * (reference_t * cell.transpose() * cell * reference - Matrix3D::unit()));
    }

    // in GPa A^2 / K
    auto BOLTZMANN = 1.38065e-2;
    auto volume_inv = reference.determinant();
    auto v_kt = 1.0 / (volume_inv * BOLTZMANN * options.temperature);
    auto compute = [&](size_t ij[2], size_t kl[2]) {
        auto i = ij[0];
        auto j = ij[1];
        auto k = kl[0];
        auto l = kl[1];

        // Multiplicative factors for cross terms yz xz xy
        double factor = 1.0;
        if (i != j) {
            factor *= 2;
        }
        if (k != l) {
            factor *= 2;
        }

        double eij = 0;
        double ekl = 0;
        double eij_ekl = 0;
        for (auto& epsilon: epsilons) {
            eij += epsilon[i][j];
            ekl += epsilon[k][l];
            eij_ekl += epsilon[i][j] * epsilon[k][l];
        }
        eij /= epsilons.size();
        ekl /= epsilons.size();
        eij_ekl /= epsilons.size();

        return factor * v_kt * (eij_ekl - eij * ekl);
    };

    auto SVoigt = Matrix6();
    for (size_t i=0; i<6; i++) {
        for (size_t j=0; j<=i; j++) {
            SVoigt(i, j) = compute(CARTESIAN_TO_VOIGT[i], CARTESIAN_TO_VOIGT[j]);
        }
    }
    // Make the matrix symetric
    for (size_t i=0; i<6; i++) {
        for (size_t j=i+1; j<6; j++) {
            SVoigt(i, j) = SVoigt(j, i);
        }
    }

    if (std::abs(SVoigt.determinant()) < 100 * DBL_EPSILON) {
        throw CFilesError("the compliance matrix is not invertible");
    }

    auto CVoigt = SVoigt.inverse();
    std::ofstream outfile(options.outfile, std::ios::out);
    if (!outfile.is_open()) {
        throw CFilesError("Could not open the '" + options.outfile + "' file.");
    }

    fmt::print(outfile, "# stiffness tensor in GPa from {}\n", options.trajectory);
    for (size_t i=0; i<6; i++) {
        for (size_t j=0; j<6; j++) {
            if (j != 0) {
                fmt::print(outfile, " ");
            }
            if (j >= i) {
                fmt::print(outfile, "{:12.5f}", CVoigt(i, j));
            } else {
                fmt::print(outfile, "            ");
            }
        }
        fmt::print(outfile, "\n");
    }

    auto eigenvalues = CVoigt.eigenvalues();
    auto sorter = [](std::complex<double> i, std::complex<double> j) {
        return std::abs(i) < std::abs(j);
    };
    std::sort(eigenvalues.data(), eigenvalues.data() + eigenvalues.size(), sorter);
    fmt::print(outfile, "# eigen values of the stiffness tensor (GPa)\n");
    for (size_t i=0; i<6; i++) {
        auto& value = eigenvalues(i);
        if (value.imag() == 0) {
            fmt::print(outfile, "{:12.5f}\n", value.real());
        } else {
            fmt::print(outfile, "{:12.5f} + {:12.5f}i\n", value.real(), value.imag());
        }
    }

    double A = (CVoigt(0, 0) + CVoigt(1, 1) + CVoigt(2, 2)) / 3.0;
    double B = (CVoigt(1, 2) + CVoigt(0, 2) + CVoigt(0, 1)) / 3.0;
    double C = (CVoigt(3, 3) + CVoigt(4, 4) + CVoigt(5, 5)) / 3.0;

    double a = (SVoigt(0, 0) + SVoigt(1, 1) + SVoigt(2, 2)) / 3.0;
    double b = (SVoigt(1, 2) + SVoigt(0, 2) + SVoigt(0, 1)) / 3.0;
    double c = (SVoigt(3, 3) + SVoigt(4, 4) + SVoigt(5, 5)) / 3.0;

    double KV = (A + 2.0 * B) / 3.0;
    double GV = (A - B + 3.0 * C) / 5.0;
    double YV = 1.0 / (1.0 / (3.0 * GV) + 1.0 / (9.0 * KV));
    double PV = (1.0 - 3.0 * GV / (3.0 * KV + GV)) / 2.0;

    double KR = 1.0 / (3.0 * a + 6.0 * b);
    double GR = 5.0 / (4.0 * a - 4.0 * b + 3.0 * c);
    double YR = 1.0 / (1.0 / (3.0 * GR) + 1.0 / (9.0 * KR));
    double PR = (1.0 - 3.0 * GR / (3.0 * KR + GR)) / 2.0;

    double KH = (KV + KR) / 2.0;
    double GH = (GV + GR) / 2.0;
    double YH = 1.0 / (1.0 / (3.0 * GH) + 1.0 / (9.0 * KH));
    double PH = (1.0 - 3.0 * GH / (3.0 * KH + GH)) / 2.0;

    fmt::print(outfile, "# Bulk modulus (GPa) | Young's modulus (GPa) | Shear modulus (GPa) | Poisson's ratio\n");
    fmt::print(outfile, "# Voigt averaging\n");
    fmt::print(outfile, "{:12.5f} {:12.5f} {:12.5f} {:12.5f}\n", KV, YV, GV, PV);
    fmt::print(outfile, "# Reuss averaging\n");
    fmt::print(outfile, "{:12.5f} {:12.5f} {:12.5f} {:12.5f}\n", KR, YR, GR, PR);
    fmt::print(outfile, "# Hill averaging\n");
    fmt::print(outfile, "{:12.5f} {:12.5f} {:12.5f} {:12.5f}\n", KH, YH, GH, PH);

    return 0;
}


//==================================================================================================
/*!
  @file

  @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_CONSTANT_PI_HPP_INCLUDED
#define BOOST_SIMD_CONSTANT_PI_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{
  /*!
    @ingroup group-constant

    Generate value \f$\pi\f$ that is the half length of a circle of radius one
    ... in normal temperature and pressure conditions.

    @par Header <boost/simd/constant/pi.hpp>

    @par Semantic:

    @code
    T r = Pi<T>();
    @endcode

    is similar to:

    @code
    T r = T(4*atan(1));
    @endcode

    @return The Pi constant for the proper type
  **/
  template<typename T> T Pi();

  namespace functional
  {
    /*!
      @ingroup group-callable-constant
      Generate the  constant pi.

      @return The Pi constant for the proper type
    **/
    Value Pi();
  }
} }
#endif

#include <boost/simd/constant/scalar/pi.hpp>
#include <boost/simd/constant/simd/pi.hpp>

#endif