#include <NTL/GF2EXFactoring.h>
#include <NTL/vec_GF2XVec.h>
#include <NTL/fileio.h>
#include <NTL/FacVec.h>

#include <stdio.h>

#include <NTL/new.h>

NTL_START_IMPL

static
void IterSqr(GF2E& c, const GF2E& a, long n)
{
   GF2E res;

   long i;

   res = a;

   for (i = 0; i < n; i++)
      sqr(res, res);

   c = res;
}

void SquareFreeDecomp(vec_pair_GF2EX_long& u, const GF2EX& ff)
{
   GF2EX f = ff;

   if (!IsOne(LeadCoeff(f)))
      Error("SquareFreeDecomp: bad args");

   GF2EX r, t, v, tmp1;
   long m, j, finished, done;

   u.SetLength(0);

   if (deg(f) == 0)
      return;

   m = 1;
   finished = 0;

   do {
      j = 1;
      diff(tmp1, f);
      GCD(r, f, tmp1);
      div(t, f, r);

      if (deg(t) > 0) {
         done = 0;
         do {
            GCD(v, r, t);
            div(tmp1, t, v);
            if (deg(tmp1) > 0) append(u, cons(tmp1, j*m));
            if (deg(v) > 0) {
               div(r, r, v);
               t = v;
               j++;
            }
            else
               done = 1;
         } while (!done);
         if (deg(r) == 0) finished = 1;
      }

      if (!finished) {
         /* r is a square */

         long k, d;
         d = deg(r)/2;
         f.rep.SetLength(d+1);
         for (k = 0; k <= d; k++)
            IterSqr(f.rep[k], r.rep[k*2], GF2E::degree()-1);
         m = m*2;
      }
   } while (!finished);
}

static
void NullSpace(long& r, vec_long& D, vec_GF2XVec& M, long verbose)
{
   long k, l, n;
   long i, j;
   long pos;
   GF2X t1, t2;
   GF2X *x, *y;

   const GF2XModulus& p = GF2E::modulus();

   n = M.length();

   D.SetLength(n);
   for (j = 0; j < n; j++) D[j] = -1;

   r = 0;

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

      if (verbose && k % 10 == 0) cerr << "+";

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

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

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

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

         for (i = l+1; i < n; i++) {
            // M[i] = M[i] + M[l]*M[i,k]

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

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

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

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

         D[k] = l;   // variable k is defined by row l
         l++;

      }
      else {
         r++;
      }
   }
}

static
void BuildMatrix(vec_GF2XVec& M, long n, const GF2EX& g, const GF2EXModulus& F,
                 long verbose)
{
   long i, j, m;
   GF2EX h;

   M.SetLength(n);
   for (i = 0; i < n; i++)
      M[i].SetSize(n, 2*GF2E::WordLength());

   set(h);
   for (j = 0; j < n; j++) {
      if (verbose && j % 10 == 0) cerr << "+";

      m = deg(h);
      for (i = 0; i < n; i++) {
         if (i <= m)
            M[i][j] = rep(h.rep[i]);
         else
            clear(M[i][j]);
      }

      if (j < n-1)
         MulMod(h, h, g, F);
   }

   for (i = 0; i < n; i++)
      add(M[i][i], M[i][i], 1);

}

static
void TraceMap(GF2EX& h, const GF2EX& a, const GF2EXModulus& F)

// one could consider making a version based on modular composition,
// as in ComposeFrobeniusMap...

{
   GF2EX res, tmp;

   res = a;
   tmp = a;

   long i;
   for (i = 0; i < GF2E::degree()-1; i++) {
      SqrMod(tmp, tmp, F);
      add(res, res, tmp);
   }

   h = res;
}

void PlainFrobeniusMap(GF2EX& h, const GF2EXModulus& F)
{
   GF2EX res;

   SetX(res);
   long i;
   for (i = 0; i < GF2E::degree(); i++)
      SqrMod(res, res, F);

   h = res;
}

long UseComposeFrobenius(long d, long n)
{
   long i;
   i = 1;
   while (i <= d) i = i << 1;
   i = i >> 1;

   i = i >> 1;
   long m = 1;

   long dz;

   if (n == 2) {
      dz = 1;
   }
   else {
      while (i) {
         long m1 = 2*m;
         if (i & d) m1++;

         if (m1 >= NTL_BITS_PER_LONG-1 || (1L << m1) >= n) break;

         m = m1;
         i = i >> 1;
      }

      dz = 1L << m;
   }

   long rootn = SqrRoot(n);
   long cnt = 0;

   if (i) {
      cnt += SqrRoot(dz+1);
      i = i >> 1;
   }

   while (i) {
      cnt += rootn;
      i = i >> 1;
   }

   return 4*cnt <= d;
}

void ComposeFrobeniusMap(GF2EX& y, const GF2EXModulus& F)
{
   long d = GF2E::degree();
   long n = deg(F);

   long i;
   i = 1;
   while (i <= d) i = i << 1;
   i = i >> 1;

   GF2EX z(INIT_SIZE, n), z1(INIT_SIZE, n);

   i = i >> 1;
   long m = 1;

   if (n == 2) {
      SetX(z);
      SqrMod(z, z, F);
   }
   else {
      while (i) {
         long m1 = 2*m;
         if (i & d) m1++;

         if (m1 >= NTL_BITS_PER_LONG-1 || (1L << m1) >= n) break;

         m = m1;
         i = i >> 1;
      }

      clear(z);
      SetCoeff(z, 1L << m);
   }

   while (i) {
      z1 = z;

      long j, k, dz;
      dz = deg(z);

      for (j = 0; j <= dz; j++)
         for (k = 0; k < m; k++)
            sqr(z1.rep[j], z1.rep[j]);

      CompMod(z, z1, z, F);
      m = 2*m;

      if (d & i) {
         SqrMod(z, z, F);
         m++;
      }

      i = i >> 1;
   }

   y = z;
}

void FrobeniusMap(GF2EX& h, const GF2EXModulus& F)
{
   long n = deg(F);
   long d = GF2E::degree();

   if (n == 1) {
      h = ConstTerm(F);
      return;
   }

   if (UseComposeFrobenius(d, n))
      ComposeFrobeniusMap(h, F);
   else
      PlainFrobeniusMap(h, F);
}

static
void RecFindRoots(vec_GF2E& x, const GF2EX& f)
{
   if (deg(f) == 0) return;

   if (deg(f) == 1) {
      long k = x.length();
      x.SetLength(k+1);
      x[k] = ConstTerm(f);
      return;
   }

   GF2EX h;

   GF2E r;

   {
      GF2EXModulus F;
      build(F, f);

      do {
         random(r);
         clear(h);
         SetCoeff(h, 1, r);
         TraceMap(h, h, F);
         GCD(h, h, f);
      } while (deg(h) <= 0 || deg(h) == deg(f));
   }

   RecFindRoots(x, h);
   div(h, f, h);
   RecFindRoots(x, h);
}

void FindRoots(vec_GF2E& x, const GF2EX& ff)
{
   GF2EX f = ff;

   if (!IsOne(LeadCoeff(f)))
      Error("FindRoots: bad args");

   x.SetMaxLength(deg(f));
   x.SetLength(0);
   RecFindRoots(x, f);
}

static
void RandomBasisElt(GF2EX& g, const vec_long& D, const vec_GF2XVec& M)
{
   static GF2X t1, t2;

   long n = D.length();

   long i, j, s;

   g.rep.SetLength(n);

   vec_GF2E& v = g.rep;

   for (j = n-1; j >= 0; j--) {
      if (D[j] == -1)
         random(v[j]);
      else {
         i = D[j];

         // v[j] = sum_{s=j+1}^{n-1} v[s]*M[i,s]

         clear(t1);

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

         conv(v[j], t1);
      }
   }

   g.normalize();
}

static
void split(GF2EX& f1, GF2EX& g1, GF2EX& f2, GF2EX& g2,
           const GF2EX& f, const GF2EX& g,
           const vec_GF2E& roots, long lo, long mid)
{
   long r = mid-lo+1;

   GF2EXModulus F;
   build(F, f);

   vec_GF2E lroots(INIT_SIZE, r);
   long i;

   for (i = 0; i < r; i++)
      lroots[i] = roots[lo+i];

   GF2EX h, a, d;
   BuildFromRoots(h, lroots);
   CompMod(a, h, g, F);

   GCD(f1, a, f);

   div(f2, f, f1);

   rem(g1, g, f1);
   rem(g2, g, f2);
}

static
void RecFindFactors(vec_GF2EX& factors, const GF2EX& f, const GF2EX& g,
                    const vec_GF2E& roots, long lo, long hi)
{
   long r = hi-lo+1;

   if (r == 0) return;

   if (r == 1) {
      append(factors, f);
      return;
   }

   GF2EX f1, g1, f2, g2;

   long mid = (lo+hi)/2;

   split(f1, g1, f2, g2, f, g, roots, lo, mid);

   RecFindFactors(factors, f1, g1, roots, lo, mid);
   RecFindFactors(factors, f2, g2, roots, mid+1, hi);
}

static
void FindFactors(vec_GF2EX& factors, const GF2EX& f, const GF2EX& g,
                 const vec_GF2E& roots)
{
   long r = roots.length();

   factors.SetMaxLength(r);
   factors.SetLength(0);

   RecFindFactors(factors, f, g, roots, 0, r-1);
}

#if 0

static
void IterFindFactors(vec_GF2EX& factors, const GF2EX& f,
                     const GF2EX& g, const vec_GF2E& roots)
{
   long r = roots.length();
   long i;
   GF2EX h;

   factors.SetLength(r);

   for (i = 0; i < r; i++) {
      add(h, g, roots[i]);
      GCD(factors[i], f, h);
   }
}

#endif

void SFBerlekamp(vec_GF2EX& factors, const GF2EX& ff, long verbose)
{
   GF2EX f = ff;

   if (!IsOne(LeadCoeff(f)))
      Error("SFBerlekamp: bad args");

   if (deg(f) == 0) {
      factors.SetLength(0);
      return;
   }

   if (deg(f) == 1) {
      factors.SetLength(1);
      factors[0] = f;
      return;
   }

   double t;

   long n = deg(f);

   GF2EXModulus F;

   build(F, f);

   GF2EX g, h;

   if (verbose) { cerr << "computing X^p..."; t = GetTime(); }
   FrobeniusMap(g, F);
   if (verbose) { cerr << (GetTime()-t) << "\n"; }

   vec_long D;
   long r;

   vec_GF2XVec M;

   if (verbose) { cerr << "building matrix..."; t = GetTime(); }
   BuildMatrix(M, n, g, F, verbose);
   if (verbose) { cerr << (GetTime()-t) << "\n"; }

   if (verbose) { cerr << "diagonalizing..."; t = GetTime(); }
   NullSpace(r, D, M, verbose);
   if (verbose) { cerr << (GetTime()-t) << "\n"; }

   if (verbose) cerr << "number of factors = " << r << "\n";

   if (r == 1) {
      factors.SetLength(1);
      factors[0] = f;
      return;
   }

   if (verbose) { cerr << "factor extraction..."; t = GetTime(); }

   vec_GF2E roots;

   RandomBasisElt(g, D, M);
   MinPolyMod(h, g, F, r);
   if (deg(h) == r) M.kill();
   FindRoots(roots, h);
   FindFactors(factors, f, g, roots);

   GF2EX g1;
   vec_GF2EX S, S1;
   long i;

   while (factors.length() < r) {
      if (verbose) cerr << "+";
      RandomBasisElt(g, D, M);
      S.kill();
      for (i = 0; i < factors.length(); i++) {
         const GF2EX& f = factors[i];
         if (deg(f) == 1) {
            append(S, f);
            continue;
         }
         build(F, f);
         rem(g1, g, F);
         if (deg(g1) <= 0) {
            append(S, f);
            continue;
         }
         MinPolyMod(h, g1, F, min(deg(f), r-factors.length()+1));
         FindRoots(roots, h);
         S1.kill();
         FindFactors(S1, f, g1, roots);
         append(S, S1);
      }
      swap(factors, S);
   }

   if (verbose) { cerr << (GetTime()-t) << "\n"; }

   if (verbose) {
      cerr << "degrees:";
      long i;
      for (i = 0; i < factors.length(); i++)
         cerr << " " << deg(factors[i]);
      cerr << "\n";
   }
}

void berlekamp(vec_pair_GF2EX_long& factors, const GF2EX& f, long verbose)
{
   double t;
   vec_pair_GF2EX_long sfd;
   vec_GF2EX x;

   if (!IsOne(LeadCoeff(f)))
      Error("berlekamp: bad args");

   if (verbose) { cerr << "square-free decomposition..."; t = GetTime(); }
   SquareFreeDecomp(sfd, f);
   if (verbose) cerr << (GetTime()-t) << "\n";

   factors.SetLength(0);

   long i, j;

   for (i = 0; i < sfd.length(); i++) {
      if (verbose) {
         cerr << "factoring multiplicity " << sfd[i].b
              << ", deg = " << deg(sfd[i].a) << "\n";
      }

      SFBerlekamp(x, sfd[i].a, verbose);

      for (j = 0; j < x.length(); j++)
         append(factors, cons(x[j], sfd[i].b));
   }
}

static
void AddFactor(vec_pair_GF2EX_long& factors, const GF2EX& g, long d, long verbose)
{
   if (verbose)
      cerr << "degree=" << d << ", number=" << deg(g)/d << "\n";
   append(factors, cons(g, d));
}

static
void ProcessTable(GF2EX& f, vec_pair_GF2EX_long& factors,
                  const GF2EXModulus& F, long limit, const vec_GF2EX& tbl,
                  long d, long verbose)

{
   if (limit == 0) return;

   if (verbose) cerr << "+";

   GF2EX t1;

   if (limit == 1) {
      GCD(t1, f, tbl[0]);
      if (deg(t1) > 0) {
         AddFactor(factors, t1, d, verbose);
         div(f, f, t1);
      }

      return;
   }

   long i;

   t1 = tbl[0];
   for (i = 1; i < limit; i++)
      MulMod(t1, t1, tbl[i], F);

   GCD(t1, f, t1);

   if (deg(t1) == 0) return;

   div(f, f, t1);

   GF2EX t2;

   i = 0;
   d = d - limit + 1;

   while (2*d <= deg(t1)) {
      GCD(t2, tbl[i], t1);
      if (deg(t2) > 0) {
         AddFactor(factors, t2, d, verbose);
         div(t1, t1, t2);
      }

      i++;
      d++;
   }

   if (deg(t1) > 0)
      AddFactor(factors, t1, deg(t1), verbose);
}

void TraceMap(GF2EX& w, const GF2EX& a, long d, const GF2EXModulus& F,
              const GF2EX& b)

