////////////////////////////////////////////////////////////////////////////////
#include <polyfem/RBFWithLinear.hpp>
#include <polyfem/Types.hpp>
#include <polyfem/Logger.hpp>

#include <igl/Timer.h>

#include <Eigen/Dense>

#include <iostream>
#include <fstream>
#include <array>
////////////////////////////////////////////////////////////////////////////////

using namespace polyfem;

namespace {

// Harmonic kernel
double kernel(const bool is_volume, const double r) {
	if (r < 1e-8) { return 0; }

	if (is_volume) {
		return 1/r;
	} else {
		return log(r);
	}
}

double kernel_prime(const bool is_volume, const double r) {
	if (r < 1e-8) { return 0; }

	if(is_volume) {
		return -1/(r*r);
	} else {
		return 1/r;
	}
}

} // anonymous namespace

////////////////////////////////////////////////////////////////////////////////

RBFWithLinear::RBFWithLinear(
		const Eigen::MatrixXd &centers,
		const Eigen::MatrixXd &samples,
		const Eigen::MatrixXd &local_basis_integral,
		const Quadrature &quadr,
		Eigen::MatrixXd &rhs,
		bool with_constraints)
	: centers_(centers)
{
	compute_weights(samples, local_basis_integral, quadr, rhs, with_constraints);
}

// -----------------------------------------------------------------------------

void RBFWithLinear::basis(const int local_index, const Eigen::MatrixXd &samples, Eigen::MatrixXd &val) const {
	Eigen::MatrixXd tmp;
	bases_values(samples, tmp);
	val = tmp.col(local_index);
}

// -----------------------------------------------------------------------------

void RBFWithLinear::grad(const int local_index, const Eigen::MatrixXd &samples, Eigen::MatrixXd &val) const {
	Eigen::MatrixXd tmp;
	const int dim = centers_.cols();
	val.resize(samples.rows(), dim);
	for (int d = 0; d < dim; ++d) {
		bases_grads(d, samples, tmp);
		val.col(d) = tmp.col(local_index);
	}
}

////////////////////////////////////////////////////////////////////////////////

void RBFWithLinear::bases_values(const Eigen::MatrixXd &samples, Eigen::MatrixXd &val) const {
	// Compute A
	Eigen::MatrixXd A;
	compute_kernels_matrix(samples, A);

	// Multiply by the weights
	val = A * weights_;
}

// -----------------------------------------------------------------------------

void RBFWithLinear::bases_grads(const int axis, const Eigen::MatrixXd &samples, Eigen::MatrixXd &val) const {
	const int num_kernels = centers_.rows();
	const int dim = centers_.cols();

	// Compute ∇xA
	Eigen::MatrixXd A_prime(samples.rows(), num_kernels + 1 + dim);
	A_prime.setZero();

	for (int j = 0; j < num_kernels; ++j) {
		A_prime.col(j) = (samples.rowwise() - centers_.row(j)).rowwise().norm().unaryExpr([this](double x)
			{ return kernel_prime(is_volume(), x) / x; });
		A_prime.col(j) = (samples.col(axis).array() - centers_(j, axis)) * A_prime.col(j).array();
	}
	// Linear terms
	A_prime.middleCols(num_kernels + 1 + axis, 1).setOnes();

	// Apply weights
	val = A_prime * weights_;
}

////////////////////////////////////////////////////////////////////////////////

void RBFWithLinear::compute_kernels_matrix(const Eigen::MatrixXd &samples, Eigen::MatrixXd &A) const {
	// Compute A
	const int num_kernels = centers_.rows();
	const int dim = centers_.cols();

	A.resize(samples.rows(), num_kernels + 1 + dim);
	for (int j = 0; j < num_kernels; ++j) {
		A.col(j) = (samples.rowwise() - centers_.row(j)).rowwise().norm().unaryExpr([this](double x)
			{ return kernel(is_volume(), x); });
	}
	A.col(num_kernels).setOnes(); // constant term
	A.rightCols(dim) = samples; // linear terms
}

// -----------------------------------------------------------------------------

void RBFWithLinear::compute_constraints_matrix(
	const int num_bases,
	const Quadrature &quadr,
	const Eigen::MatrixXd &local_basis_integral,
	Eigen::MatrixXd &L,
	Eigen::MatrixXd &t) const
{
	const int num_kernels = centers_.rows();
	const int dim = centers_.cols();

	// Compute KI
	Eigen::MatrixXd KI(num_kernels, dim);
	for (int j = 0; j < num_kernels; ++j) {
		// ∫∇x(φj)(p) = Σ_q (xq - xk) * 1/r * h'(r) * wq
		// - xq is the x coordinate of the q-th quadrature point
		// - wq is the q-th quadrature weight
		// - r is the distance from pq to the kernel center
		// - h is the RBFWithLinear RBF kernel (scalar function)
		const Eigen::MatrixXd drdp = quadr.points.rowwise() - centers_.row(j);
		const Eigen::VectorXd r = drdp.rowwise().norm();
		KI.row(j) = (drdp.array().colwise() * (quadr.weights.array() * r.unaryExpr([this](double x)
			{ return kernel_prime(is_volume(), x); }).array() / r.array())).colwise().sum();
	}
	KI /= quadr.weights.sum();

	// Compute L
	L.resize(num_kernels + dim + 1, num_kernels + 1);
	L.setZero();
	L.diagonal().setOnes();
	L.block(num_kernels + 1, 0, dim, num_kernels) = -KI.transpose();
	// std::cout << L.bottomRightCorner(10, 10) << std::endl;

	// Compute t
	t.resize(num_kernels + 1 + dim, num_bases);
	t.setZero();
	t.bottomRows(dim) = local_basis_integral.transpose().topRows(dim) / quadr.weights.sum();
}

// -----------------------------------------------------------------------------

void RBFWithLinear::compute_weights(const Eigen::MatrixXd &samples,
	const Eigen::MatrixXd &local_basis_integral, const Quadrature &quadr,
	Eigen::MatrixXd &rhs, bool with_constraints)
{
	logger().trace("#kernel centers: {}", centers_.rows());
	logger().trace("#collocation points: {}", samples.rows());
	logger().trace("#quadrature points: {}", quadr.weights.size());
	logger().trace("#non-vanishing bases: {}", rhs.cols());

	if (!with_constraints) {
		// Compute A
		Eigen::MatrixXd A;
		compute_kernels_matrix(samples, A);

		// Solve the system
		const int num_kernels = centers_.rows();
		logger().trace("-- Solving system of size {}x{}", num_kernels, num_kernels);
		weights_ = (A.transpose() * A).ldlt().solve(A.transpose() * rhs);
		logger().trace("-- Solved!");

		return;
	}

	// For each basis function f that is nonzero on the element E, we want to
	// solve the least square system A w = rhs, where:
	//     ┏                    ┓
	//     ┃ φj(pi) ... 1 xi yi ┃
	// A = ┃   ┊        ┊  ┊  ┊ ┃ ∊ ℝ^{#S x (#K+1+dim)}
	//     ┃   ┊        ┊  ┊  ┊ ┃
	//     ┗                    ┛
	//     ┏                    ┓^⊤
	// w = ┃ wj ... a00 a10 a01 ┃   ∊ ℝ^{#K+1+dim}
	//     ┗                    ┛
	// - A is the RBF kernels evaluated over the collocation points (#S)
	// - b is the expected value of the basis sampled on the boundary (#S)
	// - w is the weight of the kernels defining the basis
	// - pi = (xi, yi) is the i-th collocation point
	//
	// Moreover, we want to impose a constraint on the gradients of the kernels
	// so that the integral of the gradients over the polytope must be equal to
	// the value specified in the argument `local_basis_integral` (#K)
	//
	// Let `lb` be the precomputed expected value of ∫f over the rest of the mesh.
	// We write down the constraint as:
	//
	// ∫_{p ∊ E} Σ_j wj ∇x(φj)(p) + ∇x(a^⊤·p + c) dp = lb       (1)
	// (1) ⇔ ∫_{p ∊ E} Σ_j wj ∇x(φj)(p) + ax dp = lb
	//     ⇔ lb - Σ_j wj ∫_{p ∊ E} ∇x(φj)(p) dp = ax Vol(E)
	//
	// We now have a relationship w = Lv + t, where the weights (and esp. the
	// linear terms in the weight vector w), are expressed as an affine
	// combination of unknowns v = [wj ... c] ∊ ℝ^{#K+1} and a translation t
	//
	// After solving the new least square system A L v = rhs - A t, we can retrieve
	// w = L v

	//
	//     ┏                      ┓^⊤
	// t = ┃ 0  ┈  ┈  0 0 lbx lby ┃   / Vol(E) ∊ ℝ^{#K+1+dim}
	//     ┗                      ┛
	//
	//     ┏                  ┓
	//     ┃   1              ┃
	//     ┃       1          ┃
	//     ┃          ·       ┃
	// L = ┃             ·    ┃ ∊ ℝ^{ (#K+1+dim) x (#K+1}) }
	//     ┃                1 ┃
	//     ┃ Lx_j  ┈        0 ┃
	//     ┃ Ly_j  ┈        0 ┃
	//     ┗                  ┛
	// Where Lx_j = -∫∇xφj / Vol(E) = -∫_{p ∊ E} ∇x(φj)(p) / Vol(E) is integrated numerically
	//

	const int num_bases = rhs.cols();

	// Compute A
	Eigen::MatrixXd A;
	compute_kernels_matrix(samples, A);

	// Compute L and t
	// Note that t is stored into `weights_` for memory efficiency reasons
	Eigen::MatrixXd L;
	compute_constraints_matrix(num_bases, quadr, local_basis_integral, L, weights_);

	// Compute b = rhs - A t
	rhs -= A * weights_;

	// Solve the system
	logger().trace("-- Solving system of size {}x{}", L.cols(), L.cols());
	weights_ += L * (L.transpose() * A.transpose() * A * L).ldlt().solve(L.transpose() * A.transpose() * rhs);
	logger().trace("-- Solved!");

	// std::cout << weights_.bottomRows(10) << std::endl;

	// Eigen::MatrixXd M, x, dx;
	// grad(0, quadr.points, M);
	// for (int d = 0; d < dim; ++d) {
	// 	basis(0, quadr.points, x);
	// 	auto asd = quadr.points;
	// 	asd.col(d).array() += 1e-7;
	// 	basis(0, asd, dx);
	// 	std::cout << (dx - x) / 1e-7 - M.col(d) << std::endl;
	// 	std::cout << (M.col(d).array() * quadr.weights.array()).sum() - local_basis_integral(0, d) << std::endl;
	// }
}


#ifndef INCLUDED_STDDEFX
#include "stddefx.h"
#define INCLUDED_STDDEFX
#endif

#ifndef INCLUDED_GEO_RIKSNEIGHBOURHOOD
#include "geo_riksneighbourhood.h"
#define INCLUDED_GEO_RIKSNEIGHBOURHOOD
#endif

// Library headers.

#ifndef INCLUDED_BOOST_MATH_TR1
#include <boost/math/tr1.hpp>
#define INCLUDED_BOOST_MATH_TR1
#endif
using namespace boost::math; // use then tr1

// PCRaster library headers.
#ifndef INCLUDED_COM_MATH
#include "com_math.h"
#define INCLUDED_COM_MATH
#endif

// Module headers.

/*!
  \file
  This file contains the implementation of the RiksNeighbourhood class.
*/

//------------------------------------------------------------------------------

/*
namespace geo {

class RiksNeighbourhoodPrivate
{
public:

  RiksNeighbourhoodPrivate()
  {
  }

  ~RiksNeighbourhoodPrivate()
  {
  }

};

} // namespace geo
*/

//------------------------------------------------------------------------------
// DEFINITION OF STATIC RIKSNEIGHBOURHOOD MEMBERS
//------------------------------------------------------------------------------

boost::tuple<size_t, size_t> geo::RiksNeighbourhood::circleCell(double radius)
{
  PRECOND(radius > 0.0);

  boost::tuple<size_t, size_t> cell;
  double currentDifference(0), difference(0);

  // Determine max radius of raster which can contain a circle with given
  // radius.
  size_t maxRadius = static_cast<size_t>(std::ceil(radius));

  // Only handle one quarter of the circle.
  for(size_t row = 0; row <= maxRadius; ++row) {
    for(size_t col = 0; col <= maxRadius; ++col) {

      if(row == 0 && col == 0) {
        // First cell handled.
        difference = std::abs(radius - tr1::hypot<double>(row, col));
        cell = boost::make_tuple(row, col);
      }
      else {
        currentDifference = std::abs(radius - tr1::hypot<double>(row, col));
        if(currentDifference < difference) {
          difference = currentDifference;
          cell = boost::make_tuple(row, col);
        }
      }
    }
  }

  return cell;
}

//------------------------------------------------------------------------------
// DEFINITION OF RIKSNEIGHBOURHOOD MEMBERS
//------------------------------------------------------------------------------

geo::RiksNeighbourhood::RiksNeighbourhood(double toRadius)

  : Neighbourhood(toRadius)

{
  init();
}

