/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: t; c-basic-offset: 4 -*- vim:set ts=4 sw=4 sts=4 noet: */
#include "compiler/tac/context.h"

#include <algorithm>
#include <cassert>
#include <cstdio>
#include <iostream>
#include <memory>
#include <string>
#include <vector>

#include <boost/uuid/uuid_io.hpp>

#include "lexer.h"

namespace flint {
namespace compiler {
namespace tac {

namespace {

bool IsAnd(const cas::Compound &c)
{
	return c.keyword == "and";
}

bool IsOr(const cas::Compound &c)
{
	return c.keyword == "or";
}

bool IsAt(const cas::Compound &c)
{
	return c.keyword == "$At";
}

bool IsData(const cas::Compound &c)
{
	return c.keyword == "$Data";
}

bool IsLookback(const cas::Compound &c)
{
	return c.keyword == "$lookback";
}

bool IsPiecewise(const cas::Compound &c)
{
	return c.keyword == "piecewise";
}

bool IsPiece(const cas::Compound &c)
{
	return c.keyword == "piece";
}

bool IsOtherwise(const cas::Compound &c)
{
	return c.keyword == "otherwise";
}

bool IsTrial(const cas::Compound &c)
{
	return c.keyword == "$trial";
}

bool IsOutcome(const cas::Compound &c)
{
	return c.keyword == "$outcome";
}

struct KeyFun {
	std::string keyword;
	bool (Context::*function)(RegisterType rt, int n, cas::Compound &c);
};

const KeyFun kKeyFun[] = {
	// Keep the following entries in bibliographical order.
	{"$Mod", &Context::NaryScalarToScalar},
	{"$exponential_variate", &Context::NaryScalarToScalar},
	{"$gamma_variate", &Context::NaryScalarToScalar},
	{"$gauss_variate", &Context::NaryScalarToScalar},
	{"$lognormal_variate", &Context::NaryScalarToScalar},
	{"$poisson_variate", &Context::NaryScalarToScalar},
	{"$uniform_variate", &Context::NaryScalarToScalar},
	{"$weibull_variate", &Context::NaryScalarToScalar},
	{"abs", &Context::NaryScalarToScalar},
	{"and", &Context::NaryScalarToScalar},
	{"arccos", &Context::NaryScalarToScalar},
	{"arccosh", &Context::NaryScalarToScalar},
	{"arccot", &Context::NaryScalarToScalar},
	{"arccoth", &Context::NaryScalarToScalar},
	{"arccsc", &Context::NaryScalarToScalar},
	{"arccsch", &Context::NaryScalarToScalar},
	{"arcsec", &Context::NaryScalarToScalar},
	{"arcsech", &Context::NaryScalarToScalar},
	{"arcsin", &Context::NaryScalarToScalar},
	{"arcsinh", &Context::NaryScalarToScalar},
	{"arctan", &Context::NaryScalarToScalar},
	{"arctanh", &Context::NaryScalarToScalar},
	{"ceiling", &Context::NaryScalarToScalar},
	{"cos", &Context::NaryScalarToScalar},
	{"cosh", &Context::NaryScalarToScalar},
	{"cot", &Context::NaryScalarToScalar},
	{"coth", &Context::NaryScalarToScalar},
	{"csc", &Context::NaryScalarToScalar},
	{"csch", &Context::NaryScalarToScalar},
	{"determinant", &Context::Determinant},
	{"divide", &Context::NaryScalarToScalar},
	{"eq", &Context::Componentwise},
	{"exp", &Context::NaryScalarToScalar},
	{"floor", &Context::NaryScalarToScalar},
	{"geq", &Context::NaryScalarToScalar},
	{"gt", &Context::NaryScalarToScalar},
	{"leq", &Context::NaryScalarToScalar},
	{"ln", &Context::NaryScalarToScalar},
	{"log", &Context::NaryScalarToScalar},
	{"log10", &Context::NaryScalarToScalar},
	{"lt", &Context::NaryScalarToScalar},
	{"matrix", &Context::Matrix},
	{"max", &Context::NaryScalarToScalar},
	{"mean", &Context::NaryScalarToScalar},
	{"min", &Context::NaryScalarToScalar},
	{"minus", &Context::Componentwise},
	{"neq", &Context::Componentwise},
	{"not", &Context::NaryScalarToScalar},
	{"or", &Context::NaryScalarToScalar},
	{"outerproduct", &Context::Outerproduct},
	{"plus", &Context::Componentwise},
	{"power", &Context::NaryScalarToScalar}, // TODO: matrix
	{"rem", &Context::NaryScalarToScalar},
	{"root", &Context::NaryScalarToScalar},
	{"scalarproduct", &Context::Scalarproduct},
	{"sec", &Context::NaryScalarToScalar},
	{"sech", &Context::NaryScalarToScalar},
	{"selector", &Context::Selector},
	{"sin", &Context::NaryScalarToScalar},
	{"sinh", &Context::NaryScalarToScalar},
	{"tan", &Context::NaryScalarToScalar},
	{"tanh", &Context::NaryScalarToScalar},
	{"times", &Context::Times},
	{"transpose", &Context::Transpose},
	{"vector", &Context::Vector},
	{"vectorproduct", &Context::Vectorproduct},
	{"xor", &Context::NaryScalarToScalar}
};

}

Context::Context(const boost::uuids::uuid &uuid,
				 const char *id,
				 std::ostream *os)
	: uuid_(uuid)
	, id_(id)
	, ir_(0)
	, fr_(0)
	, l_(0)
	, os_(os)
{}

Context::~Context() = default;

bool Context::EmitCode(cas::Expr &expr)
{
	int w, col, row;
	GetType(expr, &w, &col, &row);
	if (col == 1 && row == 1) {
		fr_++;
		if (!Assign(RegisterType::kFloat, 0, expr))
			return false;
		*os_ << "  store " << id_ << " $0" << std::endl;
	} else {
		ir_++;
		if (!Assign(RegisterType::kInteger, 0, expr))
			return false;
		*os_ << "  save " << id_ << " $i0 " << (col * row) << std::endl;
	}
	return true;
}

bool Context::Componentwise(RegisterType rt, int n, cas::Compound &c)
{
	if (rt == RegisterType::kFloat)
		return NaryScalarToScalar(rt, n, c);

	int size = c.col * c.row;
	std::vector<int> args;
	for (auto &e : c.children) {
		int p = ir_++;
		if (!Assign(RegisterType::kInteger, p, e))
			return false;
		args.push_back(p);
	}
	*os_ << "  alloc $i" << n << ' ' << size << std::endl;
	for (int i=0;i<size;i++) {
		std::vector<int> params;
		params.reserve(args.size());
		for (auto k : args) {
			int m = fr_++;
			*os_ << "  deref $" << m << " $i" << k << ' ' << i << std::endl;
			params.push_back(m);
		}
		int m = fr_++;
		*os_ << "  $" << m << " = (" << c.keyword;
		for (auto k : params)
			*os_ << " $" << k;
		*os_ << ')' << std::endl;
		*os_ << "  move $i" << n << " $" << m << ' ' << i << std::endl;
	}
	return true;
}

bool Context::NaryScalarToScalar(RegisterType /*rt*/, int n, cas::Compound &c)
{
	std::vector<int> args;
	for (auto &e : c.children) {
		int m = fr_++;
		if (!Assign(RegisterType::kFloat, m, e))
			return false;
		args.push_back(m);
	}
	*os_ << "  $" << n << " = (" << c.keyword;
	for (auto i : args)
		*os_ << " $" << i;
	*os_ << ')' << std::endl;
	return true;
}

bool Context::Determinant(RegisterType rt, int n, cas::Compound &c)
{
	assert(rt == RegisterType::kFloat);
	assert(c.children.size() == 1);
	int col, row;
	GetType(c.children[0], nullptr, &col, &row);
	assert(col == row);
	int n1 = ir_++;
	if (!Assign(RegisterType::kInteger, n1, c.children[0]))
		return false;
	*os_ << "  determinant $" << n << ' ' << col << " $i" << n1 << std::endl;
	return true;
}

bool Context::Matrix(RegisterType rt, int n, cas::Compound &c)
{
	assert(rt == RegisterType::kInteger);
	*os_ << "  alloc $i" << n << ' ' << (c.col * c.row) << std::endl;
	int i=0;
	for (auto &e : c.children) {
		assert(e.which() == cas::kExprIsCompound);
		cas::Compound &rc(boost::get<cas::Compound>(e));
		assert(rc.keyword == "matrixrow");
		for (auto &rcc : rc.children) {
			int m = fr_++;
			if (!Assign(RegisterType::kFloat, m, rcc))
				return false;
			*os_ << "  move $i" << n << " $" << m << ' ' << i++ << std::endl;
		}
	}
	return true;
}

bool Context::Outerproduct(RegisterType rt, int n, cas::Compound &c)
{
	assert(rt == RegisterType::kInteger);
	int n1 = ir_++;
	if (!Assign(RegisterType::kInteger, n1, c.children[0]))
		return false;
	int n2 = ir_++;
	if (!Assign(RegisterType::kInteger, n2, c.children[1]))
		return false;
	*os_ << "  outerproduct $i" << n
		 << ' ' << c.row << " $i" << n1
		 << ' ' << c.col << " $i" << n2
		 << std::endl;
	return true;
}

bool Context::Scalarproduct(RegisterType rt, int n, cas::Compound &c)
{
	assert(rt == RegisterType::kFloat);
	int col, row;
	GetType(c.children[0], nullptr, &col, &row);
	int n1 = ir_++;
	if (!Assign(RegisterType::kInteger, n1, c.children[0]))
		return false;
	int n2 = ir_++;
	if (!Assign(RegisterType::kInteger, n2, c.children[1]))
		return false;
	int size = std::max(col, row);
	*os_ << "  scalarproduct $" << n << ' ' << size << " $i" << n1 << " $i" << n2 << std::endl;
	return true;
}

bool Context::Selector(RegisterType rt, int n, cas::Compound &c)
{
	size_t size = c.children.size();
	assert(size >= 2);
	int n0 = ir_++;
	if (!Assign(RegisterType::kInteger, n0, c.children[0]))
		return false;
	int m1 = fr_++;
	if (!Assign(RegisterType::kFloat, m1, c.children[1]))
		return false;
	if (rt == RegisterType::kInteger) {
		assert(size == 2);
		int col0, row0;
		GetType(c.children[0], nullptr, &col0, &row0);
		*os_ << "  selrow $i" << n
			 << ' ' << row0
			 << ' ' << col0
			 << " $i" << n0
			 << " $" << m1 << std::endl;
	} else if (size == 2) {
		*os_ << "  select2 $" << n << " $i" << n0 << " $" << m1 << std::endl;
	} else {
		assert(size == 3);
		int col0, row0;
		GetType(c.children[0], nullptr, &col0, &row0);
		int m2 = fr_++;
		if (!Assign(RegisterType::kFloat, m2, c.children[2]))
			return false;
		*os_ << "  select3 $" << n
			 << ' ' << row0
			 << ' ' << col0
			 << " $i" << n0
			 << " $" << m1
			 << " $" << m2
			 << std::endl;
	}
	return true;
}

bool Context::Times(RegisterType rt, int n, cas::Compound &c)
{
	if (rt == RegisterType::kFloat)
		return NaryScalarToScalar(rt, n, c);
	assert(c.children.size() == 2);
	cas::Expr &e0(c.children.at(0));
	cas::Expr &e1(c.children.at(1));
	int w0, w1;
	int col0, col1;
	int row0, row1;
	GetType(e0, &w0, &col0, &row0);
	GetType(e1, &w1, &col1, &row1);
	if (col0 == 1 && row0 == 1) { // when the 1st operand is scalar
		int m0 = fr_++;
		if (!Assign(RegisterType::kFloat, m0, e0))
			return false;
		int n1 = ir_++;
		if (!Assign(RegisterType::kInteger, n1, e1))
			return false;
		int size = col1 * row1;
		*os_ << "  mult $i" << n << ' ' << size << " $" << m0 << " $i" << n1 << std::endl;
	} else if (col1 == 1 && row1 == 1) { // when the 2nd operand is scalar
		int n0 = ir_++;
		if (!Assign(RegisterType::kInteger, n0, e0))
			return false;
		int m1 = fr_++;
		if (!Assign(RegisterType::kFloat, m1, e1))
			return false;
		int size = col0 * row0;
		*os_ << "  mult $i" << n << ' ' << size << " $" << m1 << " $i" << n0 << std::endl;
	} else { // both are matrices
		int n0 = ir_++;
		if (!Assign(RegisterType::kInteger, n0, e0))
			return false;
		int n1 = ir_++;
		if (!Assign(RegisterType::kInteger, n1, e1))
			return false;
		*os_ << "  mmul $i" << n
			 << ' ' << row0
			 << ' ' << col0 // = row1
			 << ' ' << col1
			 << " $i" << n0
			 << " $i" << n1
			 << std::endl;
	}
	return true;
}

bool Context::Transpose(RegisterType rt, int n, cas::Compound &c)
{
	assert(rt == RegisterType::kInteger);
	int n1 = ir_++;
	if (!Assign(RegisterType::kInteger, n1, c.children[0]))
		return false;
	*os_ << "  transpose $i" << n
		 << " $i" << n1
		 << ' ' << c.row
		 << ' ' << c.col
		 << std::endl;
	return true;
}

bool Context::Vector(RegisterType rt, int n, cas::Compound &c)
{
	assert(rt == RegisterType::kInteger);
	int size = c.col * c.row;
	*os_ << "  alloc $i" << n << ' ' << size << std::endl;
	for (int i=0;i<size;i++) {
		int m = fr_++;
		if (!Assign(RegisterType::kFloat, m, c.children[i]))
			return false;
		*os_ << "  move $i" << n << " $" << m << ' ' << i << std::endl;
	}
	return true;
}

bool Context::Vectorproduct(RegisterType rt, int n, cas::Compound &c)
{
	assert(rt == RegisterType::kInteger);
	int n1 = ir_++;
	if (!Assign(RegisterType::kInteger, n1, c.children[0]))
		return false;
	int n2 = ir_++;
	if (!Assign(RegisterType::kInteger, n2, c.children[1]))
		return false;
	*os_ << "  vectorproduct $i" << n << " $i" << n1 << " $i" << n2 << std::endl;
	return true;
}

/*
 * The control should reach the end when the expr is evaluated to be false.
 *
 * @param n the variable for the resulting boolean value
 * @param l the label to go when the expr is evaluated to be true
 */
bool Context::EmitCondition(int n, int l, cas::Expr &expr)
{
	if (expr.which() == cas::kExprIsCompound) {
		cas::Compound &comp(boost::get<cas::Compound>(expr));
		if (IsAnd(comp)) {
			int l1 = l_++;
			int l2 = l_++;
			if (!EmitCondition(n, l2, comp.children.at(0)))
				return false;
			*os_ << "  jmp L" << l1 << std::endl;
			*os_ << " L" << l2 << ':' << std::endl;
			if (!EmitCondition(n, l, comp.children.at(1)))
				return false;
			*os_ << " L" << l1 << ':' << std::endl;
			return bool(*os_);
		} else if (IsOr(comp)) {
			return EmitCondition(n, l, comp.children.at(0))
				&& EmitCondition(n, l, comp.children.at(1));
		}
	}

	if (!Assign(RegisterType::kFloat, n, expr))
		return false;
	*os_ << "  br $" << n << " L" << l << std::endl;
	return bool(*os_);
}

bool Context::EmitAt(int n, std::deque<cas::Expr> &children)
{
	size_t len = children.size();
	if (len > 3) {
		std::cerr << "error: more than 3 arguments: " << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	if (len < 3) {
		std::cerr << "error: EmitAt: missing arguments: " << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	if (children.at(0).which() != cas::kExprIsInteger) {
		std::cerr << "error: invalid 1st argument of At: " << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	if (children.at(1).which() != cas::kExprIsInteger) {
		std::cerr << "error: invalid 2nd argument of At: " << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	int m = fr_++;
	if (!Assign(RegisterType::kFloat, m, children.at(2))) // TODO: RegisterType::kInteger
		return false;
	*os_ << "  ld $"
		 << n
		 << ' '
		 << boost::get<int>(children.at(0))
		 << ' '
		 << boost::get<int>(children.at(1))
		 << " $"
		 << m
		 << std::endl;
	return bool(*os_);
}

bool Context::EmitData(int n, std::deque<cas::Expr> &children)
{
	size_t len = children.size();
	if (len > 2) {
		std::cerr << "error: more than 2 arguments: " << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	if (len < 2) {
		std::cerr << "error: EmitData: missing arguments: " << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	if (children.at(0).which() != cas::kExprIsInteger) {
		std::cerr << "error: invalid 1st argument of Data: " << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	int m = fr_++;
	if (!Assign(RegisterType::kFloat, m, children.at(1))) // TODO: RegisterType::kInteger
		return false;
	*os_ << "  lc $"
		 << n
		 << ' '
		 << boost::get<int>(children.at(0))
		 << " $"
		 << m
		 << std::endl;
	return bool(*os_);
}

bool Context::EmitLookback(int n, std::deque<cas::Expr> &children)
{
	size_t len = children.size();
	if (len > 2) {
		std::cerr << "error: more than 2 arguments: " << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	if (len < 2) {
		std::cerr << "error: EmitLookback: missing arguments: " << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	if (children.at(0).which() != cas::kExprIsIdentifier) {
		std::cerr << "error: invalid 1st argument of Delay/DeltaTime: " << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	int m = fr_++;
	if (!Assign(RegisterType::kFloat, m, children.at(1))) // TODO: vector/matrix
		return false;
	*os_ << "  lb $"
		 << n
		 << ' '
		 << boost::get<cas::Identifier>(children.at(0)).name.c_str()
		 << " $"
		 << m
		 << std::endl;
	return bool(*os_);
}

bool Context::Piecewise(RegisterType rt, int n, std::deque<cas::Expr> &children)
{
	int l = l_++;
	std::vector<int> v1;
	bool otherwise = false;
	for (auto &expr : children) {
		assert(expr.which() == cas::kExprIsCompound);
		cas::Compound &comp(boost::get<cas::Compound>(expr));
		if (IsPiece(comp)) {
			int l1 = l_++;
			int m = fr_++;
			if (!EmitCondition(m, l1, comp.children.at(1)))
				return false;
			v1.push_back(l1);
		} else if (IsOtherwise(comp)) {
			otherwise = true;
			if (!Assign(rt, n, comp.children.at(0)))
				return false;
			*os_ << "  jmp L" << l << std::endl;
		} else {
			assert(false);
		}
	}
	if (!otherwise)
		*os_ << "  ret" << std::endl;
	auto v1it = v1.cbegin();
	for (auto &expr : children) {
		assert(expr.which() == cas::kExprIsCompound);
		cas::Compound &comp(boost::get<cas::Compound>(expr));
		if (IsPiece(comp)) {
			*os_ << " L" << *v1it++ << ':' << std::endl;
			if (!Assign(rt, n, comp.children.at(0)))
				return false;
			if (v1it != v1.cend()) // need jmp at every branch but last
				*os_ << "  jmp L" << l << std::endl;
		} else if (IsOtherwise(comp)) {
			/* nothing to do */
		} else {
			assert(false);
		}
	}
	*os_ << " L" << l << ':' << std::endl;
	return bool(*os_);
}

bool Context::EmitTrial(int n, std::deque<cas::Expr> &children)
{
	int l = l_++;
	int p0 = fr_++;
	int p1 = fr_++;
	int p = fr_++;
	*os_ << "  loadi $" << p0 << " 0" << std::endl
		 << "  loadi $" << p1 << " 1" << std::endl
		 << "  $" << p << " = ($uniform_variate $" << p0 << " $" << p1 << ')' << std::endl;
	std::vector<int> v1;
	for (auto &expr : children) {
		assert(expr.which() == cas::kExprIsCompound);
		cas::Compound &comp(boost::get<cas::Compound>(expr));
		if (IsOutcome(comp)) {
			int l1 = l_++;
			int m0 = fr_++;
			if (!Assign(RegisterType::kFloat, m0, comp.children.at(1)))
				return false;
			int m1 = fr_++;
			*os_ << "  $" << m1 << " = (minus $" << p << " $" << m0 << ')' << std::endl;
			p = m1; // keep the register number for the decreased value as p
			int m2 = fr_++;
			*os_ << "  $" << m2 << " = (leq $" << m1 << " $" << p0 << ')' << std::endl
				 << "  br $" << m2 << " L" << l1 << std::endl;
			v1.push_back(l1);
		} else {
			std::cerr << "error: unexpected child of $trial: " << uuid_ << ' ' << id_ << std::endl;
			return false;
		}
	}
	*os_ << "  ret" << std::endl;
	auto v1it = v1.cbegin();
	for (auto &expr : children) {
		assert(expr.which() == cas::kExprIsCompound);
		cas::Compound &comp(boost::get<cas::Compound>(expr));
		if (IsOutcome(comp)) {
			*os_ << " L" << *v1it++ << ':' << std::endl;
			if (!Assign(RegisterType::kFloat, n, comp.children.at(0)))
				return false;
			if (v1it != v1.cend()) // needs jmp at every branch but last
				*os_ << "  jmp L" << l << std::endl;
		} else {
			assert(false);
		}
	}
	*os_ << " L" << l << ':' << std::endl;
	return bool(*os_);
}

bool Context::Assign(RegisterType rt, int n, cas::Expr &expr)
{
	switch (expr.which()) {
	case cas::kExprIsCompound:
		return Assign(rt, n, boost::get<cas::Compound>(expr));
	case cas::kExprIsIdentifier:
		return Assign(rt, n, boost::get<cas::Identifier>(expr));
	case cas::kExprIsInteger:
		assert(rt == RegisterType::kFloat);
		return Assign(n, boost::get<int>(expr));
	case cas::kExprIsRational:
		assert(rt == RegisterType::kFloat);
		return Assign(n, boost::get<lexer::Rational>(expr));
	case cas::kExprIsReal:
		assert(rt == RegisterType::kFloat);
		return Assign(n, boost::get<lexer::Real>(expr));
	default:
		assert(false);
		return false;
	}
}

bool Context::Assign(RegisterType rt, int n, cas::Compound &c)
{
	if (IsAt(c))
		return EmitAt(n, c.children);
	if (IsData(c))
		return EmitData(n, c.children);
	if (IsLookback(c))
		return EmitLookback(n, c.children);
	if (IsPiecewise(c))
		return Piecewise(rt, n, c.children);
	if (IsTrial(c))
		return EmitTrial(n, c.children);

	size_t len = c.children.size();
	if (len < 1) {
		std::cerr << "error: missing arguments for " << c.keyword << ": "
			 << uuid_ << ' ' << id_ << std::endl;
		return false;
	}
	for (const auto &kf : kKeyFun) { // TODO: faster search e.g. via bsearch?
		if (c.keyword == kf.keyword)
			return (this->*(kf.function))(rt, n, c);
	}
	std::cerr << "unsupported function: " << c.keyword << std::endl;
	return false;
}

bool Context::Assign(RegisterType rt, int n, const cas::Identifier &id)
{
	const std::string &s(id.name);
	if (rt == RegisterType::kInteger) {
		*os_ << "  refer $i" << n << " " << id.name << std::endl;
	} else {
		if ( (s[0] == '%' || s[0] == '@') && isalpha(s[1]) ) {
			*os_ << "  load $" << n << ' ' << s << std::endl;
		} else {
			*os_ << "  loadi $" << n << ' ' << s << std::endl;
		}
	}
	return true;
}

bool Context::Assign(int n, int i)
{
	*os_ << "  loadi $" << n << ' ' << i << std::endl;
	return true;
}

bool Context::Assign(int n, const lexer::Rational &r)
{
	*os_ << "  loadi $" << n << ' ' << r.lexeme << std::endl;
	return true;
}

bool Context::Assign(int n, const lexer::Real &r)
{
	*os_ << "  loadi $" << n << ' ' << r.lexeme << std::endl;
	return true;
}

void Context::GetType(cas::Expr &expr, int *w, int *col, int *row)
{
	int which = expr.which();
	switch (which) {
	case cas::kExprIsCompound:
		{
			cas::Compound &c(boost::get<cas::Compound>(expr));
			if (col)
				*col = c.col;
			if (row)
				*row = c.row;
		}
		break;
	case cas::kExprIsIdentifier:
		{
			const cas::Identifier &id(boost::get<const cas::Identifier>(expr));
			if (col)
				*col = id.col;
			if (row)
				*row = id.row;
		}
		break;
	default:
		if (col)
			*col = 1;
		if (row)
			*row = 1;
		break;
	}
	if (w)
		*w = which;
}

}
}
}


/** \file   pr_kmeans.cpp
    \brief  Implement kmeans method with different initialization scheme
    \author Hui Xue
*/

#include "pr_kmeans.h"
#include "log.h"

#include "hoNDArray_reductions.h"
#include "hoNDArray_elemwise.h"
#include "hoNDArray_math.h"
#include "hoNDArray_linalg.h"

#include <boost/math/special_functions/sign.hpp>

#include <random>

namespace Gadgetron {

template <typename T>
kmeans<T>::kmeans()
{
    max_iter_ = 100;
    replicates_ = 10;
    perform_online_update_ = true;

    verbose_ = false;
    perform_timing_ = false;

    gt_timer_local_.set_timing_in_destruction(false);
    gt_timer_.set_timing_in_destruction(false);
}

template <typename T>
kmeans<T>::~kmeans()
{
}

template <typename T>
void kmeans<T>::get_initial_guess_sample(const ArrayType& X, size_t K, ArrayType& C_for_initial)
{
    try
    {
        if (this->perform_timing_) gt_timer_local_.start("get_initial_guess_sample");

        size_t P = X.get_size(0);
        size_t N = X.get_size(1);

        GADGET_CHECK_THROW(N>K);

        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_real_distribution<> dis(0, 1);

        C_for_initial.create(P, K, this->replicates_);
        Gadgetron::clear(C_for_initial);

        size_t n, k;
        for (n = 0; n < this->replicates_; n++)
        {
            for (k = 0; k < K; k++)
            {
                size_t ind = (size_t)(dis(gen)*N);
                if (ind >= N) ind = N - 1;
                memcpy(&C_for_initial(0, k, n), &X(0, ind), sizeof(T)*P);
            }
        }

        if (this->perform_timing_) gt_timer_local_.stop();
    }
    catch(...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::get_initial_guess_sample(...) ... ");
    }
}

template <typename T>
void kmeans<T>::get_initial_guess_uniform(const ArrayType& X, size_t K, ArrayType& C_for_initial)
{
    try
    {
        if (this->perform_timing_) gt_timer_local_.start("get_initial_guess_uniform");

        size_t P = X.get_size(0);
        size_t N = X.get_size(1);

        GADGET_CHECK_THROW(N>K);

        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_real_distribution<> dis(0, 1);

        C_for_initial.create(P, K, this->replicates_);
        Gadgetron::clear(C_for_initial);

        // find the range of X
        std::vector<T> xmin(P, 0);
        std::vector<T> xmax(P, 0);

        size_t p, n, k;

        for (p = 0; p<P; p++)
        {
            xmin[p] = X(p, 0);
            xmax[p] = xmin[p];
            for (n = 1; n < N; n++)
            {
                T v = X(p, n);
                if (v < xmin[p]) xmin[p] = v;
                if (v > xmax[p]) xmax[p] = v;
            }
        }

        for (n = 0; n < this->replicates_; n++)
        {
            for (k = 0; k < K; k++)
            {
                for (p = 0; p < P; p++)
                {
                    T v = xmin[p] + dis(gen) * (xmax[p] - xmin[p]);
                    C_for_initial(p, k, n) = v;
                }
            }
        }

        if (this->perform_timing_) gt_timer_local_.stop();
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::get_initial_guess_uniform(...) ... ");
    }
}

template <typename T>
void kmeans<T>::get_initial_guess_cluster(const ArrayType& X, size_t K, ArrayType& C_for_initial)
{
    try
    {
        if (this->perform_timing_) gt_timer_.start("get_initial_guess_cluster");

        size_t P = X.get_size(0);
        size_t N = X.get_size(1);

        GADGET_CHECK_THROW(N>K);

        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_real_distribution<> dis(0, 1);

        C_for_initial.create(P, K, this->replicates_);
        Gadgetron::clear(C_for_initial);

        // randomly pick 20% of data
        double ratio = 0.2;
        size_t M = (size_t)(ratio*N);
        while(M<K && ratio<=1.0)
        {
            ratio = ratio + 0.1;
            M = (size_t)(ratio*N);
        }

        if(M>=N)
        {
            GWARN_STREAM("Too few samples in the data array X ... ");
            this->get_initial_guess_sample(X, K, C_for_initial);
            return;
        }

        ArrayType X_subset;
        X_subset.create(P, M);

        ArrayType C_for_initial_subset;
        ClusterType IDX;
        ArrayType C;
        T sumD;

        size_t n, m;
        for (n = 0; n < this->replicates_; n++)
        {
            for (m = 0; m < M; m++)
            {
                size_t ind = (size_t)(dis(gen)*N);
                if (ind >= N) ind = N - 1;
                memcpy(&X_subset(0, m), &X(0, ind), sizeof(T)*P);
            }

            this->get_initial_guess_sample(X_subset, K, C_for_initial_subset);

            // call kmeans
            this->run(X_subset, K, C_for_initial_subset, IDX, C, sumD);

            memcpy(&C_for_initial(0, 0, n), C.begin(), sizeof(T)*K*P);
        }

        if (this->perform_timing_) gt_timer_.stop();
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::get_initial_guess_cluster(...) ... ");
    }
}

template <typename T>
void kmeans<T>::get_initial_guess_kmeansplusplus(const ArrayType& X, size_t K, ArrayType& C_for_initial)
{
    try
    {
        if (this->perform_timing_) gt_timer_.start("get_initial_guess_kmeansplusplus");

        size_t P = X.get_size(0);
        size_t N = X.get_size(1);

        GADGET_CHECK_THROW(N>K);

        std::random_device rd;
        std::mt19937 gen(rd());
        std::uniform_real_distribution<> dis(0, 1);

        C_for_initial.create(P, K, this->replicates_);
        Gadgetron::clear(C_for_initial);

        ClusterType IDX;
        IDX.resize(N, 0);

        // find the first center
        ArrayType C;
        C.create(P, K);
        Gadgetron::clear(C);

        VectorType norm_C(K, 0);

        ArrayType D;
        D.create(P, N);

        ArrayType D_norm;
        D_norm.create(N);

        ArrayType cumsum_D_norm;
        cumsum_D_norm.create(N);

        ArrayType CX;

        size_t n, i, t, s;

        for (n = 0; n < this->replicates_; n++)
        {
            size_t ind = (size_t)(dis(gen)*N);
            if (ind >= N) ind = N - 1;
            memcpy(&C(0, 0), &X(0, ind), sizeof(T)*P);

            ArrayType aC;
            aC.create(P, C.begin());
            norm_C[0] = Gadgetron::dot(aC, aC,false);

            for (i = 1; i < K; i++)
            {
                // compute distance to the nearest centroid for all data points
                this->compute_dist(X, IDX, C, D);

                // compute norm of distance vector
                this->compute_norm_dist(D, D_norm);

                // compute accumulated distrance
                cumsum_D_norm(0) = D_norm(0);
                for (t = 1; t < N; t++)
                {
                    cumsum_D_norm(t) = cumsum_D_norm(t - 1) + D_norm(t);
                }

                if (std::abs(cumsum_D_norm(N - 1)) < FLT_EPSILON)
                {
                    GERROR_STREAM("std::abs(cumsum_D_norm(N-1))<FLT_EPSILON ... ");
                    // set centroid from i to K

                    for (s = i; s < K; s++)
                    {
                        size_t ind = (size_t)(dis(gen)*N);
                        if (ind >= N) ind = N - 1;
                        memcpy(&C(0, s), &X(0, ind), sizeof(T)*P);
                    }
                    break;
                }

                // convert to probability
                for (t = 0; t < N; t++)
                {
                    cumsum_D_norm(t) /= cumsum_D_norm(N - 1);
                }

                T v = dis(gen);
                for (t = 0; t < N; t++)
                {
                    if (cumsum_D_norm(t)>=v) break;
                }

                memcpy(&C(0, i), &X(0, t), sizeof(T)*P);

                aC.create(P, &C(0, i));
                norm_C[i] = Gadgetron::dot(aC, aC,false);

                // update the IDX
                ArrayType curr_C;
                curr_C.create(P, i+1, C.begin());

                this->update_IDX(X, curr_C, norm_C, IDX);
            }

            memcpy(&C_for_initial(0, 0, n), C.begin(), sizeof(T)*P*K);
        }

        if (this->perform_timing_) gt_timer_.stop();
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::get_initial_guess_kmeansplusplus(...) ... ");
    }
}

template <typename T>
void kmeans<T>::run_replicates(const ArrayType& X, size_t K, const ArrayType& C_for_initial, ClusterType& IDX, ArrayType& C, VectorType& sumD_rep, T& sumD)
{
    try
    {
        Gadgetron::GadgetronTimer timer;
        timer.set_timing_in_destruction(false);

        if (this->perform_timing_) gt_timer_.start("run_replicates");

        size_t P = X.get_size(0);
        size_t N = X.get_size(1);

        GADGET_CHECK_THROW(N>K);
        GADGET_CHECK_THROW(C_for_initial.get_size(0) == P);
        GADGET_CHECK_THROW(C_for_initial.get_size(1) == K);

        size_t R = C_for_initial.get_size(2);

        std::vector<ArrayType> C_rep(R);
        std::vector<ClusterType> IDX_rep(R);

        sumD_rep.resize(R, 0);

        size_t r;
        for (r=0; r<R; r++)
        {
            std::stringstream outs;
            outs << "-----> Kmeans, replicate " << r << " out of " << R;

            if (this->verbose_)
            {
                GDEBUG_STREAM(outs.str());
            }

            ArrayType curr_C_initial;
            curr_C_initial.create(P, K, const_cast<T*>(&C_for_initial(0, 0, r)) );

            if (this->perform_timing_) timer.start(outs.str().c_str());
            this->run(X, K, curr_C_initial, IDX_rep[r], C_rep[r], sumD_rep[r]);
            if (this->perform_timing_) timer.stop();

            if(this->verbose_)
            {
                GDEBUG_STREAM("Kmeans, replicate " << r << " out of " << R << " - " << sumD_rep[r]);
            }
        }

        size_t best_r = 0;
        sumD = sumD_rep[0];
        for (r = 1; r < R; r++)
        {
            if(sumD>sumD_rep[r])
            {
                sumD = sumD_rep[r];
                best_r = r;
            }
        }

        IDX = IDX_rep[best_r];
        C = C_rep[best_r];

        if (this->perform_timing_) gt_timer_.stop();
    }
    catch(...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::run_replicates(...) ... ");
    }
}

template <typename T>
void kmeans<T>::run(const ArrayType& X, size_t K, const ArrayType& C_for_initial, ClusterType& IDX, ArrayType& C, T& sumD)
{
    try
    {
        size_t P = X.get_size(0);
        size_t N = X.get_size(1);

        GADGET_CHECK_THROW(N>K);
        GADGET_CHECK_THROW(C_for_initial.get_size(0) == P);
        GADGET_CHECK_THROW(C_for_initial.get_size(1) == K);

        this->replicates_ = C_for_initial.get_size(2);

        IDX.resize(N, 0);
        C.create(P, K);
        Gadgetron::clear(C);
        sumD  = 0;

        C = C_for_initial;

        VectorType norm_C(K, 0);

        size_t k, p;
        for (k=0; k<K; k++)
        {
            T v = 0;
            for (p=0; p<P; p++)
            {
                v += C(p, k)*C(p, k);
            }

            norm_C[k] = v;
        }

        // first round of clustering
        this->update_IDX(X, C, norm_C, IDX);

        ClusterType prev_IDX;
        ArrayType D, D_norm;

        size_t num_iter = 0;
        T prev_sumD = std::numeric_limits<T>::max();

        std::vector<size_t> cluster_size;

        while (num_iter<=this->max_iter_ &&  this->is_clustering_changed(prev_IDX, IDX))
        {
            prev_IDX = IDX;

            // update the centroid
            this->update_centroid(X, IDX, C, norm_C);
            // update clustering
            this->update_IDX(X, C, norm_C, IDX);

            this->compute_dist(X, IDX, C, D);
            this->compute_norm_dist(D, D_norm);

            // if there are clusters having no member, find a point furthest away from its own cluster centroid
            // replace the empty cluster centroid with this point
            std::vector<size_t> empty_clusters;
            this->has_empty_cluster(C, IDX, cluster_size, empty_clusters);

            if (!empty_clusters.empty())
            {
                if (this->verbose_)
                {
                    GDEBUG_STREAM("Kmeas iteration, found empty cluster : iter - " << num_iter);
                }

                size_t e;
                for (e = 0; e < empty_clusters.size(); e++)
                {
                    // find the most "lonely" point
                    size_t lonely(0);
                    this->find_lonely_point(IDX, D_norm, cluster_size, lonely);

                    // replace this empty centroid with the lonely point, update IDX
                    IDX[lonely] = empty_clusters[e];

                    // updates centroid
                    this->update_centroid(X, IDX, C, norm_C);

                    // update distances
                    this->compute_dist(X, IDX, C, D);
                    this->compute_norm_dist(D, D_norm);
                }
            }

            sumD = 0;
            for (p = 0; p<N; p++)
            {
                sumD += D_norm(p)*D_norm(p);
            }

            if (sumD>prev_sumD)
            {
                if (this->verbose_)
                {
                    GDEBUG_STREAM("Kmeas iteration terminated due to increasted total distance : iter " << num_iter << " - " << sumD << " > " << prev_sumD);
                }

                IDX = prev_IDX;

                break;
            }
            else
            {
                prev_sumD = sumD;
            }

            num_iter++;
        }

        if (this->verbose_)
        {
            GDEBUG_STREAM("Kmeas iteration stopped : iter " << num_iter << " - " << sumD);
        }

        if(this->perform_online_update_)
        {
            this->perform_online_update(X, IDX, C, sumD);
            if (this->verbose_)
            {
                GDEBUG_STREAM("Kmeas online update : " << sumD);
            }
        }
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::run(...) ... ");
    }
}

template <typename T>
void kmeans<T>::compute_dist(const ArrayType& X, const ClusterType& IDX, const ArrayType& C, ArrayType& D)
{
    try
    {
        size_t P = X.get_size(0);
        size_t N = X.get_size(1);

        size_t K = C.get_size(1);

        GADGET_CHECK_THROW(N== IDX.size());

        D.create(P, N);
        Gadgetron::clear(D);

        const T* pX = X.begin();
        const T* pC = C.begin();
        T* pD = D.begin();

        long long p, n;

#pragma omp parallel for default(none) private(n, p) shared(N, IDX, K, P, pD, pX, pC)
        for (n=0; n<N; n++)
        {
            size_t nC = IDX[n];

            if (nC >= K)
            {
                // GWARN_STREAM("nC >= K :" << nC << " - " << K << " for " << n);
                continue;
            }

            for (p=0; p<P; p++)
            {
                pD[p + n*P] = pX[p + n*P] - pC[p + nC*P];
            }
        }
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::compute_dist(...) ... ");
    }
}

template <typename T>
void kmeans<T>::compute_norm_dist(const ArrayType& D, ArrayType& D_norm)
{
    try
    {
        size_t P = D.get_size(0);
        size_t N = D.get_size(1);

        D_norm.create(N);
        Gadgetron::clear(D_norm);

        const T* pD = D.begin();

        long long n, p;

#pragma omp parallel for default(none) private(n, p) shared(N, P, pD, D_norm)
        for(n=0; n<N; n++)
        {
            T v = 0;
            for (p = 0; p<P; p++)
            {
                v += pD[p + n*P] * pD[p + n*P];
            }

            D_norm(n) = std::sqrt(v);
        }
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::compute_norm_dist(...) ... ");
    }
}

template <typename T>
void kmeans<T>::update_IDX(const ArrayType& X, const ArrayType& C, const VectorType& norm_C, ClusterType& IDX)
{
    try
    {
        size_t P = X.get_size(0);
        size_t N = X.get_size(1);

        size_t K = C.get_size(1);

        IDX.resize(N);

        ArrayType CX;
        Gadgetron::gemm(CX, C, true, X, false);

        size_t t, s;

        for (t = 0; t < N; t++)
        {
            for (s = 0; s < K; s++)
            {
                CX(s, t) = 2* CX(s, t) - norm_C[s];
            }

            T maxCX = CX(0, t);
            IDX[t] = 0;
            for (s = 1; s < K; s++)
            {
                if (CX(s, t) > maxCX)
                {
                    maxCX = CX(s, t);
                    IDX[t] = s;
                }
            }
        }
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::update_IDX(...) ... ");
    }
}

template <typename T>
void kmeans<T>::update_centroid(const ArrayType& X, const ClusterType& IDX, ArrayType& C, VectorType& norm_C)
{
    try
    {
        size_t P = X.get_size(0);
        size_t N = X.get_size(1);
        const T* pX = X.begin();

        size_t K = C.get_size(1);

        norm_C.resize(K);

        std::vector<size_t> num_in_C(K, 0);

        Gadgetron::clear(C);
        T* pC = C.begin();

        size_t n, p;
        for (n=0; n<N; n++)
        {
            size_t currK = IDX[n];

            if(currK<K)
            {
                for (p=0; p<P; p++)
                {
                    pC[p+currK*P] += pX[p+n*P];
                }

                num_in_C[currK]++;
            }
            else
            {
                GERROR_STREAM("kmeans, currC>=K, in update_centroid : " << n);
            }
        }

        for (n = 0; n < K; n++)
        {
            T v = 0;
            for (p = 0; p<P; p++)
            {
                if (num_in_C[n] > 0)
                {
                    pC[p + n*P] /= num_in_C[n];
                    v += pC[p + n*P] * pC[p + n*P];
                }
            }

            norm_C[n] = v;
        }
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::update_centroid(...) ... ");
    }
}

template <typename T>
void kmeans<T>::has_empty_cluster(const ArrayType& C, const ClusterType& IDX, std::vector<size_t>& cluster_size, std::vector<size_t>& empty_clusters)
{
    try
    {
        size_t K = C.get_size(1);
        size_t N = IDX.size();

        cluster_size.resize(K, 0);

        size_t n;
        for (n=0; n<N; n++)
        {
            if (IDX[n] < K)
            {
                cluster_size[IDX[n]]++;
            }
        }

        empty_clusters.clear();
        for (n = 0; n < K; n++)
        {
            if(cluster_size[n]==0)
            {
                empty_clusters.push_back(n);
            }
        }
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::has_empty_cluster(...) ... ");
    }
}

template <typename T>
void kmeans<T>::find_lonely_point(const ClusterType& IDX, const ArrayType& D_norm, const std::vector<size_t>& cluster_size, size_t& lonely)
{
    try
    {
        size_t N = IDX.size();

        T maxD (0);
        Gadgetron::maxAbsolute(D_norm, maxD, lonely);

        size_t cluster = IDX[lonely];

        size_t n, c;

        if(cluster_size[cluster]<2)
        {
            // if the picked cluster has only one member ...
            std::vector<size_t> big_clusters;
            for (c=0; c<cluster_size.size(); c++)
            {
                if(cluster_size[c]>2)
                {
                    big_clusters.push_back(c);
                }
            }

            if(big_clusters.empty())
            {
                GADGET_THROW("All clusters have only one member ... ");
            }
            else
            {
                size_t cluster_picked = big_clusters[big_clusters.size() / 2]; // pick one cluster

                // find the furthest point in the picked cluster
                lonely = 0;
                maxD = 0;
                for (n=0; n<N; n++)
                {
                    if(IDX[n]==cluster_picked)
                    {
                        if(D_norm(n)>maxD)
                        {
                            maxD = D_norm(n);
                            lonely = n;
                        }
                    }
                }
            }
        }
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::find_lonely_point(...) ... ");
    }
}

template <typename T>
bool kmeans<T>::is_clustering_changed(const ClusterType&prev_IDX, const ClusterType& IDX)
{
    try
    {
        if (prev_IDX.size() != IDX.size()) return true;

        size_t N = IDX.size();

        size_t n;
        for (n=0; n<N; n++)
        {
            if (prev_IDX[n] != IDX[n])
            {
                return true;
            }
        }

        return false;
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::is_clustering_changed(...) ... ");
        return false;
    }
}

template <typename T>
void kmeans<T>::perform_online_update(const ArrayType& X, ClusterType& IDX, ArrayType& C, T& sumD)
{
    try
    {
        size_t P = X.get_size(0);
        size_t N = X.get_size(1);
        const T* pX = X.begin();

        size_t K = C.get_size(1);

        ArrayType del_cost; // store the delat change of sum cost if reassign a point
        del_cost.create(N, K);

        // count the number of points in each cluster
        std::vector<size_t> num_pt_clusters(K, 0);

        size_t n, p, k;

        num_pt_clusters.resize(K, 0);
        for (n = 0; n < N; n++)
        {
            num_pt_clusters[IDX[n]]++;
        }

        size_t iter(0);

        size_t lastmoved = 0;
        size_t nummoved = 0;
        ClusterType prevIDX, newIDX(IDX);

        while (iter < this->max_iter_)
        {
            // for every cluster K and every point N
            // compute change of delta sum cost
            for (k = 0; k < K; k++)
            {
                for (n = 0; n < N; n++)
                {
                    T v;
                    if (IDX[n] == k)
                    {
                        if (num_pt_clusters[k] > 1)
                            v = (T)num_pt_clusters[k] / (T)(num_pt_clusters[k] - 1);
                        else
                            v = 1;
                    }
                    else
                        v = (T)num_pt_clusters[k] / (T)(num_pt_clusters[k] + 1);

                    T t(0), d = 0;
                    for (p = 0; p < P; p++)
                    {
                        t = X(p, n) - C(p, k);
                        d += t*t;
                    }

                    del_cost(n, k) = v * d;
                }
            }

            prevIDX = IDX;

            // get the new IDX
            for (n = 0; n < N; n++)
            {
                newIDX[n] = 0;
                T min_del_cost = del_cost(n, 0);
                for (k = 1; k < K; k++)
                {
                    if(del_cost(n, k) < min_del_cost)
                    {
                        newIDX[n] = k;
                        min_del_cost = del_cost(n, k);
                    }
                }
            }

            // marked the moving points
            std::vector<size_t> moved;
            for (n = 0; n < N; n++)
            {
                if(prevIDX[n] != newIDX[n])
                {
                    moved.push_back(n);
                }
            }

            // if no candidates to move, stop
            if(moved.empty())
            {
                iter++;
                break;
            }

            // pick a point to move
            int moved_ind = N+1;
            int tt(0);
            for (size_t ii = 0; ii < moved.size(); ii++)
            {
                tt = (int)moved[ii] - (int)lastmoved - 1;
                if (tt < 0) tt += N;
                if (tt >= N) tt -= N;

                tt += lastmoved;

                if (tt < moved_ind) moved_ind = tt;
            }

            if (tt < 0) tt += N;
            if (tt >= N) tt -= N;
            moved_ind = (size_t)(tt);

            if(moved_ind<=lastmoved)
            {
                iter++;
                if (iter >= this->max_iter_) break;
                nummoved = 0;
            }

            nummoved++;
            lastmoved = moved_ind;

            size_t oidx = IDX[moved_ind];
            size_t nidx = newIDX[moved_ind];

            sumD = sumD + del_cost(moved_ind, nidx) - del_cost(moved_ind, oidx);

            IDX[moved_ind] = nidx;

            num_pt_clusters[oidx]--;
            num_pt_clusters[nidx]++;

            for (p=0; p<P; p++)
            {
                C(p, nidx) = C(p, nidx) + (X(p, moved_ind) - C(p, nidx)) / num_pt_clusters[nidx];
                C(p, oidx) = C(p, oidx) - (X(p, moved_ind) - C(p, oidx)) / num_pt_clusters[oidx];
            }
        }
    }
    catch (...)
    {
        GERROR_STREAM("Exceptions happened in kmeans<T>::perform_online_update(...) ... ");
    }
}

// ------------------------------------------------------------
// Instantiation
// ------------------------------------------------------------

template class EXPORTPR kmeans< float >;
template class EXPORTPR kmeans< double >;

}


/*

Copyright (c) 2013, Project OSRM, Dennis Luxen, others
All rights reserved.

Redistribution and use in source and binary forms, with or without modification,
are permitted provided that the following conditions are met:

Redistributions of source code must retain the above copyright notice, this list
of conditions and the following disclaimer.
Redistributions in binary form must reproduce the above copyright notice, this
list of conditions and the following disclaimer in the documentation and/or
other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

*/

// KD Tree coded by Christian Vetter, Monav Project

#ifndef STATICKDTREE_HPP
#define STATICKDTREE_HPP

#include <boost/assert.hpp>
#include <vector>
#include <algorithm>
#include <stack>
#include <limits>

namespace KDTree
{

#define KDTREE_BASESIZE (8)

template <unsigned k, typename T> class BoundingBox
{
  public:
    BoundingBox()
    {
        for (unsigned dim = 0; dim < k; ++dim)
        {
            min[dim] = std::numeric_limits<T>::min();
            max[dim] = std::numeric_limits<T>::max();
        }
    }

    T min[k];
    T max[k];
};

struct NoData
{
};

template <unsigned k, typename T> class EuclidianMetric
{
  public:
    double operator()(const T left[k], const T right[k])
    {
        double result = 0;
        for (unsigned i = 0; i < k; ++i)
        {
            double temp = (double)left[i] - (double)right[i];
            result += temp * temp;
        }
        return result;
    }

    double operator()(const BoundingBox<k, T> &box, const T point[k])
    {
        T nearest[k];
        for (unsigned dim = 0; dim < k; ++dim)
        {
            if (point[dim] < box.min[dim])
                nearest[dim] = box.min[dim];
            else if (point[dim] > box.max[dim])
                nearest[dim] = box.max[dim];
            else
                nearest[dim] = point[dim];
        }
        return operator()(point, nearest);
    }
};

template <unsigned k, typename T, typename Data = NoData, typename Metric = EuclidianMetric<k, T>>
class StaticKDTree
{
  public:
    struct InputPoint
    {
        T coordinates[k];
        Data data;
        bool operator==(const InputPoint &right)
        {
            for (int i = 0; i < k; i++)
            {
                if (coordinates[i] != right.coordinates[i])
                    return false;
            }
            return true;
        }
    };

    explicit StaticKDTree(std::vector<InputPoint> *points)
    {
        BOOST_ASSERT(k > 0);
        BOOST_ASSERT(points->size() > 0);
        size = points->size();
        kdtree = new InputPoint[size];
        for (Iterator i = 0; i != size; ++i)
        {
            kdtree[i] = points->at(i);
            for (unsigned dim = 0; dim < k; ++dim)
            {
                if (kdtree[i].coordinates[dim] < boundingBox.min[dim])
                    boundingBox.min[dim] = kdtree[i].coordinates[dim];
                if (kdtree[i].coordinates[dim] > boundingBox.max[dim])
                    boundingBox.max[dim] = kdtree[i].coordinates[dim];
            }
        }
        std::stack<Tree> s;
        s.push(Tree(0, size, 0));
        while (!s.empty())
        {
            Tree tree = s.top();
            s.pop();

            if (tree.right - tree.left < KDTREE_BASESIZE)
                continue;

            Iterator middle = tree.left + (tree.right - tree.left) / 2;
            std::nth_element(
                kdtree + tree.left, kdtree + middle, kdtree + tree.right, Less(tree.dimension));
            s.push(Tree(tree.left, middle, (tree.dimension + 1) % k));
            s.push(Tree(middle + 1, tree.right, (tree.dimension + 1) % k));
        }
    }

    ~StaticKDTree() { delete[] kdtree; }

    bool NearestNeighbor(InputPoint *result, const InputPoint &point)
    {
        Metric distance;
        bool found = false;
        double nearestDistance = std::numeric_limits<T>::max();
        std::stack<NNTree> s;
        s.push(NNTree(0, size, 0, boundingBox));
        while (!s.empty())
        {
            NNTree tree = s.top();
            s.pop();

            if (distance(tree.box, point.coordinates) >= nearestDistance)
                continue;

            if (tree.right - tree.left < KDTREE_BASESIZE)
            {
                for (unsigned i = tree.left; i < tree.right; i++)
                {
                    double newDistance = distance(kdtree[i].coordinates, point.coordinates);
                    if (newDistance < nearestDistance)
                    {
                        nearestDistance = newDistance;
                        *result = kdtree[i];
                        found = true;
                    }
                }
                continue;
            }

            Iterator middle = tree.left + (tree.right - tree.left) / 2;

            double newDistance = distance(kdtree[middle].coordinates, point.coordinates);
            if (newDistance < nearestDistance)
            {
                nearestDistance = newDistance;
                *result = kdtree[middle];
                found = true;
            }

            Less comperator(tree.dimension);
            if (!comperator(point, kdtree[middle]))
            {
                NNTree first(middle + 1, tree.right, (tree.dimension + 1) % k, tree.box);
                NNTree second(tree.left, middle, (tree.dimension + 1) % k, tree.box);
                first.box.min[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
                second.box.max[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
                s.push(second);
                s.push(first);
            }
            else
            {
                NNTree first(middle + 1, tree.right, (tree.dimension + 1) % k, tree.box);
                NNTree second(tree.left, middle, (tree.dimension + 1) % k, tree.box);
                first.box.min[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
                second.box.max[tree.dimension] = kdtree[middle].coordinates[tree.dimension];
                s.push(first);
                s.push(second);
            }
        }
        return found;
    }

  private:
    using Iterator = unsigned;
    struct Tree
    {
        Iterator left;
        Iterator right;
        unsigned dimension;
        Tree() {}
        Tree(Iterator l, Iterator r, unsigned d) : left(l), right(r), dimension(d) {}
    };
    struct NNTree
    {
        Iterator left;
        Iterator right;
        unsigned dimension;
        BoundingBox<k, T> box;
        NNTree() {}
        NNTree(Iterator l, Iterator r, unsigned d, const BoundingBox<k, T> &b)
            : left(l), right(r), dimension(d), box(b)
        {
        }
    };
    class Less
    {
      public:
        explicit Less(unsigned d)
        {
            dimension = d;
            BOOST_ASSERT(dimension < k);
        }

        bool operator()(const InputPoint &left, const InputPoint &right)
        {
            BOOST_ASSERT(dimension < k);
            return left.coordinates[dimension] < right.coordinates[dimension];
        }

      private:
        unsigned dimension;
    };

    BoundingBox<k, T> boundingBox;
    InputPoint *kdtree;
    Iterator size;
};
}

#endif // STATICKDTREE_HPP


#include "PrimesAccumulator.hpp"

#include <boost/throw_exception.hpp>
#include <sstream>
#include <stdexcept>
#include <string>

#include "primes/PrimesCalculator.hpp"
namespace {
auto createNegativeArgumentExplanation(int passedArgument) {
  std::stringstream ss;
  ss << "Number of primes to sum must be non-negative integer. Passed value: "
     << passedArgument;
  return ss.str();
}
}  // namespace

namespace adder {
PrimesAccumulator::PrimesAccumulator(
    const primes::PrimesCalculator& primesCalculator_)
    : primesCalculator{primesCalculator_} {}

unsigned PrimesAccumulator::sumNthPrimes(int numberOfPrimes) const {
  if (numberOfPrimes < 0) {
    BOOST_THROW_EXCEPTION(std::invalid_argument(
        createNegativeArgumentExplanation(numberOfPrimes)));
  }
  if (numberOfPrimes == 0) {
    return 0;
  }
  auto result = 0u;
  for (auto index = 1u; index <= unsigned(numberOfPrimes); index++) {
    result += primesCalculator.getPrime(index);
  }
  return result;
}
}  // namespace adder


#if !defined(BALSA_GEOMETRY_DATA_STRUCTURES_CONTAINER_OF_CONTAINERS_TO_STACKED_CONTIGUOUS_BUFFER)
#define BALSA_GEOMETRY_DATA_STRUCTURES_CONTAINER_OF_CONTAINERS_TO_STACKED_CONTIGUOUS_BUFFER
#include <Eigen/Core>
#include <range/v3/view/take_exactly.hpp>

#include "stacked_contiguous_buffer.hpp"
#include <concepts>
#include <range/v3/view/transform.hpp>
#include <range/v3/view/partial_sum.hpp>
#include <range/v3/view/join.hpp>
#include <range/v3/algorithm/copy.hpp>

namespace balsa::data_structures {

template<typename IndexType, typename ContainerType>
requires requires(ContainerType t, typename ContainerType::value_type u, typename ContainerType::value_type::value_type v) {
    t.size();
    u.size();
    IndexType(v);
}
StackedContiguousBuffer<IndexType> container_of_containers_to_stacked_contiguous_buffer(const ContainerType &container) {

    StackedContiguousBuffer<IndexType> ret;
    {
        ret._offsets.resize(container.size() + 1);
        ret._offsets(0) = 0;
        auto range = container | ranges::views::transform([](const auto &a) {
                         return a.size();
                     })
                     | ranges::views::partial_sum | ranges::views::take_exactly(ret._offsets.size() - 1);
        ranges::copy(range, ret._offsets.data() + 1);
    }

    {
        ret._buffer.resize(ret._offsets(Eigen::last));

        auto range = container | ranges::views::join | ranges::views::take_exactly(ret._buffer.size());
        ranges::copy(range, ret._buffer.data());
    }
    return ret;
}

}// namespace balsa::geometry::data_structures
#endif


#include <array>
#include <vector>
#include <functional>
#include <algorithm>
#include <iostream>
#include <fstream>
#include <iomanip>

#include <exception>
#include <mutex>
#include <thread>
#include <cmath>
#include <csignal>
#include <cerrno>
#include <sys/stat.h>

#include <Eigen/Dense>

#include <franka/duration.h>
#include <franka/exception.h>
#include <franka/rate_limiting.h>
#include <franka/model.h>
#include <franka/robot.h>
#include <franka/gripper.h>

#include <trajectory_msgs/JointTrajectory.h>
#include <ros/ros.h>
#include "std_msgs/Float64MultiArray.h"
#include "sensor_msgs/JointState.h"

using namespace std;

/*GLOBAL VARIABLE*/
const double object_grasp_width = 0.02;
const std::string ROBOT_IP("172.16.0.2");

const std::array<double, 7> KP = {{60.0, 60.0, 60.0, 60.0, 40.0, 20.0, 10.0}};
const std::array<double, 7> KD = {{10.0, 10.0, 10.0, 8.0, 5.0, 2.0, 2.0}};

std::array<double, 7> Q_INIT, Q_DESIR, QD_DESIR;

/*callback used for receive q_desired*/
void trajCallback(const trajectory_msgs::JointTrajectoryConstPtr& msg){
    auto points = msg->points;
    //Q_DESIR     = points[0].positions;
    cout << "i am in callback " << endl;
}

int main(int argc, char** argv) {
try{
    ros::init(argc, argv, "lowlevel_node");
    ros::NodeHandle nh;
    nh.subscribe("robot1/jointTrajectory", 4, trajCallback);
    ros::Publisher jointsPub, fkPub;
    jointsPub = nh.advertise<sensor_msgs::JointState>("/robot1/joint_states", 2);
    fkPub     = nh.advertise<std_msgs::Float64MultiArray>("/robot1/fk", 2);

    sensor_msgs::JointState msg;
    msg.name.resize(9);
    msg.position.resize(9);
    msg.velocity.resize(9);
    for(size_t i=0; i<7; i++){
        msg.name[i]     = "panda_joint" + to_string(i+1);
    }
    for(size_t i=0; i<2; i++){
        msg.name[i+7]     = "panda_finger_joint" + to_string(i+1);
    }
    std_msgs::Float64MultiArray cur_trans;
    cur_trans.data.resize(16);

    franka::Robot robot(ROBOT_IP);
    franka::Gripper gripper(ROBOT_IP);

    robot.setCollisionBehavior({{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}},
                                {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0, 100.0}},
                                {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}},
                                {{100.0, 100.0, 100.0, 100.0, 100.0, 100.0}});

    /* set stiffness frame (object relative to EE)
     Set to EE instead of the object so that the end effector force can be directly read*/
    robot.setK({{1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0}});

    /*initial pose*/
#if 0
    Q_INIT = {{0, -M_PI_4, 0, -3 * M_PI_4, 0, M_PI_2, M_PI_4}};
    MotionGenerator motion_generator(0.2, q_init);
    gripper.stop(); // to make sure homing executes successfully
#endif
    cout << "gripper"<<endl;
    gripper.homing();
    sleep(2); /*2sec wait for homing operation*/
    double width_gripper_open = gripper.readOnce().max_width;
    if (width_gripper_open < 0.079){
        ROS_ERROR("Failed initialising the gripper! Make sure nothing is between the fingers when homing.");
    }
    //else
    //    robot.control(motion_generator);

    std::thread send_data_thread;
    /*user confirmation*/
    std::cout << "WARNING: Collision thresholds are set to high values. "
              << "Make sure you have the user stop at hand!" << std::endl
              << "Press Enter to continue..." << std::endl;
    std::cin.ignore();

    /***   main func    ***/
    /*load the kinematics and dynamics model*/
    franka::Model model = robot.loadModel();
    franka::RobotState initial_state = robot.readOnce();

    /*q_cmd used in control callback*/
    std::array<double, 7> q_cmd_now = initial_state.q;

    /*define callback for the torque control loop*/
    std::function<franka::Torques(const franka::RobotState&, franka::Duration)> joint_PD_control_callback = [&model, q_cmd_now](const franka::RobotState& robot_state, franka::Duration dt)->franka::Torques
    {
            std::array<double, 7> coriolis = model.coriolis(robot_state);

            std::array<double, 7> tau_jpd;
            for (int i = 0; i < 7; i++)
                tau_jpd[i] = KP[i]*(q_cmd_now[i] - robot_state.q[i]) - KD[i]*robot_state.dq[i] + coriolis[i];
            cout << "i am running fine" << endl;
           return tau_jpd;
    };

    send_data_thread = std::thread([&]()
    {
        franka::RobotState   state;
        franka::GripperState gripper_state;
        std::chrono::steady_clock::time_point wakeup_timer = std::chrono::steady_clock::now();
        while (ros::ok()){
            wakeup_timer += std::chrono::milliseconds(50);
            std::this_thread::sleep_until(wakeup_timer);

        	state         = robot.readOnce();
        	gripper_state = gripper.readOnce();

        	for(size_t i=0; i<7; i++){
                msg.position[i] = state.q[i];
                msg.velocity[i] = state.dq[i];
            }
            for(size_t i=0; i<2; i++){
                msg.position[i+7] = gripper_state.width / 2.0;
                msg.velocity[i+7] = 0;
            }
            msg.header.stamp = ros::Time::now();
            jointsPub.publish(msg);

            cur_trans.data.insert(cur_trans.data.begin(), state.O_T_EE.begin(), state.O_T_EE.end());
            fkPub.publish(cur_trans);
            cur_trans.data.clear();

            //ros::spinOnce();
            //loop_rate.sleep();
        }
        return;
    });
   std::cout << "send_data thread created." << std::endl;
   std::this_thread::sleep_for(std::chrono::milliseconds(200));

    /*run it*/
    robot.control(joint_PD_control_callback);

    if (send_data_thread.joinable()){
        send_data_thread.join();
        std::cout << "[send_data] Thread terminated." << std::endl;
    }
}
catch(exception& e){
    cerr << "error" << endl;
}

    std::cout << "Quitted robot control loop." << std::endl;
    //franka::MotionFinished(tau_cmd_rate_limited);
    std::cout << "Program ended." << std::endl;
    return 0;
}

// define thread for receive q_d, qdot_d
//std::thread receive_thread([&Q_DSIR, &nh]() {
//});


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.
/*
 *  rscale.hpp
 *  MTL4
 *
 *  Created by Hui Li (huil@Princeton.EDU)
 *
 */

#ifndef MTL_RSCALE_INCLUDE
#define MTL_RSCALE_INCLUDE

#include <boost/numeric/mtl/mtl_fwd.hpp>
#include <boost/numeric/mtl/concept/std_concept.hpp>
#include <boost/numeric/mtl/matrix/map_view.hpp>
#include <boost/numeric/mtl/vector/map_view.hpp>
#include <boost/numeric/mtl/utility/algebraic_category.hpp>
#include <boost/numeric/mtl/interface/vpt.hpp>

namespace mtl { namespace tfunctor {

    // AlgebraicCategory is by default tag::scalar
    template <typename Value1, typename Value2, typename AlgebraicCategory>
    struct rscale
    {
	typedef typename Multiplicable<Value1, Value2>::result_type result_type;

	explicit rscale(const Value2& v2) : v2(v2) {}

	result_type operator() (const Value1& v1) const
	{
	    return v1 * v2;
	}

	Value2 value() const { return v2; }
    private:
	Value2 v2;
    };

    template <typename Matrix, typename Value2>
    struct rscale<Matrix, Value2, tag::matrix>
    {
	typedef matrix::rscaled_view<Matrix,Value2> result_type;

	explicit rscale(const Value2& v2) : v2(v2) {}

	result_type operator() (const Matrix& matrix) const
	{
	    return result_type(matrix, v2);
	}
    private:
	Value2 v2;
    };

    template <typename Vector, typename Value2>
    struct rscale<Vector, Value2, tag::vector>
    {
	typedef vector::rscaled_view<Vector, Value2> result_type;

	explicit rscale(const Value2& v2) : v2(v2) {}

	result_type operator() (const Vector& vector) const
	{
	    return result_type(vector, v2);
	}
    private:
	Value2 v2;
    };

} // namespace tfunctor

namespace matrix {

    template <typename Value1, typename Value2>
    typename tfunctor::rscale<Value1, Value2, typename traits::algebraic_category<Value1>::type>::result_type
    inline rscale(const Value1& value1, const Value2& value2)
    {
	vampir_trace<4017> tracer;
	return tfunctor::rscale<Value1, Value2, typename traits::algebraic_category<Value1>::type>(value2)(value1);
    }
}

namespace vector {

    template <typename Value1, typename Value2>
    typename tfunctor::rscale<Value1, Value2, typename traits::algebraic_category<Value1>::type>::result_type
    inline rscale(const Value1& value1, const Value2& value2)
    {
	vampir_trace<2024> tracer;
	return tfunctor::rscale<Value1, Value2, typename traits::algebraic_category<Value1>::type>(value2)(value1);
    }
}

} // namespace mtl

#endif // MTL_RSCALE_INCLUDE


/* test_hyperexponential.cpp
 *
 * Copyright Marco Guazzone 2014
 * Distributed under the Boost Software License, Version 1.0. (See
 * accompanying file LICENSE_1_0.txt or copy at
 * http://www.boost.org/LICENSE_1_0.txt)
 *
 * \author Marco Guazzone (marco.guazzone@gmail.com)
 *
 */

#include <boost/exception/diagnostic_information.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/math/distributions/hyperexponential.hpp>
#include <boost/numeric/conversion/cast.hpp>
#include <boost/random/hyperexponential_distribution.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_int.hpp>
#include <boost/random/uniform_real.hpp>
#include <boost/range/numeric.hpp>
#include <cstring>
#include <iostream>
#include <numeric>
#include <vector>

#include "statistic_tests.hpp"

// This test unit is similar to the one in "test_real_distribution.ipp", which
// has been changed for the hyperexponential distribution.
// We cannot directly use the original test suite since it doesn't work for
// distributions with vector parameters.
// Also, this implementation has been inspired by the test unit for the
// discrete_distribution class.

#ifndef BOOST_RANDOM_P_CUTOFF
# define BOOST_RANDOM_P_CUTOFF 0.99
#endif

#define BOOST_RANDOM_HYPEREXP_NUM_PHASES_MIN 1
#define BOOST_RANDOM_HYPEREXP_NUM_PHASES_MAX 3
#define BOOST_RANDOM_HYPEREXP_PROBABILITY_MIN 0.01
#define BOOST_RANDOM_HYPEREXP_PROBABILITY_MAX 1.0
#define BOOST_RANDOM_HYPEREXP_RATE_MIN 0.0001
#define BOOST_RANDOM_HYPEREXP_RATE_MAX 1000.0
#define BOOST_RANDOM_HYPEREXP_NUM_TRIALS 1000000ll

namespace /*<unnamed>*/ { namespace detail {

template <typename T>
std::vector<T>& normalize(std::vector<T>& v)
{
	if (v.size() == 0)
	{
		return v;
	}

    const T sum = std::accumulate(v.begin(), v.end(), static_cast<T>(0));
	T final_sum = 0;

    const typename std::vector<T>::iterator end = --v.end();
    for (typename std::vector<T>::iterator it = v.begin();
         it != end;
         ++it)
    {
        *it /= sum;
    }
	*end = 1 - final_sum;  // avoids round off errors, ensures the probs really do sum to 1.

    return v;
}

template <typename T>
std::vector<T> normalize_copy(std::vector<T> const& v)
{
	std::vector<T> vv(v);

    return normalize(vv);
}

template <typename T, typename DistT, typename EngineT>
std::vector<T> make_random_vector(std::size_t n, DistT const& dist, EngineT& eng)
{
    std::vector<T> res(n);

    for (std::size_t i = 0; i < n; ++i)
    {
        res[i] = dist(eng);
    }

    return res;
}

}} // Namespace <unnamed>::detail

bool do_test(std::vector<double> const& probabilities,
             std::vector<double> const& rates,
             long long max,
             boost::mt19937& gen)
{
    std::cout << "running hyperexponential(probabilities,rates) " << max << " times: " << std::flush;

    boost::math::hyperexponential_distribution<> expected(probabilities, rates);

    boost::random::hyperexponential_distribution<> dist(probabilities, rates);

    kolmogorov_experiment test(boost::numeric_cast<int>(max));
    boost::variate_generator<boost::mt19937&, boost::random::hyperexponential_distribution<> > vgen(gen, dist);

    const double prob = test.probability(test.run(vgen, expected));

    const bool result = prob < BOOST_RANDOM_P_CUTOFF;
    const char* err = result? "" : "*";
    std::cout << std::setprecision(17) << prob << err << std::endl;

    std::cout << std::setprecision(6);

    return result;
}

bool do_tests(int repeat, int max_num_phases, double max_rate, long long trials)
{
    boost::mt19937 gen;
    boost::random::uniform_int_distribution<> num_phases_dist(BOOST_RANDOM_HYPEREXP_NUM_PHASES_MIN, max_num_phases);

    int errors = 0;
    for (int i = 0; i < repeat; ++i)
    {
        const int num_phases = num_phases_dist(gen);
        boost::random::uniform_real_distribution<> prob_dist(BOOST_RANDOM_HYPEREXP_PROBABILITY_MIN, BOOST_RANDOM_HYPEREXP_PROBABILITY_MAX);
        boost::random::uniform_real_distribution<> rate_dist(BOOST_RANDOM_HYPEREXP_RATE_MIN, max_rate);

        const std::vector<double> probabilities = detail::normalize_copy(detail::make_random_vector<double>(num_phases, prob_dist, gen));
        const std::vector<double> rates = detail::make_random_vector<double>(num_phases, rate_dist, gen);

        if (!do_test(probabilities, rates, trials, gen))
        {
            ++errors;
        }
    }
    if (errors != 0)
    {
        std::cout << "*** " << errors << " errors detected ***" << std::endl;
    }

    return errors == 0;
}

int usage()
{
    std::cerr << "Usage: test_hyperexponential"
        " -r <repeat>"
        " -num_phases"
        " <max num_phases>"
        " -rate"
        " <max rate>"
        " -t <trials>" << std::endl;
    return 2;
}

template<class T>
bool handle_option(int& argc, char**& argv, const char* opt, T& value)
{
    if (std::strcmp(argv[0], opt) == 0 && argc > 1)
    {
        --argc;
        ++argv;
        value = boost::lexical_cast<T>(argv[0]);
        return true;
    }
    else
    {
        return false;
    }
}

int main(int argc, char** argv)
{
    int repeat = 1;
    int max_num_phases = BOOST_RANDOM_HYPEREXP_NUM_PHASES_MAX;
    double max_rate = BOOST_RANDOM_HYPEREXP_RATE_MAX;
    long long trials = BOOST_RANDOM_HYPEREXP_NUM_TRIALS;

    if (argc > 0)
    {
        --argc;
        ++argv;
    }
    while(argc > 0)
    {
        if (argv[0][0] != '-')
        {
            return usage();
        }
        else if (!handle_option(argc, argv, "-r", repeat)
                 && !handle_option(argc, argv, "-num_phases", max_num_phases)
                 && !handle_option(argc, argv, "-rate", max_rate)
                 && !handle_option(argc, argv, "-t", trials))
        {
            return usage();
        }
        --argc;
        ++argv;
    }

    try
    {
        if (do_tests(repeat, max_num_phases, max_rate, trials))
        {
            return 0;
        }
        else
        {
            return EXIT_FAILURE;
        }
    }
    catch(...)
    {
        std::cerr << boost::current_exception_diagnostic_information() << std::endl;
        return EXIT_FAILURE;
    }
}


#include "domain/definition.h"

#include <cmath>
#include <functional>
#include <memory>

#include <gtest/gtest.h>

#include <deal.II/grid/tria.h>
#include <deal.II/fe/fe_q.h>
#include <deal.II/grid/grid_generator.h>

#include "test_helpers/gmock_wrapper.h"
#include "domain/mesh/tests/mesh_mock.h"
#include "domain/finite_element/tests/finite_element_mock.h"
#include "problem/parameter_types.h"

namespace {

using namespace bart;

using ::testing::_;
using ::testing::NiceMock;

template <typename DimensionWrapper>
class DomainDefinitionTest : public ::testing::Test {
 protected:
  static constexpr int dim = DimensionWrapper::value;
  std::unique_ptr<bart::domain::mesh::MeshMock<dim>> mesh_ptr;
  std::unique_ptr<NiceMock<bart::domain::mesh::MeshMock<dim>>> nice_mesh_ptr;
  std::shared_ptr<bart::domain::finite_element::FiniteElementMock<dim>> fe_ptr;

  void SetUp() override;
};

TYPED_TEST_CASE(DomainDefinitionTest, bart::testing::AllDimensions);

template <typename DimensionWrapper>
void DomainDefinitionTest<DimensionWrapper>::SetUp() {
  mesh_ptr = std::make_unique<bart::domain::mesh::MeshMock<dim>>();
  nice_mesh_ptr = std::make_unique<NiceMock<bart::domain::mesh::MeshMock<dim>>>();
  fe_ptr = std::make_shared<bart::domain::finite_element::FiniteElementMock<dim>>();
}

TYPED_TEST(DomainDefinitionTest, Constructor) {

  using Discretization = problem::DiscretizationType;
  std::array<Discretization, 2> discretizations{
    Discretization::kContinuousFEM, Discretization::kDiscontinuousFEM};

  for (const auto discretization : discretizations) {
    bart::domain::Definition<this->dim> test_domain(std::move(this->mesh_ptr),
                                                    this->fe_ptr,
                                                    discretization);
    // Verify ownership has been taken by constructor
    EXPECT_EQ(this->mesh_ptr, nullptr);
    EXPECT_EQ(this->fe_ptr.use_count(), 2);
    EXPECT_EQ(test_domain.discretization_type(), discretization);
  }
}

TYPED_TEST(DomainDefinitionTest, SetUpMesh) {
  EXPECT_CALL(*this->mesh_ptr, FillTriangulation(_));
  EXPECT_CALL(*this->mesh_ptr, FillMaterialID(_));
  EXPECT_CALL(*this->mesh_ptr, has_material_mapping()).
      WillOnce(::testing::Return(true));
  EXPECT_CALL(*this->mesh_ptr, FillBoundaryID(_));

  bart::domain::Definition<this->dim> test_domain(std::move(this->mesh_ptr),
                                          this->fe_ptr);

  EXPECT_NO_THROW(test_domain.SetUpMesh(););
}

TYPED_TEST(DomainDefinitionTest, SetUpMeshMaterialMappingError) {
  EXPECT_CALL(*this->nice_mesh_ptr, has_material_mapping()).
      WillOnce(::testing::Return(false));

  bart::domain::Definition<this->dim> test_domain(std::move(this->nice_mesh_ptr),
                                            this->fe_ptr);
  EXPECT_ANY_THROW(test_domain.SetUpMesh(););
}

template <typename DimensionWrapper>
class DomainDefinitionDOFTest : public DomainDefinitionTest<DimensionWrapper> {
 protected:
  static constexpr int dim = DimensionWrapper::value;
  DomainDefinitionDOFTest() : fe(1) {};
  dealii::Triangulation<dim> triangulation;
  dealii::FE_Q<dim> fe;
  int n_cells_;

  void SetUp() override;
  static void SetTriangulation(dealii::Triangulation<dim> &to_fill) {
    dealii::GridGenerator::hyper_cube(to_fill, -1, 1);
    if (dim == 1) {
      to_fill.refine_global(4);
    } else {
      to_fill.refine_global(2);
    }
  }
};

TYPED_TEST_CASE(DomainDefinitionDOFTest, bart::testing::AllDimensions);

template <typename DimensionWrapper>
void DomainDefinitionDOFTest<DimensionWrapper>::SetUp() {
  DomainDefinitionTest<DimensionWrapper>::SetUp();
}

TYPED_TEST(DomainDefinitionDOFTest, SetUpDOFTestMPI) {
  EXPECT_CALL(*this->nice_mesh_ptr, has_material_mapping()).
      WillOnce(::testing::Return(true));
  EXPECT_CALL(*this->nice_mesh_ptr, FillTriangulation(_))
      .WillOnce(::testing::Invoke(this->SetTriangulation));
  EXPECT_CALL(*this->fe_ptr, finite_element())
      .WillOnce(::testing::Return(&this->fe));

  bart::domain::Definition<this->dim> test_domain(std::move(this->nice_mesh_ptr),
                                                  this->fe_ptr);
  test_domain.SetUpMesh();
  test_domain.SetUpDOF();

  EXPECT_EQ(test_domain.total_degrees_of_freedom(),
            test_domain.dof_handler().n_dofs());

  int total_cells = 0;
  for (auto cell = test_domain.dof_handler().begin_active();
       cell != test_domain.dof_handler().end(); ++cell) {
    if (cell->is_locally_owned())
      ++total_cells;
  }

  EXPECT_EQ(test_domain.Cells().size(), total_cells);

  EXPECT_CALL(*this->fe_ptr, dofs_per_cell())
      .Times(2)
      .WillRepeatedly(::testing::Return(4));

  auto matrix = test_domain.GetCellMatrix();
  EXPECT_EQ(matrix.n_rows(), 4);
  EXPECT_EQ(matrix.n_cols(), 4);

  auto vector = test_domain.GetCellVector();
  EXPECT_EQ(vector.size(), 4);
}

TYPED_TEST(DomainDefinitionDOFTest, SystemMatrixMPI) {
  EXPECT_CALL(*this->nice_mesh_ptr, has_material_mapping()).
      WillOnce(::testing::Return(true));
  EXPECT_CALL(*this->nice_mesh_ptr, FillTriangulation(_))
      .WillOnce(::testing::Invoke(this->SetTriangulation));
  EXPECT_CALL(*this->fe_ptr, finite_element())
      .WillOnce(::testing::Return(&this->fe));

  bart::domain::Definition<this->dim> test_domain(std::move(this->nice_mesh_ptr),
                                                  this->fe_ptr);
  test_domain.SetUpMesh();
  test_domain.SetUpDOF();

  auto system_matrix_ptr = test_domain.MakeSystemMatrix();

  ASSERT_NE(system_matrix_ptr, nullptr);
  EXPECT_EQ(system_matrix_ptr->n(), test_domain.locally_owned_dofs().size());
  EXPECT_EQ(system_matrix_ptr->m(), test_domain.locally_owned_dofs().size());
}

TYPED_TEST(DomainDefinitionDOFTest, SystemVectorMPI) {
  EXPECT_CALL(*this->nice_mesh_ptr, has_material_mapping()).
      WillOnce(::testing::Return(true));
  EXPECT_CALL(*this->nice_mesh_ptr, FillTriangulation(_))
      .WillOnce(::testing::Invoke(this->SetTriangulation));
  EXPECT_CALL(*this->fe_ptr, finite_element())
      .WillOnce(::testing::Return(&this->fe));

  bart::domain::Definition<this->dim> test_domain(std::move(this->nice_mesh_ptr),
                                                  this->fe_ptr);
  test_domain.SetUpMesh();
  test_domain.SetUpDOF();

  auto system_vector_ptr = test_domain.MakeSystemVector();

  ASSERT_NE(system_vector_ptr, nullptr);
  EXPECT_EQ(system_vector_ptr->size(), test_domain.locally_owned_dofs().size());
}

} // namespace


//
//  Copyright (c) 2012, Institue of Cancer Research.
//  All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.
//     * Neither the name of Institue of Cancer Research.
//       nor the names of its contributors may be used to endorse or promote
//       products derived from this software without specific prior written
//       permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// For more information on the Plane of Best Fit please see
// http://pubs.acs.org/doi/abs/10.1021/ci300293f
//
//  If this code has been useful to you, please include the reference
//  in any work which has made use of it:

//  Plane of Best Fit: A Novel Method to Characterize the Three-Dimensionality
//  of Molecules, Nicholas C. Firth, Nathan Brown, and Julian Blagg, Journal of
//  Chemical Information and Modeling 2012 52 (10), 2516-2525

//
//
// Created by Nicholas Firth, November 2011
// Modified by Greg Landrum for inclusion in the RDKit distribution November
// 2012
// Further modified by Greg Landrum for inclusion in the RDKit core September
// 2016
//

#include <GraphMol/RDKitBase.h>
#include <GraphMol/MolTransforms/MolTransforms.h>

#include "PBF.h"
#include <Numerics/Matrix.h>
#include <Numerics/SquareMatrix.h>
#include <Numerics/SymmMatrix.h>
#include <boost/foreach.hpp>

#include <Eigen/Dense>

namespace RDKit {
namespace Descriptors {
namespace {

double distanceFromAPlane(const RDGeom::Point3D &pt,
                          const std::vector<double> &plane, double denom) {
  double numer = 0.0;
  numer =
      std::abs(pt.x * plane[0] + pt.y * plane[1] + pt.z * plane[2] + plane[3]);

  return numer / denom;
}

bool getBestFitPlane(const Conformer &conf,
                     const std::vector<RDGeom::Point3D> &points,
                     std::vector<double> &plane,
                     const std::vector<double> *weights) {
  PRECONDITION((!weights || weights->size() >= points.size()),
               "bad weights vector");
  PRECONDITION(plane.size() >= 4, "bad plane");
  RDGeom::Point3D origin(0, 0, 0);
  double wSum = 0.0;

  for (unsigned int i = 0; i < points.size(); ++i) {
    if (weights) {
      double w = (*weights)[i];
      wSum += w;
      origin += points[i] * w;
    } else {
      wSum += 1;
      origin += points[i];
    }
  }
  origin /= wSum;

  Eigen::Matrix3d evects;
  Eigen::Vector3d evals;
  MolTransforms::computePrincipalAxesAndMomentsFromGyrationMatrix(
      conf, evects, evals, false, weights);
  RDGeom::Point3D normal;
  normal.x = evects(0, 0);
  normal.y = evects(1, 0);
  normal.z = evects(2, 0);

  plane[0] = normal.x;
  plane[1] = normal.y;
  plane[2] = normal.z;
  plane[3] = -1 * normal.dotProduct(origin);
  return true;
}

}  // end of anonymous namespace

double PBF(const ROMol &mol, int confId) {
  PRECONDITION(mol.getNumConformers() >= 1, "molecule has no conformers")
  unsigned int numAtoms = mol.getNumAtoms();
  if (numAtoms < 4) return 0;

  const Conformer &conf = mol.getConformer(confId);
  if (!conf.is3D()) return 0;

  std::vector<RDGeom::Point3D> points;
  points.reserve(numAtoms);
  for (unsigned int i = 0; i < numAtoms; ++i) {
    points.push_back(conf.getAtomPos(i));
  }

  std::vector<double> plane(4);
  if (!getBestFitPlane(conf, points, plane, NULL)) {
    // the eigenvalue calculation failed, return 0
    // FIX: throw an exception here?
    return 0.0;
  }

  double denom = 0.0;
  for (unsigned int i = 0; i < 3; ++i) {
    denom += plane[i] * plane[i];
  }
  denom = sqrt(denom);

  double res = 0.0;
  for (unsigned int i = 0; i < numAtoms; ++i) {
    res += distanceFromAPlane(points[i], plane, denom);
  }
  res /= numAtoms;

  return res;
}

}  // end of Descriptors namespace
}  // end of RDKit namespace