{
   if (d < 0) Error("TraceMap: bad args");

   GF2EX y, z, t;

   z = b;
   y = a;
   clear(w);

   while (d) {
      if (d == 1) {
         if (IsZero(w))
            w = y;
         else {
            CompMod(w, w, z, F);
            add(w, w, y);
         }
      }
      else if ((d & 1) == 0) {
         Comp2Mod(z, t, z, y, z, F);
         add(y, t, y);
      }
      else if (IsZero(w)) {
         w = y;
         Comp2Mod(z, t, z, y, z, F);
         add(y, t, y);
      }
      else {
         Comp3Mod(z, t, w, z, y, w, z, F);
         add(w, w, y);
         add(y, t, y);
      }

      d = d >> 1;
   }
}

void PowerCompose(GF2EX& y, const GF2EX& h, long q, const GF2EXModulus& F)
{
   if (q < 0) Error("powerCompose: bad args");

   GF2EX z(INIT_SIZE, F.n);
   long sw;

   z = h;
   SetX(y);

   while (q) {
      sw = 0;

      if (q > 1) sw = 2;
      if (q & 1) {
         if (IsX(y))
            y = z;
         else
            sw = sw | 1;
      }

      switch (sw) {
      case 0:
         break;

      case 1:
         CompMod(y, y, z, F);
         break;

      case 2:
         CompMod(z, z, z, F);
         break;

      case 3:
         Comp2Mod(y, z, y, z, z, F);
         break;
      }

      q = q >> 1;
   }
}

long ProbIrredTest(const GF2EX& f, long iter)
{
   long n = deg(f);

   if (n <= 0) return 0;
   if (n == 1) return 1;

   GF2EXModulus F;

   build(F, f);

   GF2EX b, r, s;

   FrobeniusMap(b, F);

   long all_zero = 1;

   long i;

   for (i = 0; i < iter; i++) {
      random(r, n);
      TraceMap(s, r, n, F, b);

      all_zero = all_zero && IsZero(s);

      if (deg(s) > 0) return 0;
   }

   if (!all_zero || (n & 1)) return 1;

   PowerCompose(s, b, n/2, F);
   return !IsX(s);
}

long GF2EX_BlockingFactor = 10;

void DDF(vec_pair_GF2EX_long& factors, const GF2EX& ff, const GF2EX& hh,
         long verbose)
{
   GF2EX f = ff;
   GF2EX h = hh;

   if (!IsOne(LeadCoeff(f)))
      Error("DDF: bad args");

   factors.SetLength(0);

   if (deg(f) == 0)
      return;

   if (deg(f) == 1) {
      AddFactor(factors, f, 1, verbose);
      return;
   }

   long CompTableSize = 2*SqrRoot(deg(f));

   long GCDTableSize = GF2EX_BlockingFactor;

   GF2EXModulus F;
   build(F, f);

   GF2EXArgument H;

   build(H, h, F, min(CompTableSize, deg(f)));

   long i, d, limit, old_n;
   GF2EX g, X;

   vec_GF2EX tbl(INIT_SIZE, GCDTableSize);

   SetX(X);

   i = 0;
   g = h;
   d = 1;
   limit = GCDTableSize;

   while (2*d <= deg(f)) {

      old_n = deg(f);
      add(tbl[i], g, X);
      i++;
      if (i == limit) {
         ProcessTable(f, factors, F, i, tbl, d, verbose);
         i = 0;
      }

      d = d + 1;
      if (2*d <= deg(f)) {
         // we need to go further

         if (deg(f) < old_n) {
            // f has changed

            build(F, f);
            rem(h, h, f);
            rem(g, g, f);
            build(H, h, F, min(CompTableSize, deg(f)));
         }

         CompMod(g, g, H, F);
      }
   }

   ProcessTable(f, factors, F, i, tbl, d-1, verbose);

   if (!IsOne(f)) AddFactor(factors, f, deg(f), verbose);
}

void RootEDF(vec_GF2EX& factors, const GF2EX& f, long verbose)
{
   vec_GF2E roots;
   double t;

   if (verbose) { cerr << "finding roots..."; t = GetTime(); }
   FindRoots(roots, f);
   if (verbose) { cerr << (GetTime()-t) << "\n"; }

   long r = roots.length();
   factors.SetLength(r);
   for (long j = 0; j < r; j++) {
      SetX(factors[j]);
      add(factors[j], factors[j], roots[j]);
   }
}

static
void EDFSplit(vec_GF2EX& v, const GF2EX& f, const GF2EX& b, long d)
{
   GF2EX a, g, h;
   GF2EXModulus F;
   vec_GF2E roots;

   build(F, f);
   long n = F.n;
   long r = n/d;
   random(a, n);
   TraceMap(g, a, d, F, b);
   MinPolyMod(h, g, F, r);
   FindRoots(roots, h);
   FindFactors(v, f, g, roots);
}

static
void RecEDF(vec_GF2EX& factors, const GF2EX& f, const GF2EX& b, long d,
            long verbose)
{
   vec_GF2EX v;
   long i;
   GF2EX bb;

   if (verbose) cerr << "+";

   EDFSplit(v, f, b, d);
   for (i = 0; i < v.length(); i++) {
      if (deg(v[i]) == d) {
         append(factors, v[i]);
      }
      else {
         GF2EX bb;
         rem(bb, b, v[i]);
         RecEDF(factors, v[i], bb, d, verbose);
      }
   }
}

void EDF(vec_GF2EX& factors, const GF2EX& ff, const GF2EX& bb,
         long d, long verbose)

{
   GF2EX f = ff;
   GF2EX b = bb;

   if (!IsOne(LeadCoeff(f)))
      Error("EDF: bad args");

   long n = deg(f);
   long r = n/d;

   if (r == 0) {
      factors.SetLength(0);
      return;
   }

   if (r == 1) {
      factors.SetLength(1);
      factors[0] = f;
      return;
   }

   if (d == 1) {
      RootEDF(factors, f, verbose);
      return;
   }

   double t;
   if (verbose) {
      cerr << "computing EDF(" << d << "," << r << ")...";
      t = GetTime();
   }

   factors.SetLength(0);

   RecEDF(factors, f, b, d, verbose);

   if (verbose) cerr << (GetTime()-t) << "\n";
}

void SFCanZass(vec_GF2EX& factors, const GF2EX& ff, long verbose)
{
   GF2EX f = ff;

   if (!IsOne(LeadCoeff(f)))
      Error("SFCanZass: bad args");

   if (deg(f) == 0) {
      factors.SetLength(0);
      return;
   }

   if (deg(f) == 1) {
      factors.SetLength(1);
      factors[0] = f;
      return;
   }

   factors.SetLength(0);

   double t;

   GF2EXModulus F;
   build(F, f);

   GF2EX h;

   if (verbose) { cerr << "computing X^p..."; t = GetTime(); }
   FrobeniusMap(h, F);
   if (verbose) { cerr << (GetTime()-t) << "\n"; }

   vec_pair_GF2EX_long u;
   if (verbose) { cerr << "computing DDF..."; t = GetTime(); }
   NewDDF(u, f, h, verbose);
   if (verbose) {
      t = GetTime()-t;
      cerr << "DDF time: " << t << "\n";
   }

   GF2EX hh;
   vec_GF2EX v;

   long i;
   for (i = 0; i < u.length(); i++) {
      const GF2EX& g = u[i].a;
      long d = u[i].b;
      long r = deg(g)/d;

      if (r == 1) {
         // g is already irreducible

         append(factors, g);
      }
      else {
         // must perform EDF

         if (d == 1) {
            // root finding
            RootEDF(v, g, verbose);
            append(factors, v);
         }
         else {
            // general case
            rem(hh, h, g);
            EDF(v, g, hh, d, verbose);
            append(factors, v);
         }
      }
   }
}

void CanZass(vec_pair_GF2EX_long& factors, const GF2EX& f, long verbose)
{
   if (!IsOne(LeadCoeff(f)))
      Error("CanZass: bad args");

   double t;
   vec_pair_GF2EX_long sfd;
   vec_GF2EX x;

   if (verbose) { cerr << "square-free decomposition..."; t = GetTime(); }
   SquareFreeDecomp(sfd, f);
   if (verbose) cerr << (GetTime()-t) << "\n";

   factors.SetLength(0);

   long i, j;

   for (i = 0; i < sfd.length(); i++) {
      if (verbose) {
         cerr << "factoring multiplicity " << sfd[i].b
              << ", deg = " << deg(sfd[i].a) << "\n";
      }

      SFCanZass(x, sfd[i].a, verbose);

      for (j = 0; j < x.length(); j++)
         append(factors, cons(x[j], sfd[i].b));
   }
}

void mul(GF2EX& f, const vec_pair_GF2EX_long& v)
{
   long i, j, n;

   n = 0;
   for (i = 0; i < v.length(); i++)
      n += v[i].b*deg(v[i].a);

   GF2EX g(INIT_SIZE, n+1);

   set(g);
   for (i = 0; i < v.length(); i++)
      for (j = 0; j < v[i].b; j++) {
         mul(g, g, v[i].a);
      }

   f = g;
}

static
long BaseCase(const GF2EX& h, long q, long a, const GF2EXModulus& F)
{
   long b, e;
   GF2EX lh(INIT_SIZE, F.n);

   lh = h;
   b = 1;
   e = 0;
   while (e < a-1 && !IsX(lh)) {
      e++;
      b *= q;
      PowerCompose(lh, lh, q, F);
   }

   if (!IsX(lh)) b *= q;

   return b;
}

static
void TandemPowerCompose(GF2EX& y1, GF2EX& y2, const GF2EX& h,
                        long q1, long q2, const GF2EXModulus& F)
{
   GF2EX z(INIT_SIZE, F.n);
   long sw;

   z = h;
   SetX(y1);
   SetX(y2);

   while (q1 || q2) {
      sw = 0;

      if (q1 > 1 || q2 > 1) sw = 4;

      if (q1 & 1) {
         if (IsX(y1))
            y1 = z;
         else
            sw = sw | 2;
      }

      if (q2 & 1) {
         if (IsX(y2))
            y2 = z;
         else
            sw = sw | 1;
      }

      switch (sw) {
      case 0:
         break;

      case 1:
         CompMod(y2, y2, z, F);
         break;

      case 2:
         CompMod(y1, y1, z, F);
         break;

      case 3:
         Comp2Mod(y1, y2, y1, y2, z, F);
         break;

      case 4:
         CompMod(z, z, z, F);
         break;

      case 5:
         Comp2Mod(z, y2, z, y2, z, F);
         break;

      case 6:
         Comp2Mod(z, y1, z, y1, z, F);
         break;

      case 7:
         Comp3Mod(z, y1, y2, z, y1, y2, z, F);
         break;
      }

      q1 = q1 >> 1;
      q2 = q2 >> 1;
   }
}

static
long RecComputeDegree(long u, const GF2EX& h, const GF2EXModulus& F,
                      FacVec& fvec)
{
   if (IsX(h)) return 1;

   if (fvec[u].link == -1) return BaseCase(h, fvec[u].q, fvec[u].a, F);

   GF2EX h1, h2;
   long q1, q2, r1, r2;

   q1 = fvec[fvec[u].link].val;
   q2 = fvec[fvec[u].link+1].val;

   TandemPowerCompose(h1, h2, h, q1, q2, F);
   r1 = RecComputeDegree(fvec[u].link, h2, F, fvec);
   r2 = RecComputeDegree(fvec[u].link+1, h1, F, fvec);
   return r1*r2;
}

long RecComputeDegree(const GF2EX& h, const GF2EXModulus& F)
   // f = F.f is assumed to be an "equal degree" polynomial
   // h = X^p mod f
   // the common degree of the irreducible factors of f is computed
{
   if (F.n == 1 || IsX(h))
      return 1;

   FacVec fvec;

   FactorInt(fvec, F.n);

   return RecComputeDegree(fvec.length()-1, h, F, fvec);
}

void FindRoot(GF2E& root, const GF2EX& ff)
// finds a root of ff.
// assumes that ff is monic and splits into distinct linear factors

{
   GF2EXModulus F;
   GF2EX h, h1, f;
   GF2E r;

   f = ff;

   if (!IsOne(LeadCoeff(f)))
      Error("FindRoot: bad args");

   if (deg(f) == 0)
      Error("FindRoot: bad args");

   while (deg(f) > 1) {
      build(F, f);
      random(r);
      clear(h);
      SetCoeff(h, 1, r);
      TraceMap(h, h, F);
      GCD(h, h, f);
      if (deg(h) > 0 && deg(h) < deg(f)) {
         if (deg(h) > deg(f)/2)
            div(f, f, h);
         else
            f = h;
      }
   }

   root = ConstTerm(f);
}

static
long power(long a, long e)
{
   long i, res;

   res = 1;
   for (i = 1; i <= e; i++)
      res = res * a;

   return res;
}

static
long IrredBaseCase(const GF2EX& h, long q, long a, const GF2EXModulus& F)
{
   long e;
   GF2EX X, s, d;

   e = power(q, a-1);
   PowerCompose(s, h, e, F);
   SetX(X);
   add(s, s, X);
   GCD(d, F.f, s);
   return IsOne(d);
}

static
long RecIrredTest(long u, const GF2EX& h, const GF2EXModulus& F,
                 const FacVec& fvec)
{
   long  q1, q2;
   GF2EX h1, h2;

   if (IsX(h)) return 0;

   if (fvec[u].link == -1) {
      return IrredBaseCase(h, fvec[u].q, fvec[u].a, F);
   }

   q1 = fvec[fvec[u].link].val;
   q2 = fvec[fvec[u].link+1].val;

   TandemPowerCompose(h1, h2, h, q1, q2, F);
   return RecIrredTest(fvec[u].link, h2, F, fvec)
          && RecIrredTest(fvec[u].link+1, h1, F, fvec);
}

long DetIrredTest(const GF2EX& f)
{
   if (deg(f) <= 0) return 0;
   if (deg(f) == 1) return 1;

   GF2EXModulus F;

   build(F, f);

   GF2EX h;

   FrobeniusMap(h, F);

   GF2EX s;
   PowerCompose(s, h, F.n, F);
   if (!IsX(s)) return 0;

   FacVec fvec;

   FactorInt(fvec, F.n);

   return RecIrredTest(fvec.length()-1, h, F, fvec);
}