geo::RiksNeighbourhood::RiksNeighbourhood(double fromRadius, double toRadius)

  : Neighbourhood(fromRadius, toRadius)

{
  init();
}

geo::RiksNeighbourhood::~RiksNeighbourhood()
{
}

//!
/*!
  \param     .
  \return    .
  \exception .
  \warning   .
  \sa        .
*/
void geo::RiksNeighbourhood::init()
{
  // Determine which cells match the given radiusses best. These are the cells
  // who's centers are closest to the circle defined by the radiusses.
  boost::tuple<size_t, size_t> fromCircleCell = boost::make_tuple(0, 0);
  boost::tuple<size_t, size_t> toCircleCell(fromCircleCell);

  if(fromRadius() > 0.0) {
    fromCircleCell = circleCell(fromRadius());
  }

  if(toRadius() > 0.0) {
    toCircleCell = circleCell(toRadius());
  }

  // Determine the radiusses of the Riks neighbourhoods of which the selected
  // cells are part.
  double fromRadius = tr1::hypot<double>(
         fromCircleCell.get<0>(), fromCircleCell.get<1>());
  double toRadius = tr1::hypot<double>(
         toCircleCell.get<0>(), toCircleCell.get<1>());

  // Make sure the current set radius is equal of larger than the selected one.
  POSTCOND(static_cast<double>(radius()) >= toRadius);

  // Determine which cells are part of the Riks neighbourhood with the
  // selected radius.
  size_t offset = radius();
  for(size_t row = 0; row <= radius(); ++row) {
    for(size_t col = 0; col <= radius(); ++col) {
      double radius = tr1::hypot<double>(row, col);
      if((radius > fromRadius && radius < toRadius) ||
          com::equal_epsilon(radius, fromRadius) ||
          com::equal_epsilon(radius, toRadius)) {
        this->cell(offset + row, offset + col) = 1.0;
        this->cell(offset + row, offset - col) = 1.0;
        this->cell(offset - row, offset + col) = 1.0;
        this->cell(offset - row, offset - col) = 1.0;
      }
    }
  }
}

//------------------------------------------------------------------------------
// DEFINITION OF FREE OPERATORS
//------------------------------------------------------------------------------

//------------------------------------------------------------------------------
// DEFINITION OF FREE FUNCTIONS
//------------------------------------------------------------------------------


/** Copyright 2017-2020 CNRS-AIST JRL and CNRS-UM LIRMM */

#include <tvm/Space.h>
#include <tvm/Variable.h>
#include <tvm/constraint/BasicLinearConstraint.h>
#include <tvm/constraint/abstract/LinearConstraint.h>
#include <tvm/hint/Substitution.h>
#include <tvm/hint/internal/DiagonalCalculator.h>
#include <tvm/hint/internal/GenericCalculator.h>
#include <tvm/hint/internal/Substitutions.h>
#include <tvm/internal/MatrixProperties.h>
#include <tvm/internal/VariableVectorPartition.h>

#include <Eigen/SVD>

//#include <iostream>
#include <vector>

#define DOCTEST_CONFIG_IMPLEMENT_WITH_MAIN
#define DOCTEST_CONFIG_SUPER_FAST_ASSERTS
#include "doctest/doctest.h"

using namespace tvm;
using namespace tvm::hint;
using namespace tvm::hint::internal;
using namespace Eigen;

TEST_CASE("GenericCalculator")
{
  VariablePtr x = Space(5).createVariable("x");
  VariablePtr y = Space(3).createVariable("y");
  MatrixXd A = MatrixXd::Random(5, 5);
  MatrixXd B = MatrixXd::Random(5, 3);
  VectorXd b = VectorXd::Random(5);

  std::shared_ptr<constraint::BasicLinearConstraint> c(
      new constraint::BasicLinearConstraint({A, B}, {x, y}, b, constraint::Type::EQUAL));

  auto calc = GenericCalculator().impl({std::static_pointer_cast<constraint::abstract::LinearConstraint>(c)}, {x}, 5);
  calc->update();

  MatrixXd AsA(5, 5);
  MatrixXd StA(0, 5);
  calc->premultiplyByASharpAndSTranspose(AsA, StA, A, false);
  FAST_CHECK_UNARY(AsA.isIdentity());

  MatrixXd AsB(5, 3);
  MatrixXd StB(0, 3);
  calc->premultiplyByASharpAndSTranspose(AsB, StB, B, true);
  auto qrA = A.colPivHouseholderQr();
  FAST_CHECK_UNARY(AsB.isApprox(-MatrixXd(qrA.solve(B))));

  MatrixXd C(3, 5);
  C << 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0;
  MatrixXd D = MatrixXd::Random(3, 3);
  VectorXd d = VectorXd::Random(3);

  std::shared_ptr<constraint::BasicLinearConstraint> c2(
      new constraint::BasicLinearConstraint({C, D}, {x, y}, d, constraint::Type::EQUAL));

  auto calc2 = GenericCalculator().impl({std::static_pointer_cast<constraint::abstract::LinearConstraint>(c2)}, {x}, 2);
  calc2->update();
  MatrixXd CsD(5, 3);
  MatrixXd StD(1, 3);
  calc2->premultiplyByASharpAndSTranspose(CsD, StD, D, true);
  FAST_CHECK_UNARY(CsD.topRows(2).isApprox(-D.topRows(2)));
  FAST_CHECK_UNARY(CsD.bottomRows(3).isZero());
  FAST_CHECK_UNARY(StD.isApprox(D.row(2)));
  FAST_CHECK_UNARY(MatrixXd(C * calc2->N()).isZero());
  FAST_CHECK_EQ(calc2->N().colPivHouseholderQr().rank(), 3);
}

TEST_CASE("Diagonal Calculator")
{
  VariablePtr x = Space(7).createVariable("x");
  VariablePtr y = Space(3).createVariable("y");

  {
    MatrixXd A = MatrixXd::Identity(7, 7);
    MatrixXd B = MatrixXd::Random(7, 3);
    VectorXd b = VectorXd::Random(7);

    std::shared_ptr<constraint::BasicLinearConstraint> c(
        new constraint::BasicLinearConstraint({A, B}, {x, y}, b, constraint::Type::EQUAL));

    auto calc =
        DiagonalCalculator().impl({std::static_pointer_cast<constraint::abstract::LinearConstraint>(c)}, {x}, 7);
    calc->update();

    MatrixXd AsA(7, 7);
    MatrixXd StA(0, 7);
    calc->premultiplyByASharpAndSTranspose(AsA, StA, A, false);
    FAST_CHECK_UNARY(AsA.isIdentity());

    MatrixXd AsB(7, 3);
    MatrixXd StB(0, 3);
    calc->premultiplyByASharpAndSTranspose(AsB, StB, B, false);
    FAST_CHECK_UNARY(AsB.isApprox(B));

    FAST_CHECK_EQ(calc->N().rows(), 7);
    FAST_CHECK_EQ(calc->N().cols(), 0);
  }

  {
    MatrixXd A(3, 7);
    A << 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0;
    MatrixXd B = MatrixXd::Random(3, 3);
    VectorXd b = VectorXd::Random(3);

    std::shared_ptr<constraint::BasicLinearConstraint> c(
        new constraint::BasicLinearConstraint({A, B}, {x, y}, b, constraint::Type::EQUAL));

    auto calc = DiagonalCalculator({1, 3, 4}).impl(
        {std::static_pointer_cast<constraint::abstract::LinearConstraint>(c)}, {x}, 3);
    calc->update();

    MatrixXd AsA(7, 7);
    MatrixXd StA(0, 7);
    calc->premultiplyByASharpAndSTranspose(AsA, StA, A, false);
    FAST_CHECK_UNARY(AsA.isApprox(A.transpose() * A));

    MatrixXd AsB(7, 3);
    MatrixXd StB(0, 3);
    calc->premultiplyByASharpAndSTranspose(AsB, StB, B, false);
    FAST_CHECK_UNARY(AsB.isApprox(A.transpose() * B));

    MatrixXd N(7, 4);
    N << 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1;
    FAST_CHECK_UNARY(N.isApprox(calc->N()));
  }

  {
    MatrixXd A(5, 7);
    A << 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0;
    MatrixXd B = MatrixXd::Random(5, 3);
    VectorXd b = VectorXd::Random(5);

    std::shared_ptr<constraint::BasicLinearConstraint> c(
        new constraint::BasicLinearConstraint({A, B}, {x, y}, b, constraint::Type::EQUAL));

    auto calc = DiagonalCalculator({1, 3, 4}, {0, 2})
                    .impl({std::static_pointer_cast<constraint::abstract::LinearConstraint>(c)}, {x}, 3);
    calc->update();

    MatrixXd AsA(7, 7);
    MatrixXd StA(2, 7);
    calc->premultiplyByASharpAndSTranspose(AsA, StA, A, false);
    FAST_CHECK_UNARY(AsA.isApprox(A.transpose() * A));
    FAST_CHECK_UNARY(StA.isZero());

    MatrixXd AsB(7, 3);
    MatrixXd StB(2, 3);
    calc->premultiplyByASharpAndSTranspose(AsB, StB, B, false);
    FAST_CHECK_UNARY(AsB.isApprox(A.transpose() * B));
    FAST_CHECK_UNARY(StB.row(0).isApprox(B.row(0)));
    FAST_CHECK_UNARY(StB.row(1).isApprox(B.row(2)));

    MatrixXd N(7, 4);
    N << 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1;
    FAST_CHECK_UNARY(N.isApprox(calc->N()));
  }
}

MatrixXd randM(int m, int n, int r = 0)
{
  if(r <= 0 || r == std::min(m, n))
    return MatrixXd::Random(m, n);
  else
    return MatrixXd::Random(m, r) * MatrixXd::Random(r, n);
}

/** check whether the systems given by cstr and by subs are equivalent.
 * The system are supposed to be feasible.
 * From cstr, we deduce a system A [x;y] = b (1)
 * From subs, we deduce C [y;z] = d and x = E y + F z + g (2)
 * We check that any solution of (1) is a solution of (2) by writing a
 * solution of (1) [x;y] = pinv(A)*b + Na u with Na a base of the nullspace of
 * A and u a vector. For size(u) different value of u, we the rewrite (2) as a
 * system of z only, and verify we can find a solution.
 * To check that any solution of (2) yield a solution of (1), we first solve
 * C [y;z] = d to get [y;z] = pinv(C)*d + Nc v. Then for size(v) value of v, we
 * compute x from x = E y + F z + g, and check that [x;y] is a solution of (1).
 */
void checkEquivalence(const std::vector<std::shared_ptr<constraint::BasicLinearConstraint>> & cstr,
                      Substitutions & subs)
{
  subs.updateSubstitutions();

  VariableVector x(subs.variables());
  VariableVector y;
  VariableVector z(subs.additionalVariables());
  tvm::internal::VariableCountingVector partition(true);
  partition.add(x);
  for(const auto & c : cstr)
  {
    partition.add(c->variables());
  }
  int m0 = 0;
  for(auto & c : cstr)
  {
    m0 += c->size();
    for(const auto & vi : tvm::internal::VariableVectorPartition(c->variables(), partition))
    {
      if(!x.contains(*vi))
      {
        y.add(vi);
      }
    }
  }

  int m1 = 0;
  for(auto & c : subs.additionalConstraints())
  {
    m1 += c->size();
  }

  // Create A and b such that the system given by cstr is A [x;y] = b
  MatrixXd A = MatrixXd::Zero(m0, x.totalSize() + y.totalSize());
  VectorXd b(m0);
  m0 = 0;
  for(const auto & c : cstr)
  {
    auto mi = c->size();
    for(const auto & xi : x.variables())
    {
      if(c->variables().contains(*xi))
      {
        auto rx = xi->getMappingIn(x);
        A.block(m0, rx.start, mi, rx.dim) = c->jacobian(*xi);
      }
    }
    for(const auto & yi : y.variables())
    {
      if(c->variables().contains(*yi))
      {
        auto ry = yi->getMappingIn(y);
        A.block(m0, ry.start + x.totalSize(), mi, ry.dim) = c->jacobian(*yi);
      }
    }
    switch(c->rhs())
    {
      case constraint::RHS::ZERO:
        b.segment(m0, mi).setZero();
        break;
      case constraint::RHS::AS_GIVEN:
        b.segment(m0, mi) = c->e();
        break;
      case constraint::RHS::OPPOSITE:
        b.segment(m0, mi) = -c->e();
        break;
    }
    m0 += mi;
  }

  // Create C and d such that the additional constraint of subs write C [y;z] = d
  MatrixXd C = MatrixXd::Zero(m1, y.totalSize() + z.totalSize());
  VectorXd d(m1);
  m1 = 0;
  for(const auto & c : subs.additionalConstraints())
  {
    auto mi = c->size();
    for(const auto & yi : y.variables())
    {
      if(c->variables().contains(*yi))
      {
        auto ry = yi->getMappingIn(y);
        C.block(m1, ry.start, mi, ry.dim) = c->jacobian(*yi);
      }
    }
    for(const auto & zi : z.variables())
    {
      if(c->variables().contains(*zi))
      {
        auto rz = zi->getMappingIn(z);
        C.block(m1, rz.start + y.totalSize(), mi, rz.dim) = c->jacobian(*zi);
      }
    }
    switch(c->rhs())
    {
      case constraint::RHS::ZERO:
        d.segment(m1, mi).setZero();
        break;
      case constraint::RHS::AS_GIVEN:
        d.segment(m1, mi) = c->e();
        break;
      case constraint::RHS::OPPOSITE:
        d.segment(m1, mi) = -c->e();
        break;
    }
    m1 += mi;
  }

  // Create E, F and g such that the substitutions write x = E y + F z + g
  MatrixXd E = MatrixXd::Zero(x.totalSize(), y.totalSize());
  MatrixXd F = MatrixXd::Zero(x.totalSize(), z.totalSize());
  VectorXd g(x.totalSize());
  for(size_t i = 0; i < x.variables().size(); ++i)
  {
    const auto & f = subs.variableSubstitutions()[i];
    auto rx = x[static_cast<int>(i)]->getMappingIn(x);
    for(const auto & yi : y.variables())
    {
      if(f->variables().contains(*yi))
      {
        auto ry = yi->getMappingIn(y);
        E.block(rx.start, ry.start, rx.dim, ry.dim) = f->jacobian(*yi);
      }
    }
    for(const auto & zi : z.variables())
    {
      if(f->variables().contains(*zi))
      {
        auto rz = zi->getMappingIn(z);
        F.block(rx.start, rz.start, rx.dim, rz.dim) = f->jacobian(*zi);
      }
    }
    g.segment(rx.start, rx.dim) = f->b();
  }

  // Solve A [x;y] = b
  auto svdA = A.jacobiSvd(ComputeFullU | ComputeFullV);
  auto sol0 = svdA.solve(b);                                      // least square solution
  auto r0 = (A * sol0 - b).norm();                                // residual
  MatrixXd Na = svdA.matrixV().rightCols(A.cols() - svdA.rank()); // nullspace of A0
  FAST_CHECK_LE(r0, 1e-9);

  // Solve C [y;z] = d
  VectorXd sol1;
  MatrixXd Nc;
  if(C.size() > 0)
  {
    auto svdC = C.jacobiSvd(ComputeFullU | ComputeFullV);
    sol1 = svdC.solve(d);                                  // least square solution
    Nc = svdC.matrixV().rightCols(C.cols() - svdC.rank()); // nullspace of C
  }
  else
  {
    sol1 = VectorXd::Zero(C.cols());
    Nc = MatrixXd::Identity(C.cols(), C.cols());
  }

  // now check the equivalence:
  // 1 - The solutions of Solve C [y;z] = d give solutions of A [x;y] = b
  for(auto i = 0; i < Nc.cols(); ++i)
  {
    // one solution to the additional constraints
    VectorXd yz = sol1 + Nc * VectorXd::Random(Nc.cols());
    // split it over y and z
    y.value(yz.head(y.totalSize()));
    z.value(yz.tail(z.totalSize()));
    // compute the corresponding x
    subs.updateVariableValues();
    // compute the residual for the current x and y
    VectorXd xy(x.totalSize() + y.totalSize());
    xy.head(x.totalSize()) = x.value();
    xy.tail(y.totalSize()) = y.value();
    double res = (A * xy - b).norm();
    FAST_CHECK_LE(res, 1e-9);
  }

  // 2 - The solutions of A [x;y] = b are such that there we can find z for which
  //   C [y;z] = d and x = E y + F z + g
  MatrixXd M(C.rows() + F.rows(), z.totalSize());
  M.topRows(C.rows()) = C.rightCols(z.totalSize());
  M.bottomRows(F.rows()) = F;
  JacobiSVD<MatrixXd> svdM;
  if(z.totalSize() > 0)
  {
    svdM.compute(M, ComputeThinU | ComputeThinV);
  }
  for(auto i = 0; i < Na.cols(); ++i)
  {
    // one solution to the original system
    VectorXd xy = sol0 + Na * VectorXd::Random(Na.cols());
    // solve C2 z = d - C1 y and F z = x - E y - g
    VectorXd u(C.rows() + F.rows());
    u.head(C.rows()) = d - C.leftCols(y.totalSize()) * xy.tail(y.totalSize());
    u.tail(F.rows()) = xy.head(x.totalSize()) - E * xy.tail(y.totalSize()) - g;
    double res;
    if(z.totalSize() > 0)
    {
      VectorXd s = svdM.solve(u);
      res = (M * s - u).norm();
    }
    else
    {
      res = u.norm();
    }
    FAST_CHECK_LE(res, 1e-9);
  }
}

TEST_CASE("Substitution construction")
{
  using BLC = constraint::BasicLinearConstraint;
  auto eq = constraint::Type::EQUAL;

  VariablePtr x = Space(5).createVariable("x");
  {
    MatrixXd A = randM(5, 5);
    VectorXd b = VectorXd::Random(5);

    auto c = std::shared_ptr<BLC>(new BLC(A, x, b, eq));
    Substitution s(c, x);
    FAST_CHECK_EQ(s.rank(), 5);
    FAST_CHECK_EQ(typeid(*s.calculator()).hash_code(), typeid(GenericCalculator::Impl).hash_code());
  }

  {
    MatrixXd A = MatrixXd::Identity(5, 5);
    VectorXd b = VectorXd::Random(5);

    auto c = std::shared_ptr<BLC>(new BLC(5, x, eq));
    c->A(A, {tvm::internal::MatrixProperties::IDENTITY});
    Substitution s(c, x);
    FAST_CHECK_EQ(s.rank(), 5);
    FAST_CHECK_EQ(typeid(*s.calculator()).hash_code(), typeid(DiagonalCalculator::Impl).hash_code());
  }

  {
    MatrixXd A = randM(5, 5);
    VectorXd b = VectorXd::Random(5);

    auto c = std::shared_ptr<BLC>(new BLC(A, x, b, constraint::Type::LOWER_THAN));
    CHECK_THROWS(Substitution(c, x));

    VariablePtr y = Space(3).createVariable("y");
    CHECK_THROWS(Substitution(c, y));

    auto c2 = std::shared_ptr<BLC>(new BLC(3, y, eq));
    CHECK_THROWS(Substitution({c, c2}, x));

    auto c3 = std::shared_ptr<BLC>(new BLC(5, y, eq));
    CHECK_THROWS(Substitution(c3, y));
  }
}

TEST_CASE("Substitution0")
{
  using BLC = constraint::BasicLinearConstraint;
  auto eq = constraint::Type::EQUAL;

  {
    // solving A x = b using substitutions
    VariablePtr x = Space(5).createVariable("x");
    MatrixXd A = randM(5, 5);
    VectorXd b = VectorXd::Random(5);
    VectorXd x0 = A.colPivHouseholderQr().solve(b);

    auto c = std::shared_ptr<BLC>(new BLC(A, x, b, eq));
    Substitution s(c, x);
    Substitutions subs;
    subs.add(s);
    subs.finalize();
    FAST_CHECK_EQ(subs.variables().size(), 1);
    FAST_CHECK_EQ(subs.variables().front(), x);
    FAST_CHECK_EQ(subs.additionalVariables().size(), 0);
    FAST_CHECK_EQ(subs.additionalConstraints().size(), 0);
    FAST_CHECK_EQ(subs.variableSubstitutions().size(), 1);

    auto f = subs.variableSubstitutions().front();
    FAST_CHECK_EQ(f->variables().totalSize(), 0);
    subs.updateSubstitutions();
    FAST_CHECK_UNARY(f->b().isApprox(x0));
    subs.updateVariableValues();
    FAST_CHECK_UNARY(x->value().isApprox(x0));
  }
}

TEST_CASE("Substitution1")
{
  int m1 = 3;
  int n1 = 4;
  int r1 = 2;
  int m2 = 3;
  int n2 = 3;
  int r2 = 3;
  int m3 = 3;
  int n3 = 6;
  int r3 = 3;
  int m4 = 4;
  int n4 = 4;
  int r4 = 4;
  int m5 = 7;
  int n5 = 8;
  int r5 = 4;
  int m6 = 3;
  int n6 = 3;
  int r6 = 3;
  int l1 = 3;
  int l2 = 7;
  int l3 = 4;
  int l4 = 4;

  VariablePtr x1 = Space(n1).createVariable("x1");
  VariablePtr x2 = Space(n2).createVariable("x2");
  VariablePtr x3 = Space(n3).createVariable("x3");
  VariablePtr x4 = Space(n4).createVariable("x4");
  VariablePtr x5 = Space(n5).createVariable("x5");
  VariablePtr x6 = Space(n6).createVariable("x6");
  VariablePtr y1 = Space(l1).createVariable("y1");
  VariablePtr y2 = Space(l2).createVariable("y2");
  VariablePtr y3 = Space(l3).createVariable("y3");
  VariablePtr y4 = Space(l4).createVariable("y4");

  VectorXd b1 = VectorXd::Random(m1);
  VectorXd b2 = VectorXd::Random(m2);
  VectorXd b3 = VectorXd::Random(m3);
  VectorXd b4 = VectorXd::Random(m4);
  VectorXd b5 = VectorXd::Random(m5);
  VectorXd b6 = VectorXd::Random(m6);

  //                                             | x1 |
  //                                             | x2 |
  // | A11 A12 A13  0   0  A16 B11  0   0   0  | | x3 |   | b1 |
  // |  0  A22  0  A24  0   0   0   0   0   0  | | x4 |   | b2 |
  // |  0   0  A33  0  A35 A36 B31  0   0   0  | | x5 |   | b3 |
  // |  0   0   0  A44 A45  0   0   0   0   0  | | x6 | = | b4 |
  // |  0   0   0   0  A55  0   0  B52  0  B54 | | y1 |   | b5 |
  // |  0   0   0   0   0  A66  0  B62 B63  0  | | y2 |   | b6 |
  //                                             | y3 |
  //                                             | y4 |

  MatrixXd A11 = randM(m1, n1, r1);
  MatrixXd A12 = randM(m1, n2);
  MatrixXd A13 = randM(m1, n3);
  MatrixXd A16 = randM(m1, n6);
  MatrixXd B11 = randM(m1, l1);
  MatrixXd A22 = randM(m2, n2, r2);
  MatrixXd A24 = randM(m2, n4);
  MatrixXd A33 = randM(m3, n3, r3);
  MatrixXd A35 = randM(m3, n5);
  MatrixXd A36 = randM(m3, n6);
  MatrixXd B31 = randM(m3, l1);
  MatrixXd A44 = randM(m4, n4, r4);
  MatrixXd A45 = randM(m4, n5);
  MatrixXd A55 = randM(m5, n5, r5);
  MatrixXd B52 = randM(m5, l2);
  MatrixXd B54 = randM(m5, l4);
  MatrixXd A66 = randM(m6, n6, r6);
  MatrixXd B62 = randM(m6, l2);
  MatrixXd B63 = randM(m6, l3);

  using BLC = constraint::BasicLinearConstraint;
  auto eq = constraint::Type::EQUAL;
  auto c1 = std::shared_ptr<BLC>(new BLC({A11, A12, A13, A16, B11}, {x1, x2, x3, x6, y1}, b1, eq));
  auto c2 = std::shared_ptr<BLC>(new BLC({A22, A24}, {x2, x4}, b2, eq));
  auto c3 = std::shared_ptr<BLC>(new BLC({A33, A35, A36, B31}, {x3, x5, x6, y1}, b3, eq));
  auto c4 = std::shared_ptr<BLC>(new BLC({A44, A45}, {x4, x5}, b4, eq));
  auto c5 = std::shared_ptr<BLC>(new BLC({A55, B52, B54}, {x5, y2, y4}, b5, eq));
  auto c6 = std::shared_ptr<BLC>(new BLC({A66, B62, B63}, {x6, y2, y3}, b6, eq));

  Substitution s1(c1, x1, r1);
  Substitution s2(c2, x2, r2);
  Substitution s3(c3, x3, r3);
  Substitution s4(c4, x4, r4);
  Substitution s5(c5, x5, r5);
  Substitution s6(c6, x6, r6);

  Substitutions subs;
  subs.add(s1);
  subs.add(s2);
  subs.add(s3);
  subs.add(s4);
  subs.add(s5);
  subs.add(s6);

  subs.finalize();

  checkEquivalence({c1, c2, c3, c4, c5, c6}, subs);
}