long IterIrredTest(const GF2EX& f)
{
   if (deg(f) <= 0) return 0;
   if (deg(f) == 1) return 1;

   GF2EXModulus F;

   build(F, f);

   GF2EX h;

   FrobeniusMap(h, F);

   long CompTableSize = 2*SqrRoot(deg(f));

   GF2EXArgument H;

   build(H, h, F, CompTableSize);

   long i, d, limit, limit_sqr;
   GF2EX g, X, t, prod;

   SetX(X);

   i = 0;
   g = h;
   d = 1;
   limit = 2;
   limit_sqr = limit*limit;

   set(prod);

   while (2*d <= deg(f)) {
      add(t, g, X);
      MulMod(prod, prod, t, F);
      i++;
      if (i == limit_sqr) {
         GCD(t, f, prod);
         if (!IsOne(t)) return 0;

         set(prod);
         limit++;
         limit_sqr = limit*limit;
         i = 0;
      }

      d = d + 1;
      if (2*d <= deg(f)) {
         CompMod(g, g, H, F);
      }
   }

   if (i > 0) {
      GCD(t, f, prod);
      if (!IsOne(t)) return 0;
   }

   return 1;
}

static
void MulByXPlusY(vec_GF2EX& h, const GF2EX& f, const GF2EX& g)
// h represents the bivariate polynomial h[0] + h[1]*Y + ... + h[n-1]*Y^k,
// where the h[i]'s are polynomials in X, each of degree < deg(f),
// and k < deg(g).
// h is replaced by the bivariate polynomial h*(X+Y) (mod f(X), g(Y)).

{
   long n = deg(g);
   long k = h.length()-1;

   if (k < 0) return;

   if (k < n-1) {
      h.SetLength(k+2);
      h[k+1] = h[k];
      for (long i = k; i >= 1; i--) {
         MulByXMod(h[i], h[i], f);
         add(h[i], h[i], h[i-1]);
      }
      MulByXMod(h[0], h[0], f);
   }
   else {
      GF2EX b, t;

      b = h[n-1];
      for (long i = n-1; i >= 1; i--) {
         mul(t, b, g.rep[i]);
         MulByXMod(h[i], h[i], f);
         add(h[i], h[i], h[i-1]);
         add(h[i], h[i], t);
      }
      mul(t, b, g.rep[0]);
      MulByXMod(h[0], h[0], f);
      add(h[0], h[0], t);
   }

   // normalize

   k = h.length()-1;
   while (k >= 0 && IsZero(h[k])) k--;
   h.SetLength(k+1);
}

static
void IrredCombine(GF2EX& x, const GF2EX& f, const GF2EX& g)
{
   if (deg(f) < deg(g)) {
      IrredCombine(x, g, f);
      return;
   }

   // deg(f) >= deg(g)...not necessary, but maybe a little more
   //                    time & space efficient

   long df = deg(f);
   long dg = deg(g);
   long m = df*dg;

   vec_GF2EX h(INIT_SIZE, dg);

   long i;
   for (i = 0; i < dg; i++) h[i].SetMaxLength(df);

   h.SetLength(1);
   set(h[0]);

   vec_GF2E a;

   a.SetLength(2*m);

   for (i = 0; i < 2*m; i++) {
      a[i] = ConstTerm(h[0]);
      if (i < 2*m-1)
         MulByXPlusY(h, f, g);
   }

   MinPolySeq(x, a, m);
}

static
void BuildPrimePowerIrred(GF2EX& f, long q, long e)
{
   long n = power(q, e);

   do {
      random(f, n);
      SetCoeff(f, n);
   } while (!IterIrredTest(f));
}

static
void RecBuildIrred(GF2EX& f, long u, const FacVec& fvec)
{
   if (fvec[u].link == -1)
      BuildPrimePowerIrred(f, fvec[u].q, fvec[u].a);
   else {
      GF2EX g, h;
      RecBuildIrred(g, fvec[u].link, fvec);
      RecBuildIrred(h, fvec[u].link+1, fvec);
      IrredCombine(f, g, h);
   }
}

void BuildIrred(GF2EX& f, long n)
{
   if (n <= 0)
      Error("BuildIrred: n must be positive");

   if (NTL_OVERFLOW(n, 1, 0))
      Error("overflow in BuildIrred");

   if (n == 1) {
      SetX(f);
      return;
   }

   FacVec fvec;

   FactorInt(fvec, n);

   RecBuildIrred(f, fvec.length()-1, fvec);
}

#if 0
void BuildIrred(GF2EX& f, long n)
{
   if (n <= 0)
      Error("BuildIrred: n must be positive");

   if (NTL_OVERFLOW(n, 1, 0))
      Error("overflow in BuildIrred");

   if (n == 1) {
      SetX(f);
      return;
   }

   GF2EX g;

   do {
      random(g, n);
      SetCoeff(g, n);
   } while (!IterIrredTest(g));

   f = g;

}
#endif

void BuildRandomIrred(GF2EX& f, const GF2EX& g)
{
   GF2EXModulus G;
   GF2EX h, ff;

   build(G, g);
   do {
      random(h, deg(g));
      IrredPolyMod(ff, h, G);
   } while (deg(ff) < deg(g));

   f = ff;
}

/************* NEW DDF ****************/

long GF2EX_GCDTableSize = 4;
char GF2EX_stem[256] = "";

double GF2EXFileThresh = 256;

static vec_GF2EX BabyStepFile;
static vec_GF2EX GiantStepFile;

static long use_files;

static
double CalcTableSize(long n, long k)
{
   double sz = GF2E::storage();
   sz = sz * n;
   sz = sz + NTL_VECTOR_HEADER_SIZE + sizeof(vec_GF2E);
   sz = sz * k;
   sz = sz/1024;
   return sz;
}

static
void GenerateBabySteps(GF2EX& h1, const GF2EX& f, const GF2EX& h, long k,
                       long verbose)

{
   double t;

   if (verbose) { cerr << "generating baby steps..."; t = GetTime(); }

   GF2EXModulus F;
   build(F, f);

   GF2EXArgument H;

#if 0
   double n2 = sqrt(double(F.n));
   double n4 = sqrt(n2);
   double n34 = n2*n4;
   long sz = long(ceil(n34/sqrt(sqrt(2.0))));
#else
   long sz = 2*SqrRoot(F.n);
#endif

   build(H, h, F, sz);

   h1 = h;

   long i;

   long HexOutput = GF2X::HexOutput;
   GF2X::HexOutput = 1;

   if (!use_files) {
      BabyStepFile.kill();
      BabyStepFile.SetLength(k-1);
   }

   for (i = 1; i <= k-1; i++) {
      if (use_files) {
         ofstream s;
         OpenWrite(s, FileName(GF2EX_stem, "baby", i));
         s << h1 << "\n";
         s.close();
      }
      else
         BabyStepFile(i) = h1;

      CompMod(h1, h1, H, F);
      if (verbose) cerr << "+";
   }

   if (verbose)
      cerr << (GetTime()-t) << "\n";

   GF2X::HexOutput = HexOutput;
}

static
void GenerateGiantSteps(const GF2EX& f, const GF2EX& h, long l, long verbose)
{

   double t;

   if (verbose) { cerr << "generating giant steps..."; t = GetTime(); }

   GF2EXModulus F;
   build(F, f);

   GF2EXArgument H;

#if 0
   double n2 = sqrt(double(F.n));
   double n4 = sqrt(n2);
   double n34 = n2*n4;
   long sz = long(ceil(n34/sqrt(sqrt(2.0))));
#else
   long sz = 2*SqrRoot(F.n);
#endif

   build(H, h, F, sz);

   GF2EX h1;

   h1 = h;

   long i;

   long HexOutput = GF2X::HexOutput;
   GF2X::HexOutput = 1;

   if (!use_files) {
      GiantStepFile.kill();
      GiantStepFile.SetLength(l);
   }

   for (i = 1; i <= l-1; i++) {
      if (use_files) {
         ofstream s;
         OpenWrite(s, FileName(GF2EX_stem, "giant", i));
         s << h1 << "\n";
         s.close();
      }
      else
         GiantStepFile(i) = h1;

      CompMod(h1, h1, H, F);
      if (verbose) cerr << "+";
   }

   if (use_files) {
      ofstream s;
      OpenWrite(s, FileName(GF2EX_stem, "giant", i));
      s << h1 << "\n";
      s.close();
   }
   else
      GiantStepFile(i) = h1;

   if (verbose)
      cerr << (GetTime()-t) << "\n";

   GF2X::HexOutput = HexOutput;
}

static
void FileCleanup(long k, long l)
{
   if (use_files) {
      long i;

      for (i = 1; i <= k-1; i++)
         remove(FileName(GF2EX_stem, "baby", i));

      for (i = 1; i <= l; i++)
         remove(FileName(GF2EX_stem, "giant", i));
   }
   else {
      BabyStepFile.kill();
      GiantStepFile.kill();
   }
}

static
void NewAddFactor(vec_pair_GF2EX_long& u, const GF2EX& g, long m, long verbose)
{
   long len = u.length();

   u.SetLength(len+1);
   u[len].a = g;
   u[len].b = m;

   if (verbose) {
      cerr << "split " << m << " " << deg(g) << "\n";
   }
}

static
void NewProcessTable(vec_pair_GF2EX_long& u, GF2EX& f, const GF2EXModulus& F,
                     vec_GF2EX& buf, long size, long StartInterval,
                     long IntervalLength, long verbose)

{
   if (size == 0) return;

   GF2EX& g = buf[size-1];

   long i;

   for (i = 0; i < size-1; i++)
      MulMod(g, g, buf[i], F);

   GCD(g, f, g);

   if (deg(g) == 0) return;

   div(f, f, g);

   long d = (StartInterval-1)*IntervalLength + 1;
   i = 0;
   long interval = StartInterval;

   while (i < size-1 && 2*d <= deg(g)) {
      GCD(buf[i], buf[i], g);
      if (deg(buf[i]) > 0) {
         NewAddFactor(u, buf[i], interval, verbose);
         div(g, g, buf[i]);
      }

      i++;
      interval++;
      d += IntervalLength;
   }

   if (deg(g) > 0) {
      if (i == size-1)
         NewAddFactor(u, g, interval, verbose);
      else
         NewAddFactor(u, g, (deg(g)+IntervalLength-1)/IntervalLength, verbose);
   }
}

static
void FetchGiantStep(GF2EX& g, long gs, const GF2EXModulus& F)
{
   if (use_files) {
      ifstream s;

      OpenRead(s, FileName(GF2EX_stem, "giant", gs));

      s >> g;
      s.close();
   }
   else
      g = GiantStepFile(gs);

   rem(g, g, F);
}

static
void FetchBabySteps(vec_GF2EX& v, long k)
{
   v.SetLength(k);

   SetX(v[0]);

   long i;
   for (i = 1; i <= k-1; i++) {
      if (use_files) {
         ifstream s;
         OpenRead(s, FileName(GF2EX_stem, "baby", i));
         s >> v[i];
         s.close();
      }
      else
         v[i] = BabyStepFile(i);
   }
}

static
void GiantRefine(vec_pair_GF2EX_long& u, const GF2EX& ff, long k, long l,
                 long verbose)

{
   double t;

   if (verbose) {
      cerr << "giant refine...";
      t = GetTime();
   }

   u.SetLength(0);

   vec_GF2EX BabyStep;

   FetchBabySteps(BabyStep, k);

   vec_GF2EX buf(INIT_SIZE, GF2EX_GCDTableSize);

   GF2EX f;
   f = ff;

   GF2EXModulus F;
   build(F, f);

   GF2EX g;
   GF2EX h;

   long size = 0;

   long first_gs;

   long d = 1;

   while (2*d <= deg(f)) {

      long old_n = deg(f);

      long gs = (d+k-1)/k;
      long bs = gs*k - d;

      if (bs == k-1) {
         size++;
         if (size == 1) first_gs = gs;
         FetchGiantStep(g, gs, F);
         add(buf[size-1], g, BabyStep[bs]);
      }
      else {
         add(h, g, BabyStep[bs]);
         MulMod(buf[size-1], buf[size-1], h, F);
      }

      if (verbose && bs == 0) cerr << "+";

      if (size == GF2EX_GCDTableSize && bs == 0) {
         NewProcessTable(u, f, F, buf, size, first_gs, k, verbose);
         if (verbose) cerr << "*";
         size = 0;
      }

      d++;

      if (2*d <= deg(f) && deg(f) < old_n) {
         build(F, f);

         long i;
         for (i = 1; i <= k-1; i++)
            rem(BabyStep[i], BabyStep[i], F);
      }
   }

   if (size > 0) {
      NewProcessTable(u, f, F, buf, size, first_gs, k, verbose);
      if (verbose) cerr << "*";
   }

   if (deg(f) > 0)
      NewAddFactor(u, f, 0, verbose);

   if (verbose) {
      t = GetTime()-t;
      cerr << "giant refine time: " << t << "\n";
   }
}

static
void IntervalRefine(vec_pair_GF2EX_long& factors, const GF2EX& ff,
                    long k, long gs, const vec_GF2EX& BabyStep, long verbose)

{
   vec_GF2EX buf(INIT_SIZE, GF2EX_GCDTableSize);

   GF2EX f;
   f = ff;

   GF2EXModulus F;
   build(F, f);

   GF2EX g;

   FetchGiantStep(g, gs, F);

   long size = 0;

   long first_d;

   long d = (gs-1)*k + 1;
   long bs = k-1;

   while (bs >= 0 && 2*d <= deg(f)) {

      long old_n = deg(f);

      if (size == 0) first_d = d;
      rem(buf[size], BabyStep[bs], F);
      add(buf[size], buf[size], g);
      size++;

      if (size == GF2EX_GCDTableSize) {
         NewProcessTable(factors, f, F, buf, size, first_d, 1, verbose);
         size = 0;
      }

      d++;
      bs--;

      if (bs >= 0 && 2*d <= deg(f) && deg(f) < old_n) {
         build(F, f);
         rem(g, g, F);
      }
   }

   NewProcessTable(factors, f, F, buf, size, first_d, 1, verbose);

   if (deg(f) > 0)
      NewAddFactor(factors, f, deg(f), verbose);
}

static
void BabyRefine(vec_pair_GF2EX_long& factors, const vec_pair_GF2EX_long& u,
                long k, long l, long verbose)

{
   double t;

   if (verbose) {
      cerr << "baby refine...";
      t = GetTime();
   }

   factors.SetLength(0);

   vec_GF2EX BabyStep;

   long i;
   for (i = 0; i < u.length(); i++) {
      const GF2EX& g = u[i].a;
      long gs = u[i].b;

      if (gs == 0 || 2*((gs-1)*k+1) > deg(g))
         NewAddFactor(factors, g, deg(g), verbose);
      else {
         if (BabyStep.length() == 0)
            FetchBabySteps(BabyStep, k);
         IntervalRefine(factors, g, k, gs, BabyStep, verbose);
      }
   }

   if (verbose) {
      t = GetTime()-t;
      cerr << "baby refine time: " << t << "\n";
   }
}

void NewDDF(vec_pair_GF2EX_long& factors,
            const GF2EX& f,
            const GF2EX& h,
            long verbose)