TEST_CASE("Substitution2")
{
  int m1 = 3;
  int n1 = 4;
  int m2 = 3;
  int n2 = 3;
  int m3 = 3;
  int n3 = 6;
  int m4 = 5;
  int n4 = 4;
  int m5 = 7;
  int n5 = 5;
  int m6 = 3;
  int n6 = 2;
  int m7 = 6;
  int n7 = 7;
  int m8 = 4;
  int n8 = 4;
  int m9 = 4;
  int n9 = 3;
  int l1 = 3;
  int l2 = 7;

  VariablePtr x1 = Space(n1).createVariable("x1");
  VariablePtr x2 = Space(n2).createVariable("x2");
  VariablePtr x3 = Space(n3).createVariable("x3");
  VariablePtr x4 = Space(n4).createVariable("x4");
  VariablePtr x5 = Space(n5).createVariable("x5");
  VariablePtr x6 = Space(n6).createVariable("x6");
  VariablePtr x7 = Space(n7).createVariable("x7");
  VariablePtr x8 = Space(n8).createVariable("x8");
  VariablePtr x9 = Space(n9).createVariable("x9");
  VariablePtr y1 = Space(l1).createVariable("y1");
  VariablePtr y2 = Space(l2).createVariable("y2");

  //                                                 | x1 |
  // |  0   0   0  A14  0   0  A17  0  A19 B11  0  | | x2 |   | b1 |
  // |  0  A22  0   0   0   0   0   0   0  B21 B22 | | x3 |   | b2 |
  // |  0   0   0   0   0  A36  0  A38  0  B31  0  | | x4 |   | b3 |
  // |  0   0   0   0  A45  0   0   0   0  B41  0  | | x5 |   | b4 |
  // |  0   0   0  A54  0   0  A57  0  A59  0   0  | | x6 | = | b5 |
  // |  0   0   0   0   0  A66  0  A68  0  B61 B62 | | x7 |   | b6 |
  // |  0   0  A73  0  A75  0   0   0   0   0  B72 | | x8 |   | b7 |
  // | A81  0   0   0   0   0  A87  0   0  B81  0  | | x9 |   | b8 |
  // |  0   0   0  A94  0   0   0   0   0   0  B92 | | y1 |   | b9 |
  //                                                 | y2 |

  MatrixXd A14 = randM(m1, n4);
  MatrixXd A17 = randM(m1, n7);
  MatrixXd A19 = randM(m1, n9);
  MatrixXd B11 = randM(m1, l1);
  MatrixXd A22 = randM(m2, n2);
  MatrixXd B21 = randM(m2, l1);
  MatrixXd B22 = randM(m2, l2);
  MatrixXd A36 = randM(m3, n6);
  MatrixXd A38 = randM(m3, n8);
  MatrixXd B31 = randM(m3, l1);
  MatrixXd A45 = randM(m4, n5);
  MatrixXd B41 = randM(m4, l1);
  MatrixXd A54 = randM(m5, n4);
  MatrixXd A57 = randM(m5, n7);
  MatrixXd A59 = randM(m5, n9);
  MatrixXd A66 = randM(m6, n6);
  MatrixXd A68 = randM(m6, n8);
  MatrixXd B61 = randM(m6, l1);
  MatrixXd B62 = randM(m6, l2);
  MatrixXd A73 = randM(m7, n3);
  MatrixXd A75 = randM(m7, n5);
  MatrixXd B72 = randM(m7, l2);
  MatrixXd A81 = randM(m8, n1);
  MatrixXd A87 = randM(m8, n7);
  MatrixXd B81 = randM(m8, l1);
  MatrixXd A94 = randM(m9, n4);
  MatrixXd B92 = randM(m9, l2);
  VectorXd b1 = VectorXd::Random(m1);
  VectorXd b2 = VectorXd::Random(m2);
  VectorXd b3 = VectorXd::Random(m3);
  VectorXd b4 = VectorXd::Random(m4);
  VectorXd b5 = VectorXd::Random(m5);
  VectorXd b6 = VectorXd::Random(m6);
  VectorXd b7 = VectorXd::Random(m7);
  VectorXd b8 = VectorXd::Random(m8);
  VectorXd b9 = VectorXd::Random(m9);

  using BLC = constraint::BasicLinearConstraint;
  auto eq = constraint::Type::EQUAL;
  auto c1 = std::shared_ptr<BLC>(new BLC({A14, A17, A19, B11}, {x4, x7, x9, y1}, b1, eq));
  auto c2 = std::shared_ptr<BLC>(new BLC({A22, B21, B22}, {x2, y1, y2}, b2, eq));
  auto c3 = std::shared_ptr<BLC>(new BLC({A36, A38, B31}, {x6, x8, y1}, b3, eq));
  auto c4 = std::shared_ptr<BLC>(new BLC({A45, B41}, {x5, y1}, b4, eq));
  auto c5 = std::shared_ptr<BLC>(new BLC({A54, A57, A59}, {x4, x7, x9}, b5, eq));
  auto c6 = std::shared_ptr<BLC>(new BLC({A66, A68, B62, B61}, {x6, x8, y2, y1}, b6, eq));
  auto c7 = std::shared_ptr<BLC>(new BLC({A73, A75, B72}, {x3, x5, y2}, b7, eq));
  auto c8 = std::shared_ptr<BLC>(new BLC({A81, A87, B81}, {x1, x7, y1}, b8, eq));
  auto c9 = std::shared_ptr<BLC>(new BLC({A94, B92}, {x4, y2}, b9, eq));

  Substitution s1(c1, x9);
  Substitution s2(c2, x2);
  Substitution s3(std::vector<LinearConstraintPtr>{c3, c6}, std::vector<VariablePtr>{x6, x8});
  Substitution s4(c4, x5);
  Substitution s5(c5, x7);
  // no s6: c3 and c6 are merged
  Substitution s7(c7, x3);
  Substitution s8(c8, x1);
  Substitution s9(c9, x4);

  Substitutions subs;
  subs.add(s1);
  subs.add(s2);
  subs.add(s3);
  subs.add(s4);
  subs.add(s5);
  subs.add(s7);
  subs.add(s8);
  subs.add(s9);

  subs.finalize();
  checkEquivalence({c1, c2, c3, c4, c5, c6, c7, c8, c9}, subs);
}

/** Generate a random number r, min <= r < max*/
int randI(int min, int max)
{
  assert(min < max);
  return (rand() % (max - min)) + min;
}

/** Return true with a probability p*/
bool randP(double p)
{
  assert(0 <= p && p <= 1);
  return static_cast<double>(rand()) / RAND_MAX < p;
}

/** Create a random system A [x;y] = b*/
void randomSubstitutions()
{
  using BLC = constraint::BasicLinearConstraint;
  auto eq = constraint::Type::EQUAL;

  int nxmin = 4;
  int nxmax = 16;
  int nymin = 0;
  int nymax = 9;
  int nmin = 2;
  int nmax = 16;
  double nonFullRankP = 0.33;
  int nx = randI(nxmin, nxmax); // number of x variables
  int ny = randI(nymin, nymax); // number of y variables
  double px = 0.4;              // probability to have a non-zero off-diagonal block in the 'x part' of A
  double py = 0.3;              // probability to have a non-zero block in the 'y part' of B

  std::vector<VariablePtr> x;
  std::vector<VariablePtr> y;
  std::vector<int> m;
  std::vector<int> n;
  std::vector<int> r;
  std::vector<int> l;
  std::vector<std::shared_ptr<BLC>> cstr;

  bool fullRankSystem = false;

  // We restart from scratch when we do not get a system that is full rank
  while(!fullRankSystem)
  {
    int ma = 0;
    int na = 0;
    x.clear();
    y.clear();
    m.clear();
    n.clear();
    r.clear();
    l.clear();
    cstr.clear();
    for(int i = 0; i < nx; ++i)
    {
      std::stringstream ss;
      ss << "x" << i;
      int ni = randI(nmin, nmax);
      int mi = 1 + randI(1, ni);
      n.push_back(ni);
      m.push_back(mi);
      x.push_back(Space(ni).createVariable(ss.str()));
      if(randP(nonFullRankP))
      {
        r.push_back(randI(mi / 2, mi + 1));
      }
      else
      {
        r.push_back(mi);
      }
      ma += mi;
      na += ni;
    }
    for(int i = 0; i < ny; ++i)
    {
      std::stringstream ss;
      ss << "y" << i;
      int ni = randI(nmin, nmax);
      l.push_back(ni);
      y.push_back(Space(ni).createVariable(ss.str()));
      na += l.back();
    }

    for(size_t i = 0; i < static_cast<size_t>(nx); ++i)
    {
      std::vector<VariablePtr> v;
      std::vector<MatrixXd> M;
      std::vector<MatrixConstRef> Mr;
      for(size_t j = 0; j < static_cast<size_t>(nx); ++j)
      {
        if(i == j)
        {
          M.push_back(randM(m[i], n[j], r[i]));
          v.push_back(x[j]);
        }
        else
        {
          if(randP(px))
          {
            M.push_back(randM(m[i], n[j]));
            v.push_back(x[j]);
          }
        }
      }
      for(size_t j = 0; j < static_cast<size_t>(ny); ++j)
      {
        if(randP(py))
        {
          M.push_back(randM(m[i], l[j]));
          v.push_back(y[j]);
        }
      }
      for(const auto & Mi : M)
      {
        Mr.push_back(Mi);
      }
      VectorXd b = VectorXd::Random(m[i]);
      auto c = std::shared_ptr<BLC>(new BLC(Mr, v, b, eq));
      cstr.push_back(c);
    }
    // check if the system is feasible
    MatrixXd A = MatrixXd::Zero(ma, na);
    int ms = 0;
    for(const auto & c : cstr)
    {
      int ns = 0;
      int mi = c->size();
      for(const auto & xi : x)
      {
        int ni = xi->size();
        if(c->variables().contains(*xi))
        {
          A.block(ms, ns, mi, ni) = c->jacobian(*xi);
        }
        ns += ni;
      }
      for(const auto & yi : y)
      {
        int ni = yi->size();
        if(c->variables().contains(*yi))
        {
          A.block(ms, ns, mi, ni) = c->jacobian(*yi);
        }
        ns += ni;
      }
      ms += mi;
    }
    auto svd = A.jacobiSvd();
    fullRankSystem = svd.rank() == A.rows();
  }

  Substitutions subs;
  for(size_t i = 0; i < cstr.size(); ++i)
  {
    Substitution s(cstr[i], x[i], r[i]);
    subs.add(s);
  }
  subs.finalize();
  checkEquivalence(cstr, subs);
}

TEST_CASE("Random substitutions")
{
  // In some very rare instance, this test can fail because of rank issues
  // during the qr decomposition of GenericCalculator. This is because of the
  // heuristic taken to get the rank of a group of constraints. This is not to
  // be regarded as an issue.
  for(int i = 0; i < 10; ++i)
  {
    randomSubstitutions();
  }
}

TEST_CASE("Substitution with subvariables")
{
  VariablePtr x = Space(8).createVariable("x");
  VariablePtr x1 = x->subvariable(3, "x1", 0);
  VariablePtr x2 = x->subvariable(5, "x2", 3);

  MatrixXd M = randM(8, 8);
  VectorXd r = VectorXd::Random(8);

  auto A1 = M.topRows(5);
  auto A2 = M.bottomRows(3);
  auto b1 = r.head(5);
  auto b2 = r.tail(3);

  using BLC = constraint::BasicLinearConstraint;
  auto eq = constraint::Type::EQUAL;
  VectorXd x0 = M.colPivHouseholderQr().solve(r);

  auto c = std::shared_ptr<BLC>(new BLC(A1, x, b1, eq));
  Substitution s(c, x2);
  Substitutions subs;
  subs.add(s);
  subs.finalize();
  FAST_CHECK_EQ(subs.variables().size(), 1);
  FAST_CHECK_EQ(*subs.variables().front(), *x2);
  FAST_CHECK_EQ(subs.additionalVariables().size(), 0);
  FAST_CHECK_EQ(subs.otherVariables().size(), 1);
  FAST_CHECK_EQ(*subs.otherVariables().front(), *x1);
  FAST_CHECK_EQ(subs.additionalConstraints().size(), 0);
  FAST_CHECK_EQ(subs.variableSubstitutions().size(), 1);

  auto f = subs.variableSubstitutions().front();
  FAST_CHECK_EQ(f->variables().numberOfVariables(), 1);
  FAST_CHECK_EQ(*f->variables()[0], *x1);
  subs.updateSubstitutions();

  MatrixXd C = A2.leftCols(3) + A2.rightCols(5) * f->jacobian(*x1);
  VectorXd d = b2 - A2.rightCols(5) * f->b();
  VectorXd y = C.colPivHouseholderQr().solve(d);
  x1 << y;

  subs.updateVariableValues();
  FAST_CHECK_UNARY(x->value().isApprox(x0));
}


/**
 * @file tests/kfn_test.cpp
 *
 * Tests for KFN (k-furthest-neighbors).
 *
 * 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/neighbor_search/neighbor_search.hpp>
#include <mlpack/core/tree/cover_tree.hpp>
#include <boost/test/unit_test.hpp>
#include "test_tools.hpp"

using namespace mlpack;
using namespace mlpack::neighbor;
using namespace mlpack::tree;
using namespace mlpack::metric;
using namespace mlpack::bound;

BOOST_AUTO_TEST_SUITE(KFNTest);

/**
 * Simple furthest-neighbors test with small, synthetic dataset.  This is an
 * exhaustive test, which checks that each method for performing the calculation
 * (dual-tree, single-tree, naive) produces the correct results.  An
 * eleven-point dataset and the ten furthest neighbors are taken.  The dataset
 * is in one dimension for simplicity -- the correct functionality of distance
 * functions is not tested here.
 */