{
   if (!IsOne(LeadCoeff(f)))
      Error("NewDDF: bad args");

   if (deg(f) == 0) {
      factors.SetLength(0);
      return;
   }

   if (deg(f) == 1) {
      factors.SetLength(0);
      append(factors, cons(f, 1L));
      return;
   }

   if (!GF2EX_stem[0])
      sprintf(GF2EX_stem, "ddf-%ld", RandomBnd(10000));

   long B = deg(f)/2;
   long k = SqrRoot(B);
   long l = (B+k-1)/k;

   GF2EX h1;

   if (CalcTableSize(deg(f), k + l - 1) > GF2EXFileThresh)
      use_files = 1;
   else
      use_files = 0;

   GenerateBabySteps(h1, f, h, k, verbose);

   GenerateGiantSteps(f, h1, l, verbose);

   vec_pair_GF2EX_long u;
   GiantRefine(u, f, k, l, verbose);
   BabyRefine(factors, u, k, l, verbose);

   FileCleanup(k, l);
}

long IterComputeDegree(const GF2EX& h, const GF2EXModulus& F)
{
   long n = deg(F);

   if (n == 1 || IsX(h)) return 1;

   long B = n/2;
   long k = SqrRoot(B);
   long l = (B+k-1)/k;

   GF2EXArgument H;

#if 0
   double n2 = sqrt(double(n));
   double n4 = sqrt(n2);
   double n34 = n2*n4;
   long sz = long(ceil(n34/sqrt(sqrt(2.0))));
#else
   long sz = 2*SqrRoot(F.n);
#endif

   build(H, h, F, sz);

   GF2EX h1;
   h1 = h;

   vec_GF2EX baby;
   baby.SetLength(k);

   SetX(baby[0]);

   long i;

   for (i = 1; i <= k-1; i++) {
      baby[i] = h1;
      CompMod(h1, h1, H, F);
      if (IsX(h1)) return i+1;
   }

   build(H, h1, F, sz);

   long j;

   for (j = 2; j <= l; j++) {
      CompMod(h1, h1, H, F);

      for (i = k-1; i >= 0; i--) {
         if (h1 == baby[i])
            return j*k-i;
      }
   }

   return n;
}

NTL_END_IMPL


#include <Eigen/Core>
#include <Eigen/Geometry>
#include <iostream>
#include <memory>
#include <unsupported/Eigen/AutoDiff>
#include <unsupported/Eigen/LevenbergMarquardt>
/* #include <unsupported/Eigen/NonLinearOptimization> */
#include <vector>

#include <fmt/format.h>
#include <fmt/printf.h>

#include "nanoflann.hpp"

using Vector3f = Eigen::Vector3f;
using Transform = Eigen::Isometry3f;
using Cloud = Eigen::MatrixX3f;

Cloud TransformPointCloud(const Cloud& source_cloud,
    const Transform& transform)
{
    // TODO(yycho0108): implement
    Cloud out(source_cloud.rows(), source_cloud.cols());
    for (int i = 0; i < source_cloud.rows(); ++i) {
        out.row(i).transpose() = transform * source_cloud.row(i).transpose();
    }
    return out;
}

//template <typename T>
//struct TXTYTZRZParametrization {
//    static const std::size_t Dimensions = 4;
//    Eigen::Matrix<T, 3, 1> ApplyTransform(const Eigen::Matrix<T, 3, 1>& point);
//    void FromTransform(const Transform&);
//    Transform ToTransform();
//};

static constexpr const int kParamSize = 4;
template <typename T>
struct ParametrizedIsometryCost : public Eigen::DenseFunctor<T> {
    ParametrizedIsometryCost(const Cloud& source, const Cloud& target)
        : Eigen::DenseFunctor<T>(kParamSize, source.size())
        , source_(source)
        , target_(target)
    {
    }

    template <typename T1>
    int operator()(const Eigen::Matrix<T1, Eigen::Dynamic, 1>& parameter,
        Eigen::Matrix<T1, Eigen::Dynamic, 1>& fvec) const
    {
        const Eigen::Matrix<T1, 3, 1>& translation = parameter.template head<3>();
        //const Eigen::Matrix<T1, 3, 1>& rvec = parameter.template tail<3>();
        //const T1 angle = rvec.norm();
        //const Eigen::AngleAxis<T1>& rotation{angle, rvec / angle};
        //Eigen::Quaternion<T1> rotation;
        const Eigen::AngleAxis<T1>& rotation{ parameter(3), Eigen::Matrix<T1, 3, 1>{ 0, 0, 1 } };

        for (int i = 0; i < source_.rows(); ++i) {
            //fvec.segment(i * 3, 3) = source_.row(i).template cast<T1>() + translation.transpose() - target_.row(i).template cast<T1>();
            fvec.segment(i * 3, 3) = (rotation * source_.row(i).template cast<T1>().transpose()).transpose() + translation.transpose() - target_.row(i).template cast<T1>();
        }
        return 0;
        //// Set to simpler version for now.
        // Eigen::Transform<T1, 3, Eigen::Isometry> transform;
        // transform.translation() = translation;

        //// Will this work?
        // std::cout << "here1" << std::endl;
        // auto transformed = (transform.linear() * source_.transpose().template
        // cast<T1>()).transpose();
        ///* transformed.rowwise() += transform.translation().transpose(); */
        // std::cout << "here2" << std::endl;
        // Eigen::Matrix<T1, Eigen::Dynamic, 3> delta = (transformed -
        // target_).template cast<T1>(); std::cout << "here2.5" << std::endl; if
        // (!(delta.Flags & Eigen::RowMajorBit)) {
        //    // 3xN memory layout -> Nx3 memory layout
        //    Eigen::Matrix<T1, 3, Eigen::Dynamic> data = delta.transpose();
        //    Eigen::Map<Eigen::Matrix<T1, Eigen::Dynamic, 1>> rhs(data.data(),
        //        fvec.rows(), 1);
        //    fvec.segment(0, fvec.rows()) = rhs;
        //} else {
        //    Eigen::Map<Eigen::Matrix<T1, Eigen::Dynamic, 1>> rhs(delta.data(),
        //        fvec.rows(), 1);
        //    fvec.segment(0, fvec.rows()) = rhs;
        //}
        ////std::cout << "here3" << std::endl;
        ////std::cout << delta.rows() << 'x' << delta.cols() << std::endl;
        ////std::cout << fvec.rows() << 'x' << fvec.cols() << std::endl;
        ////std::cout << rhs.rows() << 'x' << rhs.cols() << std::endl;
        ////fvec = rhs;
        // return 0;
    };

    int df(const Eigen::Matrix<float, kParamSize, 1>& parameter,
        Eigen::MatrixXf& jac) const
    {
        using Scalar = Eigen::AutoDiffScalar<Eigen::VectorXf>;
        using ScalarVector = Eigen::Matrix<Scalar, Eigen::Dynamic, 1>;
        ScalarVector ax = parameter.template cast<Scalar>();
        ScalarVector av(this->values());

        for (int j = 0; j < this->values(); ++j) {
            av[j].derivatives().resize(this->inputs());
        }
        for (int i = 0; i < this->inputs(); ++i) {
            ax[i].derivatives().resize(this->inputs());
            ax[i].derivatives().setZero();
            ax[i].derivatives()(i) = 1.0;
            //ax[i].derivatives() = Eigen::Vector3f::Unit(this->inputs(), i);
            //ax[i].derivatives().resize(this->inputs());
        }

        operator()(ax, av);

        // jac = (48x3)
        for (int i = 0; i < this->values(); ++i) {
            // std::cout << "==" << std::endl;
            // std::cout << jac.row(i).rows() << std::endl;           // 1
            // std::cout << jac.row(i).cols() << std::endl;           // 3
            // std::cout << av[i].derivatives().rows() << std::endl;  // 0
            // std::cout << av[i].derivatives().cols() << std::endl;  // 1
            jac.row(i) = av[i].derivatives();
        }
    }
    const Cloud& source_;
    const Cloud& target_;
    /* IsometryCost<T> cost; */
};

/**
 * Generalized ICP implementation, based on
 * "Generalized-ICP" Segal et al.
 * www.roboticsproceedings.org/rss05/p21.pdf
 */
class GeneralizedICP {
public:
    using KDTree = nanoflann::KDTreeEigenMatrixAdaptor<Cloud>;

    GeneralizedICP()
        : source(nullptr)
        , target(nullptr)
    {
    }
    void SetSourceCloud(const Cloud& source_cloud)
    {
        source = &source_cloud;
        /* source_tree = std::make_shared<KDTree>(3, std::cref(source_cloud),
         * 10); */
        /* source_tree->index->buildIndex(); */
    }
    void SetTargetCloud(const Cloud& target_cloud)
    {
        target = &target_cloud;
        target_tree = std::make_shared<KDTree>(3, std::cref(target_cloud), 10);
        target_tree->index->buildIndex();
    }

    void FindCorrespondences(const Cloud& source,
        std::vector<std::size_t>* indices,
        std::vector<float>* squared_distances)
    {
        // Reset memory.
        indices->clear();
        squared_distances->clear();
        indices->reserve(source.rows());
        squared_distances->reserve(source.rows());

        // Temps
        std::vector<long> knn_indices(5);
        std::vector<float> knn_squared_distances(5);

        for (int i = 0; i < source.rows(); ++i) {
            const Vector3f& query_point = source.row(i);
            /* fmt::print("query {}\n", query_point.transpose()); */
            /* const float* check = reinterpret_cast<const float*>(query_point.data()); */
            /* fmt::print("check {} {} {}\n", check[0], check[1], check[2]); */
            bool suc = target_tree->index->knnSearch(reinterpret_cast<const float*>(query_point.data()), 1,
                &knn_indices[0], &knn_squared_distances[0]);
            if (!suc) {
                fmt::print("unsuccessful");
                exit(1);
            }
            /* fmt::print("{},", knn_indices[0]); */

            indices->emplace_back(knn_indices[0]);
            squared_distances->emplace_back(knn_squared_distances[0]);
        }
    }

    float TransformCost(const Cloud& source, const Cloud& target,
        const Transform& transform)
    {
        // Nx3 = (3x3?) * (3xN)
        const Cloud& target_ = (transform * source.transpose()).transpose();
        const Cloud& delta = target - target_;

        /* delta = b - Ta */
        const Eigen::Matrix3f& R = transform.rotation();

        float total_cost{ 0 };
        for (int i = 0; i < source.rows(); ++i) {
            const Vector3f& d = delta.row(i);
            const float cost = d.transpose() * R * Eigen::Matrix3f::Identity() * R.transpose() * d;
            total_cost += cost;
        }
        return total_cost;
    }

    Transform OptimizeTransform(const Cloud& a, const Cloud& b, const Transform& seed = Transform::Identity())
    {
        // Initialize system.
        ParametrizedIsometryCost<float> cost_fun{ a, b };
        Eigen::LevenbergMarquardt<ParametrizedIsometryCost<float>> lm{ cost_fun };

        // Apply parametrization.
        Eigen::Matrix<float, -1, 1> param(4, 1);
        param.head<3>() = seed.translation();
        param(3) = Eigen::AngleAxis<float>{ seed.rotation() }.angle();

        // Run Optimization.
        lm.minimize(param);

        // Extract output and return.
        Transform out = Transform::Identity();
        out.translation() = param.head<3>();
        out.linear() = Eigen::AngleAxis<float>{ param(3), Eigen::Vector3f{ 0, 0, 1 } }.toRotationMatrix();
        return out;
    }

    Transform ComputeTransform()
    {
        static constexpr const float kTransformEpsilon = 0.001;
        static constexpr const float kRotationEpsilon = 0.001;
        static constexpr const float kMaxIterations = 100;

        Cloud tmp = *source;
        Transform transform = Transform::Identity();
        std::vector<std::size_t> target_indices;
        std::vector<float> squared_distances;

        bool converged = false;
        for (int count = 0; count < kMaxIterations; ++count) {
            FindCorrespondences(tmp, &target_indices, &squared_distances);
            for (std::size_t i = 0; i < source->rows(); ++i) {
                const std::size_t& source_index = i;
                const std::size_t& target_index = target_indices[i];
                const float& squared_distance = squared_distances[i];
                //fmt::print("{} {} {}\n", source_index, target_index,
                //    squared_distance);
            }

            // Make a copy for now. FIXME(yycho0108): better indexing scheme
            // const Cloud& X = target(target_indices, Eigen::internal::all);
            Cloud tmp_target(tmp.rows(), tmp.cols());
            for (std::size_t i = 0; i < source->rows(); ++i) {
                tmp_target.row(i) = target->row(target_indices[i]);
            }

            Transform target_from_tmp = OptimizeTransform(tmp, tmp_target);
            if (target_from_tmp.translation().squaredNorm() < kTransformEpsilon
                && std::abs(std::acos(0.5 * (target_from_tmp.linear().trace() - 1.0))) < kRotationEpsilon) {
                converged = true;
                break;
            }
            tmp = TransformPointCloud(tmp, target_from_tmp);
            transform = target_from_tmp * transform;

            ++count;
        }
        fmt::print("Converged : {}\n", bool(converged));
        return transform;
    }

    const Cloud *source, *target;
    /* std::shared_ptr<KDTree> source_tree; */
    std::shared_ptr<KDTree> target_tree;
};

int main()
{
    Transform ground_truth_transform = Transform::Identity();
    ground_truth_transform.translation() = Eigen::Vector3f{ 0.2, 0.2, 0.3 };
    ground_truth_transform.linear() = Eigen::AngleAxisf(0.1, Eigen::Vector3f{ 0., 0., 1. }).toRotationMatrix();

    Cloud source_cloud = Cloud::Zero(16, 3);
    source_cloud.setRandom();

    Cloud target_cloud = Cloud::Zero(17, 3);
    //target_cloud.block<16, 3>(0, 0) = source_cloud;
    target_cloud.block<16, 3>(0, 0) = TransformPointCloud(source_cloud, ground_truth_transform);

    /* fmt::print("{} vs {}", source_cloud.row(0), target_cloud.row(0)); */

    GeneralizedICP gicp;
    gicp.SetSourceCloud(source_cloud);
    gicp.SetTargetCloud(target_cloud);
    const Transform& estimated_transform = gicp.ComputeTransform();

    fmt::print("ground truth : \n{}\n", ground_truth_transform.matrix());
    fmt::print("computed : \n{}\n", estimated_transform.matrix());
    return 0;
}


#include <iostream>
#include <random>
#include <Eigen/Geometry>
#include <memory>

#include "ransac.h"

using namespace std;

Ransac::Ransac(int n_iterations, double threshold): n_iterations(n_iterations),
                                                    threshold(threshold)
{    // Seed the generator
    generator.seed(rd());
};

double Ransac::mahalanobis_distance(const point3d& X, const point3d& A, const point3d& B,
                                    const Eigen::Vector3d& D, const Eigen::Matrix3d& U)
{
    // Apply affine transform to points
    Eigen::Vector3d A_affine = ((D.array().inverse()).sqrt()).matrix().asDiagonal() * U.transpose() * (A - X);
    Eigen::Vector3d B_affine = ((D.array().inverse()).sqrt()).matrix().asDiagonal() * U.transpose() * (B - X);

    double distance = A_affine.cross(B_affine).norm()/(A_affine - B_affine).norm();

    return distance;
}

Eigen::Matrix3Xd Ransac::removeOutlierPoints(const Eigen::Matrix3Xd& linepoints, const std::vector<Eigen::Vector3d>& eig_val,
                                             const std::vector<Eigen::Matrix3d>& eig_vector, std::vector<Eigen::Matrix3d>& updated_covariance,
                                             std::vector<Eigen::Matrix3d>& updated_inv_root_covariance, const std::vector<Eigen::Matrix3d>& cov_G,
                                             const std::vector<Eigen::Matrix3d>& inv_cov_G)
{
    Eigen::Matrix3Xd best_inliers;
    int best_num_inliers = -1;
    float distance;

    // solace in not copying this over and over
    std::shared_ptr<vector<int>> ptr_to_best_inlier_indices = NULL;

    // Setup the range from which we choose random numbers
    std::uniform_int_distribution<unsigned> distribution(0,linepoints.cols()-1); // end inclusive

    for(int iter=0; iter < n_iterations; ++iter){
        int inliers = 0;
        auto ptr_to_inlier_indices = std::make_shared<std::vector<int>>();

        // re-generate i and j randomly until they are not equal
        int i = distribution(generator), j = distribution(generator);
        while(i == j){
            i = distribution(generator), j = distribution(generator);
        }

        Eigen::Matrix3Xd support_points(3, linepoints.cols());
        Eigen::Vector3d A = linepoints.col(i), B = linepoints.col(j);

        // calculate support for chosen endpoint indices
        for(int k=0; k < linepoints.cols(); ++k){
            distance = mahalanobis_distance(linepoints.col(k), A, B, eig_val[k], eig_vector[k]);
            // cout << distance << endl;
            if (distance < threshold){
                support_points.col(inliers++) = linepoints.col(k);
                ptr_to_inlier_indices->push_back(k);
            }

        }

        if(inliers > best_num_inliers){
            best_num_inliers = inliers;
            best_inliers = support_points.leftCols(inliers);
            ptr_to_best_inlier_indices = ptr_to_inlier_indices;
        }
    }

    // Take the indices of new covariances and populate updated covariances
    // Very slow - <To-do> refactor to deal with indices/pointers instead of copying over data
    for(auto i: *ptr_to_best_inlier_indices){
        updated_covariance.push_back(cov_G[i]);
        updated_inv_root_covariance.push_back(inv_cov_G[i]);
    }

    return best_inliers;
}


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

  @copyright 2015 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_ARCH_COMMON_SCALAR_FUNCTION_ACOTPI_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_SCALAR_FUNCTION_ACOTPI_HPP_INCLUDED

#ifndef BOOST_SIMD_NO_INFINITIES
#include <boost/simd/function/is_inf.hpp>
#endif
#include <boost/simd/constant/half.hpp>
#include <boost/simd/function/abs.hpp>
#include <boost/simd/function/atanpi.hpp>
#include <boost/simd/function/bitofsign.hpp>
#include <boost/simd/function/bitwise_or.hpp>
#include <boost/simd/function/if_else_zero.hpp>
#include <boost/simd/function/if_zero_else.hpp>
#include <boost/simd/function/is_nez.hpp>
#include <boost/simd/function/minus.hpp>
#include <boost/simd/detail/dispatch/function/overload.hpp>
#include <boost/config.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;
  BOOST_DISPATCH_OVERLOAD ( acotpi_
                          , (typename A0)
                          , bd::cpu_
                          , bd::scalar_<bd::floating_<A0> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() ( A0 a0) const BOOST_NOEXCEPT
    {
      A0 z =Half<A0>()-if_else_zero(is_nez(a0),atanpi(bs::abs(a0)));
      #ifndef BOOST_SIMD_NO_INFINITIES
      z = if_zero_else(is_inf(a0),z);
      #endif
      return bitwise_or(z, bitofsign(a0));
    }
  };
} } }

#endif


#pragma once

// Do not parallelise the GEMM routines as this is called by multiple threads
#ifndef NEON_PARALLEL_EIGEN_SOLVERS
#define EIGEN_DONT_PARALLELIZE
#endif

#include <Eigen/Core>
#include <Eigen/LU>

/// \file dense_matrix.hpp

namespace neon
{
/// Matrix in row major layout
using matrix = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
/// Alias to row major layout matrix
using row_matrix = matrix;
/// Matrix in column major layout
using col_matrix = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::ColMajor>;

/// 2x2 fixed size matrix
using matrix2 = Eigen::Matrix<double, 2, 2>;
/// 3x3 fixed size matrix
using matrix3 = Eigen::Matrix<double, 3, 3>;
/// 6x6 fixed size matrix
using matrix6 = Eigen::Matrix<double, 6, 6>;
/// 9x9 fixed size matrix
using matrix9 = Eigen::Matrix<double, 9, 9>;
/// 12x12 fixed size matrix
using matrix12 = Eigen::Matrix<double, 12, 12>;
/// 16x16 fixed size matrix
using matrix16 = Eigen::Matrix<double, 16, 16>;
/// 3x1 fixed size matrix for non-square Jacobians
using matrix31 = Eigen::Matrix<double, 3, 1>;
/// 3x2 fixed size matrix for non-square Jacobians
using matrix32 = Eigen::Matrix<double, 3, 2>;

template <int geometric_dimension>
using matrixxd = Eigen::Matrix<double, Eigen::Dynamic, geometric_dimension>;

template <int geometric_dimension>
using matrixdx = Eigen::Matrix<double, geometric_dimension, Eigen::Dynamic>;

/// Compile time fixed rows for nodal coordinates in two dimensions
using matrix2x = Eigen::Matrix<double, 2, Eigen::Dynamic>;
/// Compile time fixed rows for nodal coordinates in three dimensions
using matrix3x = Eigen::Matrix<double, 3, Eigen::Dynamic>;

/// Fixed size vector of variable length
using vector = Eigen::Matrix<double, Eigen::Dynamic, 1>;
/// Fixed size vector of length two
using vector2 = Eigen::Vector2d;
/// Fixed size vector of length three
using vector3 = Eigen::Vector3d;
/// Fixed size vector of length four
using vector4 = Eigen::Vector4d;
/// Fixed size vector of length five
using vector5 = Eigen::Matrix<double, 5, 1>;
/// Fixed size vector of length six
using vector6 = Eigen::Matrix<double, 6, 1>;
/// Fixed size vector of length sixteen
using vector16 = Eigen::Matrix<double, 16, 1>;
}


//
// Created by david on 2019-10-05.
//
#include <Eigen/Core>
#include <iostream>
int main(){
    Eigen::MatrixXd test(10,10);
    test.setRandom();
    std::cout << test << std::endl;
    return 0;
}


#include <cstdlib>
#include <iostream>
#include <complex>
#define BOOST_UBLAS_NO_ELEMENT_PROXIES
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/symmetric.hpp>
#include <boost/numeric/bindings/ublas/vector.hpp>
#include <boost/numeric/bindings/ublas/symmetric.hpp>
#include <boost/numeric/bindings/blas/level2.hpp>
#include "print.hpp"
#include "random.hpp"

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

int main(int argc, char *argv[]) {
  {
    typedef ublas::vector<double> vector;
    typedef ublas::symmetric_matrix<double, ublas::lower, ublas::column_major> matrix_l;
    typedef ublas::symmetric_matrix<double, ublas::upper, ublas::column_major> matrix_u;
    typedef typename vector::size_type size_type;
    rand_normal<double>::reset();
    size_type n=8;
    matrix_l A_l(n);
    matrix_u A_u(n);
    for (size_type j=0; j<n; ++j) {
      A_u(j, j)=rand_normal<double>::get();
      A_l(j, j)=A_u(j, j);
      for (size_type i=0; i<j; ++i) {
    	A_u(i, j)=rand_normal<double>::get();
	A_l(j, i)=A_u(i, j);
       }
    }
    vector x(n);
    for (size_type i=0; i<n; ++i)
      x(i)=rand_normal<double>::get();
    vector y(n);
    for (size_type i=0; i<n; ++i)
      y(i)=rand_normal<double>::get();
    double alpha(rand_normal<double>::get());
    double beta(rand_normal<double>::get());
    vector y1(alpha*ublas::prod(A_l, x)+beta*y);
    vector y2(y);
    blas::spmv(alpha, A_l, x, beta, y2);
    vector y3(y);
    blas::spmv(alpha, A_u, x, beta, y3);
    vector y4(y);
    blas::spmv(alpha, A_l, x, beta, y4);
    vector y5(y);
    blas::spmv(alpha, A_u, x, beta, y5);
    std::cout << "testing boost::ublas containers\n"
    	      << "using ublas            : " << print_vec(y1) << '\n'
    	      << "using blas spmv (lower): " << print_vec(y2) << '\n'
    	      << "using blas spmv (upper): " << print_vec(y3) << '\n'
    	      << "using blas hpmv (lower): " << print_vec(y4) << '\n'
    	      << "using blas hpmv (upper): " << print_vec(y5) << '\n'
    	      << '\n';
  }
  return EXIT_SUCCESS;
}


/*
 * =====================================================================================
 *
 *       Filename:  say.cpp
 *
 *    Description:
 *
 *        Version:  1.0
 *        Created:  13.11.2015 15:11:23
 *       Revision:  none
 *       Compiler:  gcc
 *
 *
 * =====================================================================================
 */
#include "say.h"
#include <cmath>
#include <boost/algorithm/string.hpp>
#include <stdexcept>

namespace say
{
	std::string oneNames(int num)
	{
		switch (num)
		{
			case 1: return "one";
			case 2: return "two";
			case 3: return "three";
			case 4: return "four";
			case 5: return "five";
			case 6: return "six";
			case 7: return "seven";
			case 8: return "eight";
			case 9: return "nine";
			case 10: return "ten";
			case 11: return "eleven";
			case 12: return "twelve";
			default: return "";

		}
	}
	std::string names6099(long int num)
	{
		std::string engl{oneNames(num/10) + "ty"};
		if( (num % 10) != 0)
		{
			engl += "-" + oneNames(num %10);
		}
		return engl;

	}

	std::string tenNames(long int num)
	{
		switch(num)
		{
			case 1 ... 12: return oneNames(num);
			case 13 ... 19: return oneNames((num % 10)) + "teen";
			case 20: return "twenty";
			case 21 ... 29: return "twenty-" + oneNames(num%10);
			case 30: return "thirty";
			case 31 ... 39: return "thirty-" + oneNames(num%10);
			case 40: return "forty";
			case 41 ... 49: return "forty-" + oneNames(num%10);
			case 50: return "fifty";
			case 51 ... 59: return "fifty-" + oneNames(num%10);
			case 60 ... 79: return names6099(num);
			case 80: return "eighty";
			case 81 ... 89: return "eighty-" + oneNames(num%10);
			case 90 ... 99: return names6099(num);
			default: return "";
		}
	}

	uint_fast8_t determineLength(long int num)
	{
		return static_cast<uint_fast8_t>(std::to_string(num).size());
	}

	uint_fast8_t rest(uint_fast8_t l)
	{
		if(l <= 2)
		{
			return 0;
		}
		else
		{
			return (l - (l%3));
		}
	}

	std::string addUnit( std::string const& unit, long int& num)
	{
		std::string engl{};
		uint_fast8_t l{ determineLength(num) };
		if((l%3)==0)
		{
			engl += " " + oneNames( num / pow(10, l-1)) + " hundred";
			num %= static_cast<long int>(pow(10, l-1));
			--l;
		}

		if( (l- determineLength(num)) <=1)
		{
			l = determineLength(num);
			engl += " " + tenNames( num / static_cast<long int>(pow(10, rest(l))));
			num %= static_cast<long int>(pow(10, rest(l)));
		}
		return engl + " " + unit;
	}

	std::string in_english(long int number)
	{
		if( (number < 0) or (number >= 1000ULL*1000ULL*1000ULL*1000ULL))
		{
			throw std::domain_error("Number out of range");
		}
		if( number == 0)
		{
			return "zero";
		}
		uint_fast8_t numberlength{determineLength(number)};
		std::string english{};
		if(numberlength > 9)
		{
			english += addUnit("billion", number);
			numberlength = determineLength(number);
		}
		if(numberlength > 6)
		{
			english += addUnit("million", number);
			numberlength = determineLength(number);
		}
		if(numberlength > 3)
		{
			english += addUnit("thousand", number);
			numberlength = determineLength(number);
		}

		if(numberlength > 0)
		{
			english += addUnit("", number);
		}
		boost::algorithm::trim(english);
		return english;
	}
}


/**
 * @file decision_tree_test.cpp
 * @author Ryan Curtin
 *
 * Tests for the DecisionTree class and related classes.
 *
 * mlpack is free software; you may redistribute it and/or modify it under the
 * terms of the 3-clause BSD license.  You should have received a copy of the
 * 3-clause BSD license along with mlpack.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more information.
 */
#include <mlpack/core.hpp>
#include <mlpack/methods/decision_tree/decision_tree.hpp>
#include <mlpack/methods/decision_tree/information_gain.hpp>
#include <mlpack/methods/decision_tree/gini_gain.hpp>
#include <mlpack/methods/decision_tree/random_dimension_select.hpp>
#include <mlpack/methods/decision_tree/multiple_random_dimension_select.hpp>

#include <boost/test/unit_test.hpp>
#include "test_tools.hpp"
#include "serialization.hpp"
#include "mock_categorical_data.hpp"

using namespace mlpack;
using namespace mlpack::tree;
using namespace mlpack::distribution;

BOOST_AUTO_TEST_SUITE(DecisionTreeTest);

/**
 * Make sure the Gini gain is zero when the labels are perfect.
 */
BOOST_AUTO_TEST_CASE(GiniGainPerfectTest)
{
  arma::rowvec weights(10, arma::fill::ones);
  arma::Row<size_t> labels;
  labels.zeros(10);

  // Test that it's perfect regardless of number of classes.
  for (size_t c = 1; c < 10; ++c)
    BOOST_REQUIRE_SMALL(GiniGain::Evaluate<false>(labels, c, weights), 1e-5);
}

/**
 * Make sure the Gini gain is -0.5 when the class split between two classes
 * is even.
 */