BOOST_AUTO_TEST_CASE(ExhaustiveSyntheticTest)
{
  // Set up our data.
  arma::mat data(1, 11);
  data[0] = 0.05; // Row addressing is unnecessary (they are all 0).
  data[1] = 0.35;
  data[2] = 0.15;
  data[3] = 1.25;
  data[4] = 5.05;
  data[5] = -0.22;
  data[6] = -2.00;
  data[7] = -1.30;
  data[8] = 0.45;
  data[9] = 0.90;
  data[10] = 1.00;

  typedef BinarySpaceTree<EuclideanDistance,
      NeighborSearchStat<FurthestNeighborSort>, arma::mat> TreeType;

  // We will loop through three times, one for each method of performing the
  // calculation.  We'll always use 10 neighbors, so set that parameter.
  std::vector<size_t> oldFromNew;
  std::vector<size_t> newFromOld;
  TreeType tree(data, oldFromNew, newFromOld, 1);
  KFN kfn(std::move(tree));

  for (int i = 0; i < 3; ++i)
  {
    switch (i)
    {
      case 0: // Use the dual-tree method.
        kfn.SearchMode() = DUAL_TREE_MODE;
        break;
      case 1: // Use the single-tree method.
        kfn.SearchMode() = SINGLE_TREE_MODE;
        break;
      case 2: // Use the naive method.
        kfn.SearchMode() = NAIVE_MODE;
        break;
    }

    // Now perform the actual calculation.
    arma::Mat<size_t> neighbors;
    arma::mat distances;
    kfn.Search(10, neighbors, distances);

    // Now the exhaustive check for correctness.  This will be long.  We must
    // also remember that the distances returned are squared distances.  As a
    // result, distance comparisons are written out as (distance * distance) for
    // readability.

    // Neighbors of point 0.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[0]), newFromOld[2]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[0]), 0.10, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[0]), newFromOld[5]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[0]), 0.27, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[0]), newFromOld[1]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[0]), 0.30, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[0]), newFromOld[8]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[0]), 0.40, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[0]), newFromOld[9]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[0]), 0.85, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[0]), newFromOld[10]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[0]), 0.95, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[0]), newFromOld[3]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[0]), 1.20, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[0]), newFromOld[7]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[0]), 1.35, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[0]), newFromOld[6]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[0]), 2.05, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[0]), newFromOld[4]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[0]), 5.00, 1e-5);

    // Neighbors of point 1.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[1]), newFromOld[8]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[1]), 0.10, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[1]), newFromOld[2]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[1]), 0.20, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[1]), newFromOld[0]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[1]), 0.30, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[1]), newFromOld[9]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[1]), 0.55, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[1]), newFromOld[5]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[1]), 0.57, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[1]), newFromOld[10]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[1]), 0.65, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[1]), newFromOld[3]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[1]), 0.90, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[1]), newFromOld[7]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[1]), 1.65, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[1]), newFromOld[6]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[1]), 2.35, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[1]), newFromOld[4]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[1]), 4.70, 1e-5);

    // Neighbors of point 2.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[2]), newFromOld[0]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[2]), 0.10, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[2]), newFromOld[1]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[2]), 0.20, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[2]), newFromOld[8]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[2]), 0.30, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[2]), newFromOld[5]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[2]), 0.37, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[2]), newFromOld[9]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[2]), 0.75, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[2]), newFromOld[10]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[2]), 0.85, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[2]), newFromOld[3]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[2]), 1.10, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[2]), newFromOld[7]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[2]), 1.45, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[2]), newFromOld[6]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[2]), 2.15, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[2]), newFromOld[4]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[2]), 4.90, 1e-5);

    // Neighbors of point 3.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[3]), newFromOld[10]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[3]), 0.25, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[3]), newFromOld[9]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[3]), 0.35, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[3]), newFromOld[8]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[3]), 0.80, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[3]), newFromOld[1]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[3]), 0.90, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[3]), newFromOld[2]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[3]), 1.10, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[3]), newFromOld[0]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[3]), 1.20, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[3]), newFromOld[5]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[3]), 1.47, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[3]), newFromOld[7]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[3]), 2.55, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[3]), newFromOld[6]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[3]), 3.25, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[3]), newFromOld[4]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[3]), 3.80, 1e-5);

    // Neighbors of point 4.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[4]), newFromOld[3]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[4]), 3.80, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[4]), newFromOld[10]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[4]), 4.05, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[4]), newFromOld[9]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[4]), 4.15, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[4]), newFromOld[8]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[4]), 4.60, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[4]), newFromOld[1]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[4]), 4.70, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[4]), newFromOld[2]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[4]), 4.90, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[4]), newFromOld[0]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[4]), 5.00, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[4]), newFromOld[5]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[4]), 5.27, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[4]), newFromOld[7]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[4]), 6.35, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[4]), newFromOld[6]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[4]), 7.05, 1e-5);

    // Neighbors of point 5.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[5]), newFromOld[0]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[5]), 0.27, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[5]), newFromOld[2]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[5]), 0.37, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[5]), newFromOld[1]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[5]), 0.57, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[5]), newFromOld[8]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[5]), 0.67, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[5]), newFromOld[7]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[5]), 1.08, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[5]), newFromOld[9]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[5]), 1.12, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[5]), newFromOld[10]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[5]), 1.22, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[5]), newFromOld[3]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[5]), 1.47, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[5]), newFromOld[6]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[5]), 1.78, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[5]), newFromOld[4]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[5]), 5.27, 1e-5);

    // Neighbors of point 6.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[6]), newFromOld[7]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[6]), 0.70, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[6]), newFromOld[5]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[6]), 1.78, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[6]), newFromOld[0]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[6]), 2.05, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[6]), newFromOld[2]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[6]), 2.15, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[6]), newFromOld[1]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[6]), 2.35, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[6]), newFromOld[8]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[6]), 2.45, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[6]), newFromOld[9]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[6]), 2.90, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[6]), newFromOld[10]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[6]), 3.00, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[6]), newFromOld[3]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[6]), 3.25, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[6]), newFromOld[4]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[6]), 7.05, 1e-5);

    // Neighbors of point 7.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[7]), newFromOld[6]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[7]), 0.70, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[7]), newFromOld[5]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[7]), 1.08, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[7]), newFromOld[0]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[7]), 1.35, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[7]), newFromOld[2]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[7]), 1.45, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[7]), newFromOld[1]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[7]), 1.65, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[7]), newFromOld[8]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[7]), 1.75, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[7]), newFromOld[9]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[7]), 2.20, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[7]), newFromOld[10]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[7]), 2.30, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[7]), newFromOld[3]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[7]), 2.55, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[7]), newFromOld[4]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[7]), 6.35, 1e-5);

    // Neighbors of point 8.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[8]), newFromOld[1]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[8]), 0.10, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[8]), newFromOld[2]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[8]), 0.30, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[8]), newFromOld[0]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[8]), 0.40, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[8]), newFromOld[9]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[8]), 0.45, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[8]), newFromOld[10]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[8]), 0.55, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[8]), newFromOld[5]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[8]), 0.67, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[8]), newFromOld[3]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[8]), 0.80, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[8]), newFromOld[7]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[8]), 1.75, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[8]), newFromOld[6]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[8]), 2.45, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[8]), newFromOld[4]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[8]), 4.60, 1e-5);

    // Neighbors of point 9.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[9]), newFromOld[10]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[9]), 0.10, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[9]), newFromOld[3]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[9]), 0.35, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[9]), newFromOld[8]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[9]), 0.45, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[9]), newFromOld[1]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[9]), 0.55, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[9]), newFromOld[2]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[9]), 0.75, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[9]), newFromOld[0]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[9]), 0.85, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[9]), newFromOld[5]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[9]), 1.12, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[9]), newFromOld[7]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[9]), 2.20, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[9]), newFromOld[6]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[9]), 2.90, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[9]), newFromOld[4]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[9]), 4.15, 1e-5);

    // Neighbors of point 10.
    BOOST_REQUIRE_EQUAL(neighbors(9, newFromOld[10]), newFromOld[9]);
    BOOST_REQUIRE_CLOSE(distances(9, newFromOld[10]), 0.10, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(8, newFromOld[10]), newFromOld[3]);
    BOOST_REQUIRE_CLOSE(distances(8, newFromOld[10]), 0.25, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(7, newFromOld[10]), newFromOld[8]);
    BOOST_REQUIRE_CLOSE(distances(7, newFromOld[10]), 0.55, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(6, newFromOld[10]), newFromOld[1]);
    BOOST_REQUIRE_CLOSE(distances(6, newFromOld[10]), 0.65, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(5, newFromOld[10]), newFromOld[2]);
    BOOST_REQUIRE_CLOSE(distances(5, newFromOld[10]), 0.85, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(4, newFromOld[10]), newFromOld[0]);
    BOOST_REQUIRE_CLOSE(distances(4, newFromOld[10]), 0.95, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(3, newFromOld[10]), newFromOld[5]);
    BOOST_REQUIRE_CLOSE(distances(3, newFromOld[10]), 1.22, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(2, newFromOld[10]), newFromOld[7]);
    BOOST_REQUIRE_CLOSE(distances(2, newFromOld[10]), 2.30, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(1, newFromOld[10]), newFromOld[6]);
    BOOST_REQUIRE_CLOSE(distances(1, newFromOld[10]), 3.00, 1e-5);
    BOOST_REQUIRE_EQUAL(neighbors(0, newFromOld[10]), newFromOld[4]);
    BOOST_REQUIRE_CLOSE(distances(0, newFromOld[10]), 4.05, 1e-5);
  }
}

/**
 * Test the dual-tree furthest-neighbors method with the naive method.  This
 * uses both a query and reference dataset.
 *
 * Errors are produced if the results are not identical.
 */
BOOST_AUTO_TEST_CASE(DualTreeVsNaive1)
{
  arma::mat dataset;

  // Hard-coded filename: bad?
  if (!data::Load("test_data_3_1000.csv", dataset))
    BOOST_FAIL("Cannot load test dataset test_data_3_1000.csv!");

  KFN kfn(dataset);

  KFN naive(dataset, NAIVE_MODE);

  arma::Mat<size_t> neighborsTree;
  arma::mat distancesTree;
  kfn.Search(dataset, 15, neighborsTree, distancesTree);

  arma::Mat<size_t> neighborsNaive;
  arma::mat distancesNaive;
  naive.Search(dataset, 15, neighborsNaive, distancesNaive);

  for (size_t i = 0; i < neighborsTree.n_elem; ++i)
  {
    BOOST_REQUIRE(neighborsTree[i] == neighborsNaive[i]);
    BOOST_REQUIRE_CLOSE(distancesTree[i], distancesNaive[i], 1e-5);
  }
}

/**
 * Test the dual-tree furthest-neighbors method with the naive method.  This
 * uses only a reference dataset.
 *
 * Errors are produced if the results are not identical.
 */
BOOST_AUTO_TEST_CASE(DualTreeVsNaive2)
{
  arma::mat dataset;

  // Hard-coded filename: bad?
  // Code duplication: also bad!
  if (!data::Load("test_data_3_1000.csv", dataset))
    BOOST_FAIL("Cannot load test dataset test_data_3_1000.csv!");

  KFN kfn(dataset);

  KFN naive(dataset, NAIVE_MODE);

  arma::Mat<size_t> neighborsTree;
  arma::mat distancesTree;
  kfn.Search(15, neighborsTree, distancesTree);

  arma::Mat<size_t> neighborsNaive;
  arma::mat distancesNaive;
  naive.Search(15, neighborsNaive, distancesNaive);

  for (size_t i = 0; i < neighborsTree.n_elem; ++i)
  {
    BOOST_REQUIRE_EQUAL(neighborsTree[i], neighborsNaive[i]);
    BOOST_REQUIRE_CLOSE(distancesTree[i], distancesNaive[i], 1e-5);
  }
}

/**
 * Test the single-tree furthest-neighbors method with the naive method.  This
 * uses only a reference dataset.
 *
 * Errors are produced if the results are not identical.
 */
BOOST_AUTO_TEST_CASE(SingleTreeVsNaive)
{
  arma::mat dataset;

  // Hard-coded filename: bad!
  // Code duplication: also bad!
  if (!data::Load("test_data_3_1000.csv", dataset))
    BOOST_FAIL("Cannot load test dataset test_data_3_1000.csv!");

  KFN kfn(dataset, SINGLE_TREE_MODE);

  KFN naive(dataset, NAIVE_MODE);

  arma::Mat<size_t> neighborsTree;
  arma::mat distancesTree;
  kfn.Search(15, neighborsTree, distancesTree);

  arma::Mat<size_t> neighborsNaive;
  arma::mat distancesNaive;
  naive.Search(15, neighborsNaive, distancesNaive);

  for (size_t i = 0; i < neighborsTree.n_elem; ++i)
  {
    BOOST_REQUIRE_EQUAL(neighborsTree[i], neighborsNaive[i]);
    BOOST_REQUIRE_CLOSE(distancesTree[i], distancesNaive[i], 1e-5);
  }
}

/**
 * Test the cover tree single-tree furthest-neighbors method against the naive
 * method.  This uses only a random reference dataset.
 *
 * Errors are produced if the results are not identical.
 */
BOOST_AUTO_TEST_CASE(SingleCoverTreeTest)
{
  arma::mat data;
  data.randu(75, 1000); // 75 dimensional, 1000 points.

  // This depends on the cover tree not mapping points.
  CoverTree<LMetric<2>, NeighborSearchStat<FurthestNeighborSort>, arma::mat,
      FirstPointIsRoot> tree(data);

  NeighborSearch<FurthestNeighborSort, LMetric<2>, arma::mat, StandardCoverTree>
      coverTreeSearch(std::move(tree), SINGLE_TREE_MODE);

  KFN naive(data, NAIVE_MODE);

  arma::Mat<size_t> coverTreeNeighbors;
  arma::mat coverTreeDistances;
  coverTreeSearch.Search(data, 15, coverTreeNeighbors, coverTreeDistances);

  arma::Mat<size_t> naiveNeighbors;
  arma::mat naiveDistances;
  naive.Search(data, 15, naiveNeighbors, naiveDistances);

  for (size_t i = 0; i < coverTreeNeighbors.n_elem; ++i)
  {
    BOOST_REQUIRE_EQUAL(coverTreeNeighbors[i], naiveNeighbors[i]);
    BOOST_REQUIRE_CLOSE(coverTreeDistances[i], naiveDistances[i], 1e-5);
  }
}

/**
 * Test the cover tree dual-tree furthest neighbors method against the naive
 * method.
 */
BOOST_AUTO_TEST_CASE(DualCoverTreeTest)
{
  arma::mat dataset;
  data::Load("test_data_3_1000.csv", dataset);

  KFN tree(dataset);

  arma::Mat<size_t> kdNeighbors;
  arma::mat kdDistances;
  tree.Search(dataset, 5, kdNeighbors, kdDistances);

  typedef CoverTree<LMetric<2, true>, NeighborSearchStat<FurthestNeighborSort>,
      arma::mat, FirstPointIsRoot> TreeType;

  TreeType referenceTree(dataset);

  NeighborSearch<FurthestNeighborSort, LMetric<2, true>, arma::mat,
      StandardCoverTree> coverTreeSearch(std::move(referenceTree));

  arma::Mat<size_t> coverNeighbors;
  arma::mat coverDistances;
  coverTreeSearch.Search(dataset, 5, coverNeighbors, coverDistances);

  for (size_t i = 0; i < coverNeighbors.n_elem; ++i)
  {
    BOOST_REQUIRE_EQUAL(coverNeighbors(i), kdNeighbors(i));
    BOOST_REQUIRE_CLOSE(coverDistances(i), kdDistances(i), 1e-5);
  }
}

/**
 * Test the ball tree single-tree furthest-neighbors method against the naive
 * method.  This uses only a random reference dataset.
 *
 * Errors are produced if the results are not identical.
 */
BOOST_AUTO_TEST_CASE(SingleBallTreeTest)
{
  arma::mat data;
  data.randu(75, 1000); // 75 dimensional, 1000 points.

  typedef BallTree<EuclideanDistance, NeighborSearchStat<FurthestNeighborSort>,
      arma::mat> TreeType;
  TreeType tree(data);

  KFN naive(tree.Dataset(), NAIVE_MODE);

  // BinarySpaceTree modifies data. Use modified data to maintain the
  // correspondence between points in the dataset for both methods. The order of
  // query points in both methods should be same.
  NeighborSearch<FurthestNeighborSort, LMetric<2>, arma::mat, BallTree>
      ballTreeSearch(std::move(tree), SINGLE_TREE_MODE);

  arma::Mat<size_t> ballTreeNeighbors;
  arma::mat ballTreeDistances;
  ballTreeSearch.Search(15, ballTreeNeighbors, ballTreeDistances);

  arma::Mat<size_t> naiveNeighbors;
  arma::mat naiveDistances;
  naive.Search(15, naiveNeighbors, naiveDistances);

  for (size_t i = 0; i < ballTreeNeighbors.n_elem; ++i)
  {
    BOOST_REQUIRE_EQUAL(ballTreeNeighbors[i], naiveNeighbors[i]);
    BOOST_REQUIRE_CLOSE(ballTreeDistances[i], naiveDistances[i], 1e-5);
  }
}

/**
 * Test the ball tree dual-tree furthest neighbors method against the naive
 * method.
 */
BOOST_AUTO_TEST_CASE(DualBallTreeTest)
{
  arma::mat dataset;
  data::Load("test_data_3_1000.csv", dataset);

  KFN tree(dataset);

  arma::Mat<size_t> kdNeighbors;
  arma::mat kdDistances;
  tree.Search(5, kdNeighbors, kdDistances);

  NeighborSearch<FurthestNeighborSort, LMetric<2, true>, arma::mat, BallTree>
      ballTreeSearch(dataset);

  arma::Mat<size_t> ballNeighbors;
  arma::mat ballDistances;
  ballTreeSearch.Search(5, ballNeighbors, ballDistances);

  for (size_t i = 0; i < ballNeighbors.n_elem; ++i)
  {
    BOOST_REQUIRE_EQUAL(ballNeighbors(i), kdNeighbors(i));
    BOOST_REQUIRE_CLOSE(ballDistances(i), kdDistances(i), 1e-5);
  }
}

BOOST_AUTO_TEST_SUITE_END();


#include <iostream>
#include <fstream>

#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

#include <imgproc/gradient_adapter.hpp>
#include <imgproc/derivative_gradient.hpp>
#include <imgproc/susan.hpp>
#include <imgproc/rcmg.hpp>
#include <imgproc/quadratureG2.hpp>
#include <imgproc/quadratureS.hpp>
#include <imgproc/quadratureSF.hpp>
#include <imgproc/quadratureLGF.hpp>
#include <imgproc/pc_sqf.hpp>
#include <imgproc/pc_lgf.hpp>
#include <imgproc/pc_matlab.hpp>
#include <imgproc/laplace.hpp>

#include <boost/filesystem.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/format.hpp>

using namespace lsfm;
using namespace std;
namespace fs = boost::filesystem;

constexpr int ENTRY_RGB = 2;

constexpr int runs = 10;

template<class GT, class MT, class FT>
struct Entry {
    Entry() {}

    Entry(const cv::Ptr<GradientI<uchar,GT,MT,FT>>& g, const std::string& n, int f = 0)
        : gradient(g), name(n), time(0), images(0), flags(f) {}

    cv::Ptr<GradientI<uchar, GT, MT, FT>> gradient;
    std::string name;
    int64 time;
    int images, flags;

    inline bool rgb() const {
        return flags & ENTRY_RGB;
    }

    inline void process(const cv::Mat src) {
        gradient->process(src);
        gradient->magnitude();
        cv::Mat tmp;
        int64 start;
        for (int i = 0; i != runs; ++i) {
            start = cv::getTickCount();
            gradient->process(src);
            tmp = gradient->magnitude();
            time += cv::getTickCount() - start;
            ++images;
        }
    }
};

void parseFolder(const fs::path &folder, std::vector<fs::path> &files) {
    fs::directory_iterator end_iter;
    for_each(fs::directory_iterator(folder), fs::directory_iterator(), [&files](const fs::path& file) {
        if (fs::is_regular_file(file))
        {
            std::string ext = file.extension().generic_string();
            boost::algorithm::to_lower(ext);
            if (ext == ".jpg" || ext == ".png") {
                files.push_back(file);
            }
        }
        if (fs::is_directory(file))
            parseFolder(file, files);
    });
}

template<class GT, class MT, class DT>
void processPath(std::vector<Entry<GT,MT,DT>> &entries, const std::pair<fs::path, std::string> &path) {
    std::cout << "processing " << path.first << std::endl;
    std::vector<fs::path> files;
    parseFolder(path.first, files);

    for_each(entries.begin(), entries.end(), [&](Entry<GT, MT, DT> &e) {
        e.time = 0;
        e.images = 0;
    });

    std::for_each(files.begin(), files.end(), [&](const fs::path& file) {
        cv::Mat rgb = cv::imread(file.generic_string());
        if (rgb.empty())
        {
            cout << "Can not open " << file.generic_string() << endl;
        }
        std::cout << file << std::endl;
        cv::Mat src;
        cv::cvtColor(rgb, src, CV_BGR2GRAY);

        for_each(entries.begin(), entries.end(), [&](Entry<GT, MT, DT> &e) {
            e.process(e.rgb() ? rgb : src);
        });
    });
}