BOOST_AUTO_TEST_CASE(GiniGainEvenSplitTest)
{
  arma::rowvec weights = arma::ones<arma::rowvec>(10);
  arma::Row<size_t> labels(10);
  for (size_t i = 0; i < 5; ++i)
    labels[i] = 0;
  for (size_t i = 5; i < 10; ++i)
    labels[i] = 1;

  // Test that it's -0.5 regardless of the number of classes.
  for (size_t c = 2; c < 10; ++c)
  {
    BOOST_REQUIRE_CLOSE(
        GiniGain::Evaluate<false>(labels, c, weights), -0.5, 1e-5);
    double weightedGain = GiniGain::Evaluate<true>(labels, c, weights);

    // The weighted gain should stay the same with unweight one
    BOOST_REQUIRE_EQUAL(
        GiniGain::Evaluate<false>(labels, c, weights), weightedGain);
  }
}

/**
 * The Gini gain of an empty vector is 0.
 */
BOOST_AUTO_TEST_CASE(GiniGainEmptyTest)
{
  arma::rowvec weights = arma::ones<arma::rowvec>(10);
  // Test across some numbers of classes.
  arma::Row<size_t> labels;
  for (size_t c = 1; c < 10; ++c)
    BOOST_REQUIRE_SMALL(GiniGain::Evaluate<false>(labels, c, weights), 1e-5);

  for (size_t c = 1; c < 10; ++c)
    BOOST_REQUIRE_SMALL(GiniGain::Evaluate<true>(labels, c, weights), 1e-5);
}

/**
 * The Gini gain is -(1 - 1/k) for k classes evenly split.
 */
BOOST_AUTO_TEST_CASE(GiniGainEvenSplitManyClassTest)
{
  // Try with many different classes.
  for (size_t c = 2; c < 30; ++c)
  {
    arma::Row<size_t> labels(c);
    arma::rowvec weights(c);
    for (size_t i = 0; i < c; ++i)
    {
      labels[i] = i;
      weights[i] = 1;
    }

    // Calculate Gini gain and make sure it is correct.
    BOOST_REQUIRE_CLOSE(GiniGain::Evaluate<false>(labels, c, weights),
        -(1.0 - 1.0 / c), 1e-5);
    BOOST_REQUIRE_CLOSE(GiniGain::Evaluate<true>(labels, c, weights),
        -(1.0 - 1.0 / c), 1e-5);
  }
}

/**
 * The Gini gain should not be sensitive to the number of points.
 */
BOOST_AUTO_TEST_CASE(GiniGainManyPoints)
{
  for (size_t i = 1; i < 20; ++i)
  {
    const size_t numPoints = 100 * i;
    arma::rowvec weights(numPoints);
    weights.ones();
    arma::Row<size_t> labels(numPoints);
    for (size_t j = 0; j < numPoints / 2; ++j)
      labels[j] = 0;
    for (size_t j = numPoints / 2; j < numPoints; ++j)
      labels[j] = 1;

    BOOST_REQUIRE_CLOSE(GiniGain::Evaluate<false>(labels, 2, weights), -0.5,
        1e-5);
    BOOST_REQUIRE_CLOSE(GiniGain::Evaluate<true>(labels, 2, weights), -0.5,
        1e-5);
  }
}

/**
 * To make sure the Gini gain can been cacluate proporately with weight.
 */
BOOST_AUTO_TEST_CASE(GiniGainWithWeight)
{
  arma::Row<size_t> labels(10);
  arma::rowvec weights(10);
  for (size_t i = 0; i < 5; ++i)
  {
    labels[i] = 0;
    weights[i] = 0.3;
  }
  for (size_t i = 5; i < 10; ++i)
  {
    labels[i] = 1;
    weights[i] = 0.7;
  }

  BOOST_REQUIRE_CLOSE(
      GiniGain::Evaluate<true>(labels, 2, weights), -0.42, 1e-5);
}

/**
 * The information gain should be zero when the labels are perfect.
 */
BOOST_AUTO_TEST_CASE(InformationGainPerfectTest)
{
  arma::rowvec weights;
  arma::Row<size_t> labels;
  labels.zeros(10);

  // Test that it's perfect regardless of number of classes.
  for (size_t c = 1; c < 10; ++c)
  {
    BOOST_REQUIRE_SMALL(
        InformationGain::Evaluate<false>(labels, c, weights), 1e-5);
  }
}

/**
 * If we have an even split, the information gain should be -1.
 */
BOOST_AUTO_TEST_CASE(InformationGainEvenSplitTest)
{
  arma::Row<size_t> labels(10);
  arma::rowvec weights(10);
  weights.ones();
  for (size_t i = 0; i < 5; ++i)
    labels[i] = 0;
  for (size_t i = 5; i < 10; ++i)
    labels[i] = 1;

  // Test that it's -1 regardless of the number of classes.
  for (size_t c = 2; c < 10; ++c)
  {
    // Weighted and unweighted result should be the same.
    BOOST_REQUIRE_CLOSE(InformationGain::Evaluate<false>(labels, c, weights),
        -1.0, 1e-5);
    BOOST_REQUIRE_CLOSE(InformationGain::Evaluate<true>(labels, c, weights),
        -1.0, 1e-5);
  }
}

/**
 * The information gain of an empty vector is 0.
 */
BOOST_AUTO_TEST_CASE(InformationGainEmptyTest)
{
  arma::Row<size_t> labels;
  arma::rowvec weights = arma::ones<arma::rowvec>(10);
  for (size_t c = 1; c < 10; ++c)
  {
    BOOST_REQUIRE_SMALL(InformationGain::Evaluate<false>(labels, c, weights),
        1e-5);
    BOOST_REQUIRE_SMALL(InformationGain::Evaluate<true>(labels, c, weights),
        1e-5);
  }
}

/**
 * The information gain is log2(1/k) when splitting equal classes.
 */
BOOST_AUTO_TEST_CASE(InformationGainEvenSplitManyClassTest)
{
  arma::rowvec weights;
  // Try with many different numbers of classes.
  for (size_t c = 2; c < 30; ++c)
  {
    arma::Row<size_t> labels(c);
    for (size_t i = 0; i < c; ++i)
      labels[i] = i;

    // Calculate information gain and make sure it is correct.
    BOOST_REQUIRE_CLOSE(InformationGain::Evaluate<false>(labels, c, weights),
        std::log2(1.0 / c), 1e-5);
  }
}

/**
 * Test the information gain with weighted labels
 */
BOOST_AUTO_TEST_CASE(InformationWithWeight)
{
  arma::Row<size_t> labels(10);
  arma::rowvec weights("1 1 1 1 1 0 0 0 0 0");
  for (size_t i = 0; i < 5; ++i)
    labels[i] = 0;
  for (size_t i = 5; i < 10; ++i)
    labels[i] = 1;

  // Zero is not a good result as gain, but we just need to prove
  // cacluation works.
  BOOST_REQUIRE_CLOSE(
      InformationGain::Evaluate<true>(labels, 2, weights), 0, 1e-5);
}

/**
 * The information gain should not be sensitive to the number of points.
 */
BOOST_AUTO_TEST_CASE(InformationGainManyPoints)
{
  for (size_t i = 1; i < 20; ++i)
  {
    const size_t numPoints = 100 * i;
    arma::Row<size_t> labels(numPoints);
    arma::rowvec weights = arma::ones<arma::rowvec>(numPoints);
    for (size_t j = 0; j < numPoints / 2; ++j)
      labels[j] = 0;
    for (size_t j = numPoints / 2; j < numPoints; ++j)
      labels[j] = 1;

    BOOST_REQUIRE_CLOSE(InformationGain::Evaluate<false>(labels, 2, weights),
        -1.0, 1e-5);
    // It should make no difference between a weighted and unweighted
    // calculation.
    BOOST_REQUIRE_CLOSE(InformationGain::Evaluate<true>(labels, 2, weights),
        -1.0, 1e-5);
  }
}

/**
 * Check that the BestBinaryNumericSplit will split on an obviously splittable
 * dimension.
 */
BOOST_AUTO_TEST_CASE(BestBinaryNumericSplitSimpleSplitTest)
{
  arma::vec values("0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0");
  arma::Row<size_t> labels("0 0 0 0 0 1 1 1 1 1 1");
  arma::rowvec weights(labels.n_elem);
  weights.ones();

  arma::vec classProbabilities;
  BestBinaryNumericSplit<GiniGain>::template AuxiliarySplitInfo<double> aux;

  // Call the method to do the splitting.
  const double bestGain = GiniGain::Evaluate<false>(labels, 2, weights);
  const double gain = BestBinaryNumericSplit<GiniGain>::SplitIfBetter<false>(
      bestGain, values, labels, 2, weights, 3, 1e-7, classProbabilities,
      aux);
  const double weightedGain =
      BestBinaryNumericSplit<GiniGain>::SplitIfBetter<true>(bestGain, values,
      labels, 2, weights, 3, 1e-7, classProbabilities, aux);

  // Make sure that a split was made.
  BOOST_REQUIRE_GT(gain, bestGain);

  // Make sure weight works and is not different than the unweighted one.
  BOOST_REQUIRE_EQUAL(gain, weightedGain);

  // The split is perfect, so we should be able to accomplish a gain of 0.
  BOOST_REQUIRE_SMALL(gain, 1e-5);

  // The class probabilities, for this split, hold the splitting point, which
  // should be between 4 and 5.
  BOOST_REQUIRE_EQUAL(classProbabilities.n_elem, 1);
  BOOST_REQUIRE_GT(classProbabilities[0], 0.4);
  BOOST_REQUIRE_LT(classProbabilities[0], 0.5);
}

/**
 * Check that the BestBinaryNumericSplit won't split if not enough points are
 * given.
 */
BOOST_AUTO_TEST_CASE(BestBinaryNumericSplitMinSamplesTest)
{
  arma::vec values("0.0 0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.0");
  arma::Row<size_t> labels("0 0 0 0 0 1 1 1 1 1 1");
  arma::rowvec weights(labels.n_elem);

  arma::vec classProbabilities;
  BestBinaryNumericSplit<GiniGain>::template AuxiliarySplitInfo<double> aux;

  // Call the method to do the splitting.
  const double bestGain = GiniGain::Evaluate<false>(labels, 2, weights);
  const double gain = BestBinaryNumericSplit<GiniGain>::SplitIfBetter<false>(
      bestGain, values, labels, 2, weights, 8, 1e-7, classProbabilities,
      aux);
  // This should make no difference because it won't split at all.
  const double weightedGain =
      BestBinaryNumericSplit<GiniGain>::SplitIfBetter<true>(bestGain, values,
      labels, 2, weights, 8, 1e-7, classProbabilities, aux);

  // Make sure that no split was made.
  BOOST_REQUIRE_EQUAL(gain, DBL_MAX);
  BOOST_REQUIRE_EQUAL(gain, weightedGain);
  BOOST_REQUIRE_EQUAL(classProbabilities.n_elem, 0);
}

/**
 * Check that the BestBinaryNumericSplit doesn't split a dimension that gives no
 * gain.
 */
BOOST_AUTO_TEST_CASE(BestBinaryNumericSplitNoGainTest)
{
  arma::vec values(100);
  arma::Row<size_t> labels(100);
  arma::rowvec weights;
  for (size_t i = 0; i < 100; i += 2)
  {
    values[i] = i;
    labels[i] = 0;
    values[i + 1] = i;
    labels[i + 1] = 1;
  }

  arma::vec classProbabilities;
  BestBinaryNumericSplit<GiniGain>::template AuxiliarySplitInfo<double> aux;

  // Call the method to do the splitting.
  const double bestGain = GiniGain::Evaluate<false>(labels, 2, weights);
  const double gain = BestBinaryNumericSplit<GiniGain>::SplitIfBetter<false>(
      bestGain, values, labels, 2, weights, 10, 1e-7, classProbabilities,
      aux);

  // Make sure there was no split.
  BOOST_REQUIRE_EQUAL(gain, DBL_MAX);
  BOOST_REQUIRE_EQUAL(classProbabilities.n_elem, 0);
}

/**
 * Check that the AllCategoricalSplit will split when the split is obviously
 * better.
 */
BOOST_AUTO_TEST_CASE(AllCategoricalSplitSimpleSplitTest)
{
  arma::vec values("0 0 0 1 1 1 2 2 2 3 3 3");
  arma::Row<size_t> labels("0 0 0 2 2 2 1 1 1 2 2 2");
  arma::rowvec weights(labels.n_elem);
  weights.ones();

  arma::vec classProbabilities;
  AllCategoricalSplit<GiniGain>::template AuxiliarySplitInfo<double> aux;

  // Call the method to do the splitting.
  const double bestGain = GiniGain::Evaluate<false>(labels, 3, weights);
  const double gain = AllCategoricalSplit<GiniGain>::SplitIfBetter<false>(
      bestGain, values, 4, labels, 3, weights, 3, 1e-7, classProbabilities,
      aux);
  const double weightedGain =
      AllCategoricalSplit<GiniGain>::SplitIfBetter<true>(bestGain, values, 4,
      labels, 3, weights, 3, 1e-7, classProbabilities, aux);

  // Make sure that a split was made.
  BOOST_REQUIRE_GT(gain, bestGain);

  // Since the split is perfect, make sure the new gain is 0.
  BOOST_REQUIRE_SMALL(gain, 1e-5);

  BOOST_REQUIRE_EQUAL(gain, weightedGain);

  // Make sure the class probabilities now hold the number of children.
  BOOST_REQUIRE_EQUAL(classProbabilities.n_elem, 1);
  BOOST_REQUIRE_EQUAL((size_t) classProbabilities[0], 4);
}

/**
 * Make sure that AllCategoricalSplit respects the minimum number of samples
 * required to split.
 */
BOOST_AUTO_TEST_CASE(AllCategoricalSplitMinSamplesTest)
{
  arma::vec values("0 0 0 1 1 1 2 2 2 3 3 3");
  arma::Row<size_t> labels("0 0 0 2 2 2 1 1 1 2 2 2");
  arma::rowvec weights(labels.n_elem);
  weights.ones();

  arma::vec classProbabilities;
  AllCategoricalSplit<GiniGain>::template AuxiliarySplitInfo<double> aux;

  // Call the method to do the splitting.
  const double bestGain = GiniGain::Evaluate<false>(labels, 3, weights);
  const double gain = AllCategoricalSplit<GiniGain>::SplitIfBetter<false>(
      bestGain, values, 4, labels, 3, weights, 4, 1e-7, classProbabilities,
      aux);

  // Make sure it's not split.
  BOOST_REQUIRE_EQUAL(gain, DBL_MAX);
  BOOST_REQUIRE_EQUAL(classProbabilities.n_elem, 0);
}

/**
 * Check that no split is made when it doesn't get us anything.
 */