int main(int argc, char** argv)
{
    char c;
    std::cin >> c;

    std::vector<std::pair<fs::path, std::string>> sets;
    //sets.push_back(std::pair<fs::path, std::string>("../../images/Selection", "Selection"));
    //sets.push_back(std::pair<fs::path, std::string>("../../images/BSDS500", "BSDS500"));
    sets.push_back(std::pair<fs::path, std::string>("../../images/MDB/MiddEval3-Q", "MDB-Q"));
    //sets.push_back(std::pair<fs::path, std::string>("../../images/MDB/MiddEval3-H", "MDB-H"));
    //sets.push_back(std::pair<fs::path, std::string>("../../images/MDB/MiddEval3-F", "MDB-F"));

    std::vector<Entry<short, int, float>> gradI;
    /*gradI.push_back(Entry<short, int, float>(new DerivativeGradient<uchar, short, int, float, RobertsDerivative<uchar, short>, QuadraticMagnitude>, "Roberts (2x2)"));
    gradI.push_back(Entry<short, int, float>(new DerivativeGradient<uchar, short, int, float, PrewittDerivative<uchar, short>, QuadraticMagnitude>, "Prewitt (3x3)"));
    gradI.push_back(Entry<short, int, float>(new DerivativeGradient<uchar, short, int, float, ScharrDerivative<uchar, short>, QuadraticMagnitude>, "Scharr (3x3)"));

    gradI.push_back(Entry<short, int, float>(new DerivativeGradient<uchar, short, int, float, SobelDerivative, QuadraticMagnitude>, "Sobel (3x3)"));
    gradI.push_back(Entry<short, int, float>(new DerivativeGradient<uchar, short, int, float, SobelDerivative, QuadraticMagnitude>({ NV("grad_kernel_size",5) }), "Sobel (5x5)"));
    gradI.push_back(Entry<short, int, float>(new DerivativeGradient<uchar, short, int, float, SobelDerivative, QuadraticMagnitude>({ NV("grad_kernel_size",7) }), "Sobel (7x7)"));
    gradI.push_back(Entry<short, int, float>(new DerivativeGradient<uchar, short, int, float, SobelDerivative, QuadraticMagnitude>({ NV("grad_kernel_size",9) }), "Sobel (9x9)"));*/
    gradI.push_back(Entry<short, int, float>(new DerivativeGradient<uchar, short, int, float, SobelDerivative, QuadraticMagnitude>({ NV("grad_kernel_size",11) }), "Sobel (11x11)"));
    gradI.push_back(Entry<short, int, float>(new DerivativeGradient<uchar, short, int, float, SobelDerivative, QuadraticMagnitude>({ NV("grad_kernel_size",13) }), "Sobel (13x13)"));
    gradI.push_back(Entry<short, int, float>(new DerivativeGradient<uchar, short, int, float, SobelDerivative, QuadraticMagnitude>({ NV("grad_kernel_size",15) }), "Sobel (15x15)"));

    /*gradI.push_back(Entry<short, int, float>(new SusanGradient<short, int>, "Susan (37)"));
    gradI.push_back(Entry<short, int, float>(new SusanGradient<short, int>(20, true), "Susan (3x3)"));

    gradI.push_back(Entry<short, int, float>(new RCMGradient<uchar, 1, short, int>(3, 1), "RMG (3x3)"));
    gradI.push_back(Entry<short, int, float>(new RCMGradient<uchar, 1, short, int>(5, 3), "RMG (5x5)"));
    gradI.push_back(Entry<short, int, float>(new RCMGradient<uchar, 3, short, int>(3, 1), "RCMG (3x3)", ENTRY_RGB));*/

    std::vector<Entry<float, float, float>> gradF;
    /*gradF.push_back(Entry<float, float, float>(new DerivativeGradient<uchar, float, float, float, GaussianDerivative, Magnitude>({ NV("grad_kernel_size",3), NV("grad_range",1.5) }), "Gauss (3x3)"));
    gradF.push_back(Entry<float, float, float>(new DerivativeGradient<uchar, float, float, float, GaussianDerivative, Magnitude>({ NV("grad_kernel_size",5), NV("grad_range",2.3) }), "Gauss (5x5)"));
    gradF.push_back(Entry<float, float, float>(new DerivativeGradient<uchar, float, float, float, GaussianDerivative, Magnitude>({ NV("grad_kernel_size",7), NV("grad_range",3.0) }), "Gauss (7x7)"));
    gradF.push_back(Entry<float, float, float>(new DerivativeGradient<uchar, float, float, float, GaussianDerivative, Magnitude>({ NV("grad_kernel_size",9), NV("grad_range",3.5) }), "Gauss (9x9)"));*/
    gradF.push_back(Entry<float, float, float>(new DerivativeGradient<uchar, float, float, float, GaussianDerivative, Magnitude>({ NV("grad_kernel_size",11), NV("grad_range",4) }), "Gauss (11x11)"));
    gradF.push_back(Entry<float, float, float>(new DerivativeGradient<uchar, float, float, float, GaussianDerivative, Magnitude>({ NV("grad_kernel_size",13), NV("grad_range",4.5) }), "Gauss (13x13)"));
    gradF.push_back(Entry<float, float, float>(new DerivativeGradient<uchar, float, float, float, GaussianDerivative, Magnitude>({ NV("grad_kernel_size",15), NV("grad_range",5) }), "Gauss (15x15)"));

    /*gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureG2<uchar, float, PolarCV>>({ NV("grad_kernel_size",3), NV("grad_kernel_spacing",1.24008) }), "QF_StG (3x3)"));
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureG2<uchar, float, PolarCV>>({ NV("grad_kernel_size", 5), NV("grad_kernel_spacing", 1.008) }), "QF_StG (5x5)"));
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureG2<uchar, float, PolarCV>>({ NV("grad_kernel_size", 7), NV("grad_kernel_spacing", 0.873226) }), "QF_StG (7x7)"));
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureG2<uchar, float, PolarCV>>({ NV("grad_kernel_size", 9), NV("grad_kernel_spacing", 0.781854) }), "QF_StG (9x9)"));*/
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureG2<uchar, float, PolarCV>>({ NV("grad_kernel_size", 11), NV("grad_kernel_spacing", 0.7) }), "QF_StG (11x11)"));
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureG2<uchar, float, PolarCV>>({ NV("grad_kernel_size", 13), NV("grad_kernel_spacing", 0.65) }), "QF_StG (13x13)"));
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureG2<uchar, float, PolarCV>>({ NV("grad_kernel_size", 15), NV("grad_kernel_spacing", 0.6) }), "QF_StG (15x15)"));

    /*gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureS<uchar, float, float, PolarCV>>({ NV("grad_scale", 1), NV("grad_muls", 2), NV("grad_kernel_size", 3), NV("grad_kernel_spacing", 1.2) }), "SQF PO (3x3)"));
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureS<uchar, float, float, PolarCV>>({ NV("grad_scale", 1), NV("grad_muls", 2), NV("grad_kernel_size", 5), NV("grad_kernel_spacing", 1.2) }), "SQF PO (5x5)"));
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureS<uchar, float, float, PolarCV>>({ NV("grad_scale", 1), NV("grad_muls", 2), NV("grad_kernel_size", 7), NV("grad_kernel_spacing", 1.2) }), "SQF PO (7x7)"));
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureS<uchar, float, float, PolarCV>>({ NV("grad_scale", 1), NV("grad_muls", 2), NV("grad_kernel_size", 9), NV("grad_kernel_spacing", 1.2) }), "SQF PO (9x9)"));*/
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureS<uchar, float, float, PolarCV>>({ NV("grad_scale", 1), NV("grad_muls", 2), NV("grad_kernel_size", 11), NV("grad_kernel_spacing", 1.2) }), "SQF PO (11x11)"));
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureS<uchar, float, float, PolarCV>>({ NV("grad_scale", 1), NV("grad_muls", 2), NV("grad_kernel_size", 13), NV("grad_kernel_spacing", 1.2) }), "SQF PO (13x13)"));
    gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureS<uchar, float, float, PolarCV>>({ NV("grad_scale", 1), NV("grad_muls", 2), NV("grad_kernel_size", 15), NV("grad_kernel_spacing", 1.2) }), "SQF PO (15x15)"));

    //gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureLGF<uchar, float, PolarCV>>({ NV("grad_waveLength", 3), NV("grad_sigmaOnf", 0.55) }), "SQFF LG"));
    //gradF.push_back(Entry<float, float, float>(new GradientEnergy<QuadratureSF<uchar, float, PolarCV>>({ NV("grad_scale", 1), NV("grad_muls", 2), NV("grad_kernel_spacing", 1.2) }), "SQFF PO"));

    //gradF.push_back(Entry<float, float, float>(new GradientPC<PCLgf<uchar, float, PolarCV>>, "PC LG"));

    //gradF.push_back(Entry<float, float, float>(new GradientPC<PCSqf<uchar, float, PolarCV>>({ NV("grad_scale", 1), NV("grad_muls", 2), NV("grad_kernel_spacing", 1.2) }), "PC PO"));

    std::vector<Entry<double, double, double>> gradD;
    //gradD.push_back(Entry<double, double, double>(new GradientPC<PCMatlab<uchar>>, "PC ML"));

    int cols = gradI.size() + gradF.size() + gradD.size() + 1;
    int rows = sets.size() + 1;
    std::vector<std::vector<std::string>> table;
    table.resize(cols);
    for_each(table.begin(), table.end(), [&](std::vector<std::string> &row) {
        row.resize(rows);
    });

    table[0][0] = "Method";

    int row = 1;
    for_each(gradI.begin(), gradI.end(), [&](const Entry<short, int, float> &e) {
        table[row++][0] = e.name;
    });
    for_each(gradF.begin(), gradF.end(), [&](const Entry<float, float, float> &e) {
        table[row++][0] = e.name;
    });
    for_each(gradD.begin(), gradD.end(), [&](const Entry<double, double, double> &e) {
        table[row++][0] = e.name;
    });

    int col = 1;
    for_each(sets.begin(), sets.end(), [&](const std::pair<fs::path, std::string> &data) {

        processPath(gradI, data);
        processPath(gradF, data);
        processPath(gradD, data);

        table[0][col] = data.second;
        row = 1;
        for_each(gradI.begin(), gradI.end(), [&](const Entry<short, int, float> &e) {
            table[row++][col] = boost::str(boost::format("%.3f") % (static_cast<double>(e.time * 1000) / (e.images * cv::getTickFrequency()))) + "ms";
        });

        for_each(gradF.begin(), gradF.end(), [&](const Entry<float, float, float> &e) {
            table[row++][col] = boost::str(boost::format("%.3f") % (static_cast<double>(e.time * 1000) / (e.images * cv::getTickFrequency()))) + "ms";
        });

        for_each(gradD.begin(), gradD.end(), [&](const Entry<double, double, double> &e) {
            table[row++][col] = boost::str(boost::format("%.3f") % (static_cast<double>(e.time * 1000) / (e.images * cv::getTickFrequency()))) + "ms";
        });
        ++col;
    });

    std::ofstream ofs;
    ofs.open("gradient_profiling.csv");

    for_each(table.begin(), table.end(), [&](const std::vector<std::string> &row) {
        for_each(row.begin(), row.end(), [&](const std::string &cell) {
            std::cout << cell << "\t";
            ofs << cell << ";";
        });
        std::cout << std::endl;
        ofs << std::endl;
    });

    ofs.close();

    return 0;
}


#pragma once

#include <list>
#include <string>
#include <boost/spirit/home/x3/support/ast/variant.hpp>
#include <boost/foreach.hpp>

namespace client {
typedef unsigned int uint;
namespace ast {
    namespace x3 = boost::spirit::x3;
    struct nil {};
    struct binary_op;
    struct conditional_op;
    struct expression;

    struct binary_operator {
        std::string name;
        std::function<uint(uint, uint)> op;

        uint
        operator()(uint lhs, uint rhs) const {
            return op(lhs, rhs);
        }
    };

    struct operand : x3::variant<
                         nil,
                         uint,
                         std::string,
                         x3::forward_ast<binary_op>,
                         x3::forward_ast<conditional_op>,
                         x3::forward_ast<expression>> {
        using base_type::base_type;
        using base_type::operator=;
    };

    struct binary_op {
        binary_operator op;
        operand lhs;
        operand rhs;
    };

    struct conditional_op {
        operand lhs;
        operand rhs_true;
        operand rhs_false;
    };

    struct operation {
        binary_operator op;
        operand rhs;
    };

    struct expression {
        operand lhs;
        std::list<operation> rhs;
    };

    struct printer {
        typedef void result_type;

        result_type
        operator()(operand const& ast) const {
            boost::apply_visitor(*this, ast.get());
        }

        result_type
        operator()(nil) const {}

        result_type
        operator()(expression const& ast) const {
            if (ast.rhs.size() > 0) {
                std::cout << '(';
            }
            boost::apply_visitor(*this, ast.lhs);
            BOOST_FOREACH (operation const& op, ast.rhs) { (*this)(op); }
            if (ast.rhs.size() > 0) {
                std::cout << ')';
            }
        }

        result_type
        operator()(operation const& ast) const {
            std::cout << ' ' << ast.op.name << ' ';
            boost::apply_visitor(*this, ast.rhs);
        }

        result_type
        operator()(binary_op const& ast) const {
            std::cout << '(';
            boost::apply_visitor(*this, ast.lhs);
            boost::apply_visitor(*this, ast.rhs);
            std::cout << ')';
        }

        result_type
        operator()(conditional_op const& ast) const {
            std::cout << '(';
            boost::apply_visitor(*this, ast.lhs);
            std::cout << " ? ";
            boost::apply_visitor(*this, ast.rhs_true);
            std::cout << " : ";
            boost::apply_visitor(*this, ast.rhs_false);
            std::cout << ')';
        }

        result_type
        operator()(uint const& ast) const {
            std::cout << ast;
        }

        result_type
        operator()(std::string const& ast) const {
            std::cout << ast;
        }
    };

    struct evaluator {
        typedef uint result_type;

        evaluator(const result_type variable) : variable(variable) {}
        result_type variable;

        result_type
        operator()(operand const& ast) const {
            return boost::apply_visitor(*this, ast.get());
        }

        result_type
        operator()(nil) const {
            BOOST_ASSERT(0);
            return 0;
        }

        result_type
        operator()(expression const& ast) const {
            result_type state = boost::apply_visitor(*this, ast.lhs);
            BOOST_FOREACH (operation const& op, ast.rhs) {
                state = (*this)(op, state);
            }
            return state;
        }

        result_type
        operator()(operation const& ast, uint lhs) const {
            result_type rhs = boost::apply_visitor(*this, ast.rhs);
            return ast.op(lhs, rhs);
        }

        result_type
        operator()(binary_op const& ast) const {
            result_type lhs = boost::apply_visitor(*this, ast.lhs);
            result_type rhs = boost::apply_visitor(*this, ast.rhs);
            return ast.op(lhs, rhs);
        }

        result_type
        operator()(conditional_op const& ast) const {
            bool lhs = boost::apply_visitor(*this, ast.lhs);
            if (lhs) {
                return boost::apply_visitor(*this, ast.rhs_true);
            }
            return boost::apply_visitor(*this, ast.rhs_false);
        }

        result_type
        operator()(uint const& ast) const {
            return ast;
        }

        result_type
        operator()(std::string const& ast) const {
            return variable;
        }
    };

} // namespace ast
} // namespace client


#include <stan/math/rev/mat.hpp>
#include <gtest/gtest.h>
#include <stan/math.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/math/distributions.hpp>
#include <test/unit/math/rev/mat/prob/lkj_corr_cholesky_test_functors.hpp>
#include <test/unit/math/rev/mat/prob/test_gradients.hpp>
#include <test/unit/math/rev/mat/util.hpp>

TEST(ProbDistributionsLkjCorr, var) {
  using stan::math::var;
  boost::random::mt19937 rng;
  int K = 4;
  Eigen::Matrix<var, Eigen::Dynamic, Eigen::Dynamic> Sigma(K, K);
  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Sigma_d(K, K);
  Sigma.setZero();
  Sigma.diagonal().setOnes();
  Sigma_d.setZero();
  Sigma_d.diagonal().setOnes();
  var eta = stan::math::uniform_rng(0, 2, rng);
  var f = stan::math::do_lkj_constant(eta, K);
  EXPECT_FLOAT_EQ(f.val(), stan::math::lkj_corr_log(Sigma, eta).val());
  EXPECT_FLOAT_EQ(f.val(), stan::math::lkj_corr_log(Sigma_d, eta).val());
  eta = 1.0;
  double eta_d = 1.0;
  f = stan::math::do_lkj_constant(eta, K);
  EXPECT_FLOAT_EQ(f.val(), stan::math::lkj_corr_log(Sigma, eta).val());
  EXPECT_FLOAT_EQ(f.val(), stan::math::lkj_corr_log(Sigma, eta_d).val());
}