BOOST_AUTO_TEST_CASE(AllCategoricalSplitNoGainTest)
{
  arma::vec values(300);
  arma::Row<size_t> labels(300);
  arma::rowvec weights = arma::ones<arma::rowvec>(300);

  for (size_t i = 0; i < 300; i += 3)
  {
    values[i] = int(i / 3) % 10;
    labels[i] = 0;
    values[i + 1] = int(i / 3) % 10;
    labels[i + 1] = 1;
    values[i + 2] = int(i / 3) % 10;
    labels[i + 2] = 2;
  }

  arma::vec classProbabilities;
  AllCategoricalSplit<GiniGain>::template AuxiliarySplitInfo<double> aux;

  // Call the method to do the splitting.
  const double bestGain = GiniGain::Evaluate<false>(labels, 3, weights);
  const double gain = AllCategoricalSplit<GiniGain>::SplitIfBetter<false>(
      bestGain, values, 10, labels, 3, weights, 10, 1e-7,
      classProbabilities, aux);
  const double weightedGain =
      AllCategoricalSplit<GiniGain>::SplitIfBetter<true>(bestGain, values, 10,
      labels, 3, weights, 10, 1e-7, classProbabilities, aux);

  // Make sure that there was no split.
  BOOST_REQUIRE_EQUAL(gain, DBL_MAX);
  BOOST_REQUIRE_EQUAL(gain, weightedGain);
  BOOST_REQUIRE_EQUAL(classProbabilities.n_elem, 0);
}

/**
 * A basic construction of the decision tree---ensure that we can create the
 * tree and that it split at least once.
 */
BOOST_AUTO_TEST_CASE(BasicConstructionTest)
{
  arma::mat dataset(10, 100, arma::fill::randu);
  arma::Row<size_t> labels(100);
  for (size_t i = 0; i < 50; ++i)
  {
    dataset(3, i) = 0.0;
    labels[i] = 0;
  }
  for (size_t i = 50; i < 100; ++i)
  {
    dataset(3, i) = 1.0;
    labels[i] = 1;
  }

  // Use default parameters.
  DecisionTree<> d(dataset, labels, 2, 10);

  // Now require that we have some children.
  BOOST_REQUIRE_GT(d.NumChildren(), 0);
}

/**
 * Construct a tree with weighted labels.
 */
BOOST_AUTO_TEST_CASE(BasicConstructionTestWithWeight)
{
  arma::mat dataset(10, 100, arma::fill::randu);
  arma::Row<size_t> labels(100);
  for (size_t i = 0; i < 50; ++i)
  {
    dataset(3, i) = 0.0;
    labels[i] = 0;
  }
  for (size_t i = 50; i < 100; ++i)
  {
    dataset(3, i) = 1.0;
    labels[i] = 1;
  }
  arma::rowvec weights(labels.n_elem);
  weights.ones();

  // Use default parameters.
  DecisionTree<> wd(dataset, labels, 2, weights, 10);
  DecisionTree<> d(dataset, labels, 2, 10);

  // Now require that we have some children.
  BOOST_REQUIRE_GT(wd.NumChildren(), 0);
  BOOST_REQUIRE_EQUAL(wd.NumChildren(), d.NumChildren());
}

/**
 * Construct the decision tree on numeric data only and see that we can fit it
 * exactly and achieve perfect performance on the training set.
 */
BOOST_AUTO_TEST_CASE(PerfectTrainingSet)
{
  arma::mat dataset(10, 100, arma::fill::randu);
  arma::Row<size_t> labels(100);
  for (size_t i = 0; i < 50; ++i)
  {
    dataset(3, i) = 0.0;
    labels[i] = 0;
  }
  for (size_t i = 50; i < 100; ++i)
  {
    dataset(3, i) = 1.0;
    labels[i] = 1;
  }

  DecisionTree<> d(dataset, labels, 2, 1, 0.0); // Minimum leaf size of 1.

  // Make sure that we can get perfect accuracy on the training set.
  for (size_t i = 0; i < 100; ++i)
  {
    size_t prediction;
    arma::vec probabilities;
    d.Classify(dataset.col(i), prediction, probabilities);

    BOOST_REQUIRE_EQUAL(prediction, labels[i]);
    BOOST_REQUIRE_EQUAL(probabilities.n_elem, 2);
    for (size_t j = 0; j < 3; ++j)
    {
      if (labels[i] == j)
        BOOST_REQUIRE_CLOSE(probabilities[j], 1.0, 1e-5);
      else
        BOOST_REQUIRE_SMALL(probabilities[j], 1e-5);
    }
  }
}

/**
 * Construct the decision tree with weighted labels
 */
BOOST_AUTO_TEST_CASE(PerfectTrainingSetWithWeight)
{
  // Completely random dataset with no structure.
  arma::mat dataset(10, 100, arma::fill::randu);
  arma::Row<size_t> labels(100);
  for (size_t i = 0; i < 50; ++i)
  {
    dataset(3, i) = 0.0;
    labels[i] = 0;
  }
  for (size_t i = 50; i < 100; ++i)
  {
    dataset(3, i) = 1.0;
    labels[i] = 1;
  }
  arma::rowvec weights(labels.n_elem);
  weights.ones();

  // Minimum leaf size of 1.
  DecisionTree<> d(dataset, labels, 2, weights, 1, 0.0);

  // This part of code is dupliacte with no weighted one.
  for (size_t i = 0; i < 100; ++i)
  {
    size_t prediction;
    arma::vec probabilities;
    d.Classify(dataset.col(i), prediction, probabilities);

    BOOST_REQUIRE_EQUAL(prediction, labels[i]);
    BOOST_REQUIRE_EQUAL(probabilities.n_elem, 2);
    for (size_t j = 0; j < 3; ++j)
    {
      if (labels[i] == j)
        BOOST_REQUIRE_CLOSE(probabilities[j], 1.0, 1e-5);
      else
        BOOST_REQUIRE_SMALL(probabilities[j], 1e-5);
    }
  }
}

/**
 * Make sure class probabilities are computed correctly in the root node.
 */
BOOST_AUTO_TEST_CASE(ClassProbabilityTest)
{
  arma::mat dataset(5, 100, arma::fill::randu);
  arma::Row<size_t> labels(100);
  for (size_t i = 0; i < 100; i += 2)
  {
    labels[i] = 0;
    labels[i + 1] = 1;
  }

  // Create a decision tree that can't split.
  DecisionTree<> d(dataset, labels, 2, 1000);

  BOOST_REQUIRE_EQUAL(d.NumChildren(), 0);

  // Estimate a point's probabilities.
  arma::vec probabilities;
  size_t prediction;
  d.Classify(dataset.col(0), prediction, probabilities);

  BOOST_REQUIRE_EQUAL(probabilities.n_elem, 2);
  BOOST_REQUIRE_CLOSE(probabilities[0], 0.5, 1e-5);
  BOOST_REQUIRE_CLOSE(probabilities[1], 0.5, 1e-5);
}

/**
 * Test that the decision tree generalizes reasonably.
 */
BOOST_AUTO_TEST_CASE(SimpleGeneralizationTest)
{
  arma::mat inputData;
  if (!data::Load("vc2.csv", inputData))
    BOOST_FAIL("Cannot load test dataset vc2.csv!");

  arma::Row<size_t> labels;
  if (!data::Load("vc2_labels.txt", labels))
    BOOST_FAIL("Cannot load labels for vc2_labels.txt");

  // Initialize an all-ones weight matrix.
  arma::rowvec weights(labels.n_cols, arma::fill::ones);

  // Build decision tree.
  DecisionTree<> d(inputData, labels, 3, 10); // Leaf size of 10.
  DecisionTree<> wd(inputData, labels, 3, weights, 10); // Leaf size of 10.

  // Load testing data.
  arma::mat testData;
  if (!data::Load("vc2_test.csv", testData))
    BOOST_FAIL("Cannot load test dataset vc2_test.csv!");

  arma::Mat<size_t> trueTestLabels;
  if (!data::Load("vc2_test_labels.txt", trueTestLabels))
    BOOST_FAIL("Cannot load labels for vc2_test_labels.txt");

  // Get the predicted test labels.
  arma::Row<size_t> predictions;
  d.Classify(testData, predictions);

  BOOST_REQUIRE_EQUAL(predictions.n_elem, testData.n_cols);

  // Figure out the accuracy.
  double correct = 0.0;
  for (size_t i = 0; i < predictions.n_elem; ++i)
    if (predictions[i] == trueTestLabels[i])
      ++correct;
  correct /= predictions.n_elem;

  BOOST_REQUIRE_GT(correct, 0.75);

  // reset the prediction
  predictions.zeros();
  wd.Classify(testData, predictions);

  BOOST_REQUIRE_EQUAL(predictions.n_elem, testData.n_cols);

  // Figure out the accuracy.
  double wdcorrect = 0.0;
  for (size_t i = 0; i < predictions.n_elem; ++i)
    if (predictions[i] == trueTestLabels[i])
      ++wdcorrect;
  wdcorrect /= predictions.n_elem;

  BOOST_REQUIRE_GT(wdcorrect, 0.75);
}

/**
 * Test that we can build a decision tree on a simple categorical dataset.
 */
BOOST_AUTO_TEST_CASE(CategoricalBuildTest)
{
  arma::mat d;
  arma::Row<size_t> l;
  data::DatasetInfo di;
  MockCategoricalData(d, l, di);

  // Split into a training set and a test set.
  arma::mat trainingData = d.cols(0, 1999);
  arma::mat testData = d.cols(2000, 3999);
  arma::Row<size_t> trainingLabels = l.subvec(0, 1999);
  arma::Row<size_t> testLabels = l.subvec(2000, 3999);

  // Build the tree.
  DecisionTree<> tree(trainingData, di, trainingLabels, 5, 10);

  // Now evaluate the accuracy of the tree.
  arma::Row<size_t> predictions;
  tree.Classify(testData, predictions);

  BOOST_REQUIRE_EQUAL(predictions.n_elem, testData.n_cols);
  size_t correct = 0;
  for (size_t i = 0; i < testData.n_cols; ++i)
    if (testLabels[i] == predictions[i])
      ++correct;

  // Make sure we got at least 70% accuracy.
  const double correctPct = double(correct) / double(testData.n_cols);
  BOOST_REQUIRE_GT(correctPct, 0.70);
}

/**
 * Test that we can build a decision tree with weights on a simple categorical
 * dataset.
 */
BOOST_AUTO_TEST_CASE(CategoricalBuildTestWithWeight)
{
  arma::mat d;
  arma::Row<size_t> l;
  data::DatasetInfo di;
  MockCategoricalData(d, l, di);

  // Split into a training set and a test set.
  arma::mat trainingData = d.cols(0, 1999);
  arma::mat testData = d.cols(2000, 3999);
  arma::Row<size_t> trainingLabels = l.subvec(0, 1999);
  arma::Row<size_t> testLabels = l.subvec(2000, 3999);

  arma::Row<double> weights = arma::ones<arma::Row<double>>(
      trainingLabels.n_elem);

  // Build the tree.
  DecisionTree<> tree(trainingData, di, trainingLabels, 5, weights, 10);

  // Now evaluate the accuracy of the tree.
  arma::Row<size_t> predictions;
  tree.Classify(testData, predictions);

  BOOST_REQUIRE_EQUAL(predictions.n_elem, testData.n_cols);
  size_t correct = 0;
  for (size_t i = 0; i < testData.n_cols; ++i)
    if (testLabels[i] == predictions[i])
      ++correct;

  // Make sure we got at least 70% accuracy.
  const double correctPct = double(correct) / double(testData.n_cols);
  BOOST_REQUIRE_GT(correctPct, 0.70);
}

/**
 * Make sure that when we ask for a decision stump, we get one.
 */
BOOST_AUTO_TEST_CASE(DecisionStumpTest)
{
  // Use a random dataset.
  arma::mat dataset(10, 1000, arma::fill::randu);
  arma::Row<size_t> labels(1000);
  for (size_t i = 0; i < 1000; ++i)
    labels[i] = i % 3; // 3 classes.

  // Build a decision stump.
  DecisionTree<GiniGain, BestBinaryNumericSplit, AllCategoricalSplit,
      AllDimensionSelect, double, true> stump(dataset, labels, 3, 1);

  // Check that it has children.
  BOOST_REQUIRE_EQUAL(stump.NumChildren(), 2);
  // Check that its children doesn't have children.
  BOOST_REQUIRE_EQUAL(stump.Child(0).NumChildren(), 0);
  BOOST_REQUIRE_EQUAL(stump.Child(1).NumChildren(), 0);
}

/**
 * Test that we can build a decision tree using weighted data (where the
 * low-weighted data is random noise), and that the tree still builds correctly
 * enough to get good results.
 */
BOOST_AUTO_TEST_CASE(WeightedDecisionTreeTest)
{
  arma::mat dataset;
  arma::Row<size_t> labels;
  data::Load("vc2.csv", dataset);
  data::Load("vc2_labels.txt", labels);

  // Add some noise.
  arma::mat noise(dataset.n_rows, 1000, arma::fill::randu);
  arma::Row<size_t> noiseLabels(1000);
  for (size_t i = 0; i < noiseLabels.n_elem; ++i)
    noiseLabels[i] = math::RandInt(3); // Random label.

  // Concatenate data matrices.
  arma::mat data = arma::join_rows(dataset, noise);
  arma::Row<size_t> fullLabels = arma::join_rows(labels, noiseLabels);

  // Now set weights.
  arma::rowvec weights(dataset.n_cols + 1000);
  for (size_t i = 0; i < dataset.n_cols; ++i)
    weights[i] = math::Random(0.9, 1.0);
  for (size_t i = dataset.n_cols; i < dataset.n_cols + 1000; ++i)
    weights[i] = math::Random(0.0, 0.01); // Low weights for false points.

  // Now build the decision tree.  I think the syntax is right here.
  DecisionTree<> d(data, fullLabels, 3, weights, 10);

  // Now we can check that we get good performance on the VC2 test set.
  arma::mat testData;
  arma::Row<size_t> testLabels;
  data::Load("vc2_test.csv", testData);
  data::Load("vc2_test_labels.txt", testLabels);

  arma::Row<size_t> predictions;
  d.Classify(testData, predictions);

  BOOST_REQUIRE_EQUAL(predictions.n_elem, testData.n_cols);

  // Figure out the accuracy.
  double correct = 0.0;
  for (size_t i = 0; i < predictions.n_elem; ++i)
    if (predictions[i] == testLabels[i])
      ++correct;
  correct /= predictions.n_elem;

  BOOST_REQUIRE_GT(correct, 0.75);
}
/**
 * Test that we can build a decision tree on a simple categorical dataset using
 * weights, with low-weight noise added.
 */