TEST(ProbDistributionsLkjCorrCholesky, var) {
  using stan::math::var;
  boost::random::mt19937 rng;
  int K = 4;
  Eigen::Matrix<var, Eigen::Dynamic, Eigen::Dynamic> Sigma(K, K);
  Sigma.setZero();
  Sigma.diagonal().setOnes();
  Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic> Sigma_d(K, K);
  Sigma_d.setZero();
  Sigma_d.diagonal().setOnes();
  var eta = stan::math::uniform_rng(0, 2, rng);
  var f = stan::math::do_lkj_constant(eta, K);
  EXPECT_FLOAT_EQ(f.val(), stan::math::lkj_corr_cholesky_log(Sigma, eta).val());
  EXPECT_FLOAT_EQ(f.val(),
                  stan::math::lkj_corr_cholesky_log(Sigma_d, eta).val());
  eta = 1.0;
  double eta_d = 1.0;
  f = stan::math::do_lkj_constant(eta, K);
  EXPECT_FLOAT_EQ(f.val(), stan::math::lkj_corr_cholesky_log(Sigma, eta).val());
  EXPECT_FLOAT_EQ(f.val(),
                  stan::math::lkj_corr_cholesky_log(Sigma, eta_d).val());
}

TEST(ProbDistributionsLkjCorrCholesky, gradients) {
  using stan::math::var;
  int dim_mat = 3;
  Eigen::Matrix<double, Eigen::Dynamic, 1> x1(dim_mat);
  Eigen::Matrix<double, Eigen::Dynamic, 1> x2(1);
  Eigen::Matrix<double, Eigen::Dynamic, 1> x3(dim_mat + 1);

  x2(0) = 2.0;

  for (int i = 0; i < dim_mat; ++i) {
    x1(i) = i / 10.0;
    x3(i + 1) = x1(i);
  }
  x3(0) = 0.5;

  stan::math::lkj_corr_cholesky_dc test_func_1(dim_mat);
  stan::math::lkj_corr_cholesky_cd test_func_2(dim_mat);
  stan::math::lkj_corr_cholesky_dd test_func_3(dim_mat);

  using stan::math::finite_diff_gradient;
  using stan::math::gradient;

  Eigen::Matrix<double, Eigen::Dynamic, 1> grad;
  double fx;
  Eigen::Matrix<double, Eigen::Dynamic, 1> grad_ad;
  double fx_ad;

  finite_diff_gradient(test_func_3, x3, fx, grad);
  gradient(test_func_3, x3, fx_ad, grad_ad);

  test_grad_eq(grad, grad_ad);
  EXPECT_FLOAT_EQ(fx, fx_ad);

  finite_diff_gradient(test_func_2, x2, fx, grad);
  gradient(test_func_2, x2, fx_ad, grad_ad);
  test_grad_eq(grad, grad_ad);
  EXPECT_FLOAT_EQ(fx, fx_ad);

  Eigen::Matrix<double, Eigen::Dynamic, 1> grad_1;
  double fx_1;
  Eigen::Matrix<double, Eigen::Dynamic, 1> grad_ad_1;
  double fx_ad_1;

  finite_diff_gradient(test_func_1, x1, fx_1, grad_1);
  gradient(test_func_1, x1, fx_ad_1, grad_ad_1);
  test_grad_eq(grad_1, grad_ad_1);
  EXPECT_FLOAT_EQ(fx, fx_ad);
}


// Copyright 2015-2020 The ALMA Project Developers
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.

/// @file
/// Test the functionality of the psi_calculator.

#include <iostream>
#include <functional>
#include <boost/filesystem.hpp>
#include <boost/mpi.hpp>
#include "gtest/gtest.h"
#include <utilities.hpp>
#include <vasp_io.hpp>
#include <qpoint_grid.hpp>
#include <processes.hpp>
#include <bulk_hdf5.hpp>
#include <analytic1d.hpp>

TEST(analytic1D_psi_case, bulk100_test) {
    auto basedir = boost::filesystem::path(TEST_RESOURCE_DIR);
    auto si_dir = basedir / boost::filesystem::path("Si");
    auto hdf5_path = si_dir / boost::filesystem::path("Si_threeph.h5");
    boost::mpi::communicator world;
    auto my_id = world.rank();

    // obtain phonon data from HDF5 file
    auto hdf5_data = alma::load_bulk_hdf5(hdf5_path.string().c_str(), world);
    auto description = std::get<0>(hdf5_data);
    auto poscar = std::move(std::get<1>(hdf5_data));
    auto syms = std::move(std::get<2>(hdf5_data));
    auto grid = std::move(std::get<3>(hdf5_data));
    auto processes = std::move(std::get<4>(hdf5_data));

    double T = 300.0;

    Eigen::ArrayXXd w(alma::calc_w0_threeph(*grid, *processes, T, world));

    // create transport direction
    Eigen::Vector3d u100(1.0, 0.0, 0.0);

    // general variables
    int Nxi;

    if (my_id == 0) { // analytic1D code is not parallellised; avoid
                      // running duplicates.
        Nxi = 4;
        alma::analytic1D::psi_calculator psiCalc(
            poscar.get(), grid.get(), &w, T);
        psiCalc.setLogGrid(1e3, 1e9, Nxi);
        Eigen::VectorXd xigrid = psiCalc.getSpatialFrequencies();

        // Calculate psi functions in bulk Si for 3 directions
        psiCalc.normaliseOutput(true); // turn on optional rescaling
                                       // (output gets divided by
                                       // Dbulk*xi^2)

        psiCalc.setDirection(u100);
        Eigen::VectorXd psi100 = psiCalc.getPsi();

        Eigen::VectorXd psi100_ref(4);
        psi100_ref << 1.000585888, 0.9864015985, 0.4049977681, 0.005287157147;
        Eigen::VectorXd ratio100 = psi100.array() / psi100_ref.array();

        for (int nxi = 0; nxi < Nxi; nxi++) {
            EXPECT_NEAR(ratio100(nxi), 1.0, 5e-3);
        }
    }
}

TEST(analytic1D_psi_case, bulk110_test) {
    auto basedir = boost::filesystem::path(TEST_RESOURCE_DIR);
    auto si_dir = basedir / boost::filesystem::path("Si");
    auto hdf5_path = si_dir / boost::filesystem::path("Si_threeph.h5");
    boost::mpi::communicator world;
    auto my_id = world.rank();

    // obtain phonon data from HDF5 file
    auto hdf5_data = alma::load_bulk_hdf5(hdf5_path.string().c_str(), world);
    auto description = std::get<0>(hdf5_data);
    auto poscar = std::move(std::get<1>(hdf5_data));
    auto syms = std::move(std::get<2>(hdf5_data));
    auto grid = std::move(std::get<3>(hdf5_data));
    auto processes = std::move(std::get<4>(hdf5_data));

    double T = 300.0;

    Eigen::ArrayXXd w(alma::calc_w0_threeph(*grid, *processes, T, world));

    // create transport direction
    Eigen::Vector3d u110(1.0, 1.0, 0.0);

    // general variables
    int Nxi;

    if (my_id == 0) { // analytic1D code is not parallellised; avoid
                      // running duplicates.
        Nxi = 4;
        alma::analytic1D::psi_calculator psiCalc(
            poscar.get(), grid.get(), &w, T);
        psiCalc.setLogGrid(1e3, 1e9, Nxi);
        Eigen::VectorXd xigrid = psiCalc.getSpatialFrequencies();

        // Calculate psi functions in bulk Si for 3 directions
        psiCalc.normaliseOutput(true); // turn on optional rescaling
                                       // (output gets divided by
                                       // Dbulk*xi^2)

        psiCalc.setDirection(u110);
        Eigen::VectorXd psi110 = psiCalc.getPsi();

        Eigen::VectorXd psi110_ref(4);
        psi110_ref << 1.000585549, 0.9834471977, 0.3966867118, 0.005560183293;
        Eigen::VectorXd ratio110 = psi110.array() / psi110_ref.array();

        for (int nxi = 0; nxi < Nxi; nxi++) {
            EXPECT_NEAR(ratio110(nxi), 1.0, 5e-3);
        }
    }
}

TEST(analytic1D_psi_case, bulk111_test) {
    auto basedir = boost::filesystem::path(TEST_RESOURCE_DIR);
    auto si_dir = basedir / boost::filesystem::path("Si");
    auto hdf5_path = si_dir / boost::filesystem::path("Si_threeph.h5");
    boost::mpi::communicator world;
    auto my_id = world.rank();

    // obtain phonon data from HDF5 file
    auto hdf5_data = alma::load_bulk_hdf5(hdf5_path.string().c_str(), world);
    auto description = std::get<0>(hdf5_data);
    auto poscar = std::move(std::get<1>(hdf5_data));
    auto syms = std::move(std::get<2>(hdf5_data));
    auto grid = std::move(std::get<3>(hdf5_data));
    auto processes = std::move(std::get<4>(hdf5_data));

    double T = 300.0;

    Eigen::ArrayXXd w(alma::calc_w0_threeph(*grid, *processes, T, world));

    // create transport direction
    Eigen::Vector3d u111(1.0, 1.0, 1.0);

    // general variables
    int Nxi;

    if (my_id == 0) { // analytic1D code is not parallellised; avoid
                      // running duplicates.
        Nxi = 4;
        alma::analytic1D::psi_calculator psiCalc(
            poscar.get(), grid.get(), &w, T);
        psiCalc.setLogGrid(1e3, 1e9, Nxi);
        Eigen::VectorXd xigrid = psiCalc.getSpatialFrequencies();

        // Calculate psi functions in bulk Si for 3 directions
        psiCalc.normaliseOutput(true); // turn on optional rescaling
                                       // (output gets divided by
                                       // Dbulk*xi^2)

        psiCalc.setDirection(u111);
        Eigen::VectorXd psi111 = psiCalc.getPsi();

        Eigen::VectorXd psi111_ref(4);
        psi111_ref << 1.000585193, 0.981217029, 0.4013087323, 0.006630052249;
        Eigen::VectorXd ratio111 = psi111.array() / psi111_ref.array();

        for (int nxi = 0; nxi < Nxi; nxi++) {
            EXPECT_NEAR(ratio111(nxi), 1.0, 5e-3);
        }
    }
}

int main(int argc, char** argv) {
    boost::mpi::environment env;

    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}


/**
 *  Copyright (C) 2012
 *    Ekaterina Potapova
 *    Automation and Control Institute
 *    Vienna University of Technology
 *    Gusshausstraße 25-29
 *    1040 Vienna, Austria
 *    potapova(at)acin.tuwien.ac.at
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see http://www.gnu.org/licenses/
 */

#include "v4r/attention_segmentation/PCA.h"
#include <Eigen/Dense>

namespace v4r
{

Eigen::Vector4f getMean(pcl::PointCloud<pcl::Normal>::ConstPtr cloud)
{
  Eigen::Vector4f mean;
  mean.setZero();

  for (unsigned int pi=0; pi<cloud->size(); pi++)
  {
    mean[0] += cloud->points.at(pi).normal[0];
    mean[1] += cloud->points.at(pi).normal[1];
    mean[2] += cloud->points.at(pi).normal[2];
  }

  mean /= (float)cloud->size();

  return(mean);

}

bool computeCovarianceMatrix(pcl::PointCloud<pcl::Normal>::ConstPtr cloud, const Eigen::Vector4f &mean, Eigen::Matrix3f &cov)
{
  bool done = false;
  cov.setZero ();

  for (unsigned pi = 0; pi < cloud->size (); ++pi)
  {
    float x = cloud->points.at(pi).normal[0] - mean[0];
    float y = cloud->points.at(pi).normal[1] - mean[1];
    float z = cloud->points.at(pi).normal[2] - mean[2];

    cov(0,0) += x*x;
    cov(0,1) += x*y;
    cov(0,2) += x*z;

    cov(1,0) += y*x;
    cov(1,1) += y*y;
    cov(1,2) += y*z;

    cov(2,0) += z*x;
    cov(2,1) += z*y;
    cov(2,2) += z*z;

    done = true;
  }

  return(done);
}

void principleAxis(pcl::PointCloud<pcl::Normal>::ConstPtr cloud, std::vector<pcl::Normal> &axis)
{
  Eigen::Vector4f mean;
  EIGEN_ALIGN16 Eigen::Matrix3f cov;
  EIGEN_ALIGN16 Eigen::Vector3f eigen_values;
  EIGEN_ALIGN16 Eigen::Matrix3f eigen_vectors;

  axis.clear();
  axis.resize(3);

  mean = getMean(cloud);

  if(computeCovarianceMatrix(cloud,mean,cov))
  {
    Eigen::SelfAdjointEigenSolver<Eigen::Matrix3f> eigensolver(cov);
    if (eigensolver.info() != Eigen::Success)
      return;

    eigen_values = eigensolver.eigenvalues();
    eigen_vectors = eigensolver.eigenvectors();

    //pcl::eigen33 (cov, eigen_vectors, eigen_values);

    //std::cerr << "inside principleAxis" << std::endl;

    axis.at(0).normal[0] = eigen_vectors (0,0);
    axis.at(0).normal[1] = eigen_vectors (1,0);
    axis.at(0).normal[2] = eigen_vectors (2,0);

    axis.at(1).normal[0] = eigen_vectors (0,1);
    axis.at(1).normal[1] = eigen_vectors (1,1);
    axis.at(1).normal[2] = eigen_vectors (2,1);

    axis.at(2).normal[0] = eigen_vectors (0,2);
    axis.at(2).normal[1] = eigen_vectors (1,2);
    axis.at(2).normal[2] = eigen_vectors (2,2);
  }

}

} //namespace v4r