BOOST_AUTO_TEST_CASE(CategoricalWeightedBuildTest)
{
  arma::mat d;
  arma::Row<size_t> l;
  data::DatasetInfo di;
  MockCategoricalData(d, l, di);

  // Split into a training set and a test set.
  arma::mat trainingData = d.cols(0, 1999);
  arma::mat testData = d.cols(2000, 3999);
  arma::Row<size_t> trainingLabels = l.subvec(0, 1999);
  arma::Row<size_t> testLabels = l.subvec(2000, 3999);

  // Now create random points.
  arma::mat randomNoise(4, 2000);
  arma::Row<size_t> randomLabels(2000);
  for (size_t i = 0; i < 2000; ++i)
  {
    randomNoise(0, i) = math::Random();
    randomNoise(1, i) = math::Random();
    randomNoise(2, i) = math::RandInt(4);
    randomNoise(3, i) = math::RandInt(2);
    randomLabels[i] = math::RandInt(5);
  }

  // Generate weights.
  arma::rowvec weights(4000);
  for (size_t i = 0; i < 2000; ++i)
    weights[i] = math::Random(0.9, 1.0);
  for (size_t i = 2000; i < 4000; ++i)
    weights[i] = math::Random(0.0, 0.001);

  arma::mat fullData = arma::join_rows(trainingData, randomNoise);
  arma::Row<size_t> fullLabels = arma::join_rows(trainingLabels, randomLabels);

  // Build the tree.
  DecisionTree<> tree(fullData, di, fullLabels, 5, weights, 10);

  // Now evaluate the accuracy of the tree.
  arma::Row<size_t> predictions;
  tree.Classify(testData, predictions);

  BOOST_REQUIRE_EQUAL(predictions.n_elem, testData.n_cols);
  size_t correct = 0;
  for (size_t i = 0; i < testData.n_cols; ++i)
    if (testLabels[i] == predictions[i])
      ++correct;

  // Make sure we got at least 70% accuracy.
  const double correctPct = double(correct) / double(testData.n_cols);
  BOOST_REQUIRE_GT(correctPct, 0.70);
}

/**
 * Test that we can build a decision tree using weighted data (where the
 * low-weighted data is random noise) with information gain, and that the tree
 * still builds correctly enough to get good results.
 */
BOOST_AUTO_TEST_CASE(WeightedDecisionTreeInformationGainTest)
{
  arma::mat dataset;
  arma::Row<size_t> labels;
  data::Load("vc2.csv", dataset);
  data::Load("vc2_labels.txt", labels);

  // Add some noise.
  arma::mat noise(dataset.n_rows, 1000, arma::fill::randu);
  arma::Row<size_t> noiseLabels(1000);
  for (size_t i = 0; i < noiseLabels.n_elem; ++i)
    noiseLabels[i] = math::RandInt(3); // Random label.

  // Concatenate data matrices.
  arma::mat data = arma::join_rows(dataset, noise);
  arma::Row<size_t> fullLabels = arma::join_rows(labels, noiseLabels);

  // Now set weights.
  arma::rowvec weights(dataset.n_cols + 1000);
  for (size_t i = 0; i < dataset.n_cols; ++i)
    weights[i] = math::Random(0.9, 1.0);
  for (size_t i = dataset.n_cols; i < dataset.n_cols + 1000; ++i)
    weights[i] = math::Random(0.0, 0.01); // Low weights for false points.

  // Now build the decision tree.  I think the syntax is right here.
  DecisionTree<InformationGain> d(data, fullLabels, 3, weights, 10);

  // Now we can check that we get good performance on the VC2 test set.
  arma::mat testData;
  arma::Row<size_t> testLabels;
  data::Load("vc2_test.csv", testData);
  data::Load("vc2_test_labels.txt", testLabels);

  arma::Row<size_t> predictions;
  d.Classify(testData, predictions);

  BOOST_REQUIRE_EQUAL(predictions.n_elem, testData.n_cols);

  // Figure out the accuracy.
  double correct = 0.0;
  for (size_t i = 0; i < predictions.n_elem; ++i)
    if (predictions[i] == testLabels[i])
      ++correct;
  correct /= predictions.n_elem;

  BOOST_REQUIRE_GT(correct, 0.75);
}
/**
 * Test that we can build a decision tree using information gain on a simple
 * categorical dataset using weights, with low-weight noise added.
 */
BOOST_AUTO_TEST_CASE(CategoricalInformationGainWeightedBuildTest)
{
  arma::mat d;
  arma::Row<size_t> l;
  data::DatasetInfo di;
  MockCategoricalData(d, l, di);

  // Split into a training set and a test set.
  arma::mat trainingData = d.cols(0, 1999);
  arma::mat testData = d.cols(2000, 3999);
  arma::Row<size_t> trainingLabels = l.subvec(0, 1999);
  arma::Row<size_t> testLabels = l.subvec(2000, 3999);

  // Now create random points.
  arma::mat randomNoise(4, 2000);
  arma::Row<size_t> randomLabels(2000);
  for (size_t i = 0; i < 2000; ++i)
  {
    randomNoise(0, i) = math::Random();
    randomNoise(1, i) = math::Random();
    randomNoise(2, i) = math::RandInt(4);
    randomNoise(3, i) = math::RandInt(2);
    randomLabels[i] = math::RandInt(5);
  }

  // Generate weights.
  arma::rowvec weights(4000);
  for (size_t i = 0; i < 2000; ++i)
    weights[i] = math::Random(0.9, 1.0);
  for (size_t i = 2000; i < 4000; ++i)
    weights[i] = math::Random(0.0, 0.001);

  arma::mat fullData = arma::join_rows(trainingData, randomNoise);
  arma::Row<size_t> fullLabels = arma::join_rows(trainingLabels, randomLabels);

  // Build the tree.
  DecisionTree<InformationGain> tree(fullData, di, fullLabels, 5, weights, 10);

  // Now evaluate the accuracy of the tree.
  arma::Row<size_t> predictions;
  tree.Classify(testData, predictions);

  BOOST_REQUIRE_EQUAL(predictions.n_elem, testData.n_cols);
  size_t correct = 0;
  for (size_t i = 0; i < testData.n_cols; ++i)
    if (testLabels[i] == predictions[i])
      ++correct;

  // Make sure we got at least 70% accuracy.
  const double correctPct = double(correct) / double(testData.n_cols);
  BOOST_REQUIRE_GT(correctPct, 0.70);
}

/**
 * Make sure that the random dimension selector only has one element.
 */
BOOST_AUTO_TEST_CASE(RandomDimensionSelectTest)
{
  RandomDimensionSelect r;
  r.Dimensions() = 10;

  BOOST_REQUIRE_LT(r.Begin(), 10);
  BOOST_REQUIRE_EQUAL(r.Next(), r.End());
  BOOST_REQUIRE_EQUAL(r.Next(), r.End());
  BOOST_REQUIRE_EQUAL(r.Next(), r.End());
}

/**
 * Make sure that the random dimension selector selects different values.
 */
BOOST_AUTO_TEST_CASE(RandomDimensionSelectRandomTest)
{
  // We'll check that 4 values are not all the same.
  RandomDimensionSelect r1, r2, r3, r4;
  r1.Dimensions() = 100000;
  r2.Dimensions() = 100000;
  r3.Dimensions() = 100000;
  r4.Dimensions() = 100000;

  BOOST_REQUIRE((r1.Begin() != r2.Begin()) ||
                (r1.Begin() != r3.Begin()) ||
                (r1.Begin() != r4.Begin()));
}

/**
 * Make sure that the multiple random dimension select only has the right number
 * of elements.
 */
BOOST_AUTO_TEST_CASE(MultipleRandomDimensionSelectTest)
{
  MultipleRandomDimensionSelect r(5);
  r.Dimensions() = 10;

  // Make sure we get five elements.
  BOOST_REQUIRE_LT(r.Begin(), 10);
  BOOST_REQUIRE_LT(r.Next(), 10);
  BOOST_REQUIRE_LT(r.Next(), 10);
  BOOST_REQUIRE_LT(r.Next(), 10);
  BOOST_REQUIRE_LT(r.Next(), 10);
  BOOST_REQUIRE_EQUAL(r.Next(), r.End());
}

/**
 * Make sure we get every element from the distribution.
 */
BOOST_AUTO_TEST_CASE(MultipleRandomDimensionAllSelectTest)
{
  MultipleRandomDimensionSelect r(3);
  r.Dimensions() = 3;

  bool found[3];
  found[0] = found[1] = found[2] = false;

  found[r.Begin()] = true;
  found[r.Next()] = true;
  found[r.Next()] = true;

  BOOST_REQUIRE_EQUAL(found[0], true);
  BOOST_REQUIRE_EQUAL(found[1], true);
  BOOST_REQUIRE_EQUAL(found[2], true);
}

/**
 * Make sure the right number of classes is returned for an empty tree (1).
 */
BOOST_AUTO_TEST_CASE(NumClassesEmptyTreeTest)
{
  DecisionTree<> dt;
  BOOST_REQUIRE_EQUAL(dt.NumClasses(), 1);
}

/**
 * Make sure the right number of classes is returned for a nonempty tree.
 */
BOOST_AUTO_TEST_CASE(NumClassesTest)
{
  // Load a dataset to train with.
  arma::mat dataset;
  arma::Row<size_t> labels;
  data::Load("vc2.csv", dataset);
  data::Load("vc2_labels.txt", labels);

  DecisionTree<> dt(dataset, labels, 3);

  BOOST_REQUIRE_EQUAL(dt.NumClasses(), 3);
}

/*
 * Test that we can pass const data into DecisionTree constructors.
 */
BOOST_AUTO_TEST_CASE(ConstDataTest)
{
  arma::mat data;
  arma::Row<size_t> labels;
  data::DatasetInfo datasetInfo;
  MockCategoricalData(data, labels, datasetInfo);

  const arma::mat& constData = data;
  const arma::Row<size_t>& constLabels = labels;
  const arma::rowvec constWeights(labels.n_elem, arma::fill::randu);
  const size_t numClasses = 5;

  DecisionTree<> dt(constData, constLabels, numClasses);
  DecisionTree<> dt2(constData, datasetInfo, constLabels, numClasses);
  DecisionTree<> dt3(constData, constLabels, numClasses, constWeights);
  DecisionTree<> dt4(constData, datasetInfo, constLabels, numClasses,
      constWeights);
}

/**
 * Construct the decision tree with splitting only if gain is more than
 * threshold.
 */
BOOST_AUTO_TEST_CASE(RegularisedDecisionTree)
{
  // Completely random dataset with no structure.
  arma::mat dataset(10, 1000, arma::fill::randu);
  arma::Row<size_t> labels(1000);
  for (size_t i = 0; i < 1000; ++i)
    labels[i] = i % 3; // 3 classes.
  arma::rowvec weights(labels.n_elem);
  weights.ones();

  // Minimum leaf size of 1.
  DecisionTree<> d(dataset, labels, 3, weights, 1, 1e-7);

  // Minimum leaf size of 1 and Minimum gain split of 0.01.
  DecisionTree<> dRegularised(dataset, labels, 3, weights, 1, 0.01);

  size_t count = 0;
  // This part of code is dupliacte with no weighted one.
  for (size_t i = 0; i < 1000; ++i)
  {
    size_t prediction, predictionsregularised;
    arma::vec probabilities, probabilitiesRegularised;

    d.Classify(dataset.col(i), prediction, probabilities);
    dRegularised.Classify(dataset.col(i), predictionsregularised,
                          probabilitiesRegularised);

    if (prediction != predictionsregularised)
      count++;

    BOOST_REQUIRE_EQUAL(probabilities.n_elem, 3);
    BOOST_REQUIRE_EQUAL(probabilitiesRegularised.n_elem, 3);
  }

  BOOST_REQUIRE_GT(count, 0);
}

/**
 * Test that DecisionTree::Train() returns finite entropy on numeric dataset.
 */
BOOST_AUTO_TEST_CASE(DecisionTreeNumericTrainReturnEntropy)
{
  arma::mat dataset(10, 1000, arma::fill::randu);
  arma::Row<size_t> labels(1000);
  arma::rowvec weights(labels.n_elem);
  weights.ones();

  for (size_t i = 0; i < 1000; ++i)
    labels[i] = i % 3; // 3 classes.

  // Train a simpe tree on numeric dataset.
  DecisionTree<> d(3);
  double entropy = d.Train(dataset, labels, 3, 50);

  BOOST_REQUIRE_EQUAL(std::isfinite(entropy), true);

  // Train a tree with weights on numeric dataset.
  DecisionTree<> wd(3);
  entropy = wd.Train(dataset, labels, 3, weights, 50);

  BOOST_REQUIRE_EQUAL(std::isfinite(entropy), true);
}

/**
 * Test that DecisionTree::Train() returns finite entropy on categorical
 * dataset.
 */
BOOST_AUTO_TEST_CASE(DecisionTreeCategoricalTrainReturnEntropy)
{
  arma::mat d;
  arma::Row<size_t> l;
  data::DatasetInfo di;
  MockCategoricalData(d, l, di);

  arma::Row<double> weights = arma::ones<arma::Row<double>>(l.n_elem);

  // Train a simple tree on categorical dataset.
  DecisionTree<> dtree(5);
  double entropy = dtree.Train(d, di, l, 5, 10);

  BOOST_REQUIRE_EQUAL(std::isfinite(entropy), true);

  // Train a tree with weights on categorical dataset.
  DecisionTree<> wdtree(5);
  entropy = wdtree.Train(d, di, l, 5, weights, 10);

  BOOST_REQUIRE_EQUAL(std::isfinite(entropy), true);
}

/**
 * Make sure different maximum depth values give different numbers of children.
 */
BOOST_AUTO_TEST_CASE(DifferentMaximumDepthTest)
{
  arma::mat dataset;
  arma::Row<size_t> labels;
  data::Load("vc2.csv", dataset);
  data::Load("vc2_labels.txt", labels);

  DecisionTree<> d(dataset, labels, 3, 10, 1e-7, 1);

  DecisionTree<> d1(dataset, labels, 3, 10, 1e-7, 2);

  DecisionTree<> d2(dataset, labels, 3, 10, 1e-7);

  // Now require that we have zero children.
  BOOST_REQUIRE_EQUAL(d.NumChildren(), 0);

  // Now require that we have two children.
  BOOST_REQUIRE_EQUAL(d1.NumChildren(), 2);
  BOOST_REQUIRE_EQUAL(d1.Child(0).NumChildren(), 0);
  BOOST_REQUIRE_EQUAL(d1.Child(1).NumChildren(), 0);

  // Now require that we have two children.
  BOOST_REQUIRE_EQUAL(d2.NumChildren(), 2);
  BOOST_REQUIRE_EQUAL(d2.Child(0).NumChildren(), 2);
  BOOST_REQUIRE_EQUAL(d2.Child(1).NumChildren(), 2);
}

BOOST_AUTO_TEST_SUITE_END();