#ifndef BSPLINEFITTERIMPL_HPP_
#define BSPLINEFITTERIMPL_HPP_

#include "DiffManifoldBSplineTools.hpp"
#include "../DynamicOrTemplateInt.hpp"
#include "../DiffManifoldBSpline.hpp"

#include <memory>
#include <map>
#include <Eigen/Cholesky>
#include <Eigen/LU>
#include <Eigen/QR>
#include <sparse_block_matrix/linear_solver_cholmod.h>

namespace bsplines{
#define _TEMPLATE template <typename TSpline>
#define _CLASS BSplineFitter <TSpline>

//TODO improve : support different containers of times and points
_TEMPLATE
inline void _CLASS::initUniformSpline(TSpline & spline, const std::vector<time_t> & times, const std::vector<point_t> & points, int numSegments, double lambda, FittingBackend backend){
	SM_ASSERT_GE(Exception, times.size(), 2, "There must be at least two time point pairs");
	initUniformSpline(spline, times.front(), times.back(), times, points, numSegments, lambda, backend);
}

_TEMPLATE
inline void _CLASS::initUniformSpline(TSpline & spline, time_t minT, time_t maxT, const std::vector<time_t> & times, const std::vector<point_t> & points, int numSegments, double lambda, FittingBackend backend){
	spline.assertConstructing();
	SM_ASSERT_GE(Exception, times.front(), minT, "Times must all be >= minT");
	SM_ASSERT_LE(Exception, times.back(), maxT, "Times must all be <= maxT");
	IntervalUniformKnotGenerator<TimePolicy> knotGenerator(spline.getSplineOrder(), minT, maxT, numSegments);
	initSpline(spline, knotGenerator, times, points, lambda, backend);
}

_TEMPLATE
inline DeltaUniformKnotGenerator<typename _CLASS::TimePolicy> _CLASS::initUniformSplineWithKnotDelta(TSpline & spline, const std::vector<time_t> & times, const std::vector<point_t> & points, const duration_t knotDelta, double lambda, FittingBackend backend){
	spline.assertConstructing();
	SM_ASSERT_GE(Exception, times.size(), 2, "There must be at least two time point pairs");
	DeltaUniformKnotGenerator<TimePolicy> knotGenerator(spline.getSplineOrder(), *times.begin(), *--times.end(), knotDelta);
	initSpline(spline, knotGenerator, times, points, lambda, backend);
	return knotGenerator;
}

namespace internal{
	template<typename T>
	class PointerGuard {
		bool _deleteIt;
		T * _ptr;
	 public:
		PointerGuard() : _deleteIt(false), _ptr(nullptr) {}
		inline void init(T* ptr, bool deleteIt = true) {
			_deleteIt = deleteIt;
			_ptr = ptr;
		}
		inline ~PointerGuard(){
			if(_deleteIt) delete _ptr;
		}
		inline T& operator *(){
			return *_ptr;
		}
	};

	template<typename TTime>
	struct MapKnotIndexResolver : public KnotIndexResolver<TTime>{
		std::map<TTime, int> knots;
		inline int getKnotIndexAtTime(TTime t) const { return (--knots.upper_bound(t))->second; };
	};

	template <typename TException, typename TTime>
	inline void checkTimesBounds(const std::vector<TTime>& times)
	{
		if(times.size() > 1){
			TTime min = *times.begin(), max=*--times.end();
			for(TTime t : times){
				SM_ASSERT_LE(TException, min, t, "The time sequence must be bounded by its bounding elements. But " << t << " is smaller than the first element :" << min);
				SM_ASSERT_LE(TException, t, max, "The time sequence must be bounded by its bounding elements. But " << t << " is bigger than the last element :" << max);
			}
		}
	}
}

	_TEMPLATE
	void _CLASS::initSpline(TSpline & spline, KnotGenerator<time_t> & knotGenerator, const std::vector<time_t> & times, const std::vector<point_t> & points, double lambda, FittingBackend fittingBackend){
		spline.assertConstructing();
		const size_t numPoints = points.size();

		SM_ASSERT_EQ(Exception, times.size(), numPoints, "The number of times and the number of points must be equal");
		SM_ASSERT_GE(Exception, numPoints, 2, "There must be at least two time point pairs");
		internal::checkTimesBounds<Exception>(times);

		internal::PointerGuard<const KnotIndexResolver<time_t> > resolverPtr;

		if(knotGenerator.hasKnotIndexResolver()){
			for(size_t i = 0; knotGenerator.hasNextKnot(); i++){
				spline.addKnot(knotGenerator.getNextKnot());
			}
			resolverPtr.init(&knotGenerator.getKnotIndexResolver(), false);
		}else{
			auto mapResolver = new internal::MapKnotIndexResolver<time_t>();
			resolverPtr.init(mapResolver, true);
			for(size_t i = 0; knotGenerator.hasNextKnot(); i++){
				time_t nextKnot = knotGenerator.getNextKnot();
				spline.addKnot(nextKnot);
				mapResolver->knots.insert(std::make_pair(nextKnot, i));
			}
		}

		spline.init();

		calcFittedControlVertices(spline, *resolverPtr, times, points, std::function<scalar_t(int) > (), lambda, 0, fittingBackend);
	}

	_TEMPLATE
	void _CLASS::fitSpline(TSpline & spline, const std::vector<time_t> & times, const std::vector<point_t> & points, double lambda, int fixNFirstRelevantControlVertices, std::function<scalar_t(int i) > weights, FittingBackend fittingBackend, const bool calculateControlVertexOffsets){
		spline.assertEvaluable();
		const size_t numPoints = points.size();
		SM_ASSERT_GE(Exception, numPoints, 1, "The must be at least one time point pair!");
		SM_ASSERT_EQ(Exception, times.size(), numPoints, "The number of times and the number of points must be equal");
		internal::checkTimesBounds<Exception>(times);

		auto lastTime = times[numPoints - 1];
		SM_ASSERT_GE(Exception, times[0], spline.getMinTime(), "The values in times may not exceed the time domain of the spline.");
		SM_ASSERT_LE(Exception, lastTime, spline.getMaxTime(), "The values in times may not exceed the time domain of the spline.");

		internal::MapKnotIndexResolver<time_t> mapResolver;
		int i = 0;
		for(auto it = spline.getFirstRelevantSegmentByLast(spline.getSegmentIterator(times[0])), end = spline.getAbsoluteEnd(); it != end; ++it){
			time_t nextKnot = it->getKnot();
			if(nextKnot > lastTime) break;
			spline.addKnot(nextKnot);
			mapResolver.knots.insert(std::make_pair(nextKnot, i++));
		}

		calcFittedControlVertices(spline, mapResolver, times, points, weights, lambda, fixNFirstRelevantControlVertices, fittingBackend, calculateControlVertexOffsets);
	}

namespace internal{
	template<typename TSpline>
	unsigned int getNumberOfRelevantControlVertices(TSpline& spline, typename TSpline::time_t startTime, const typename TSpline::time_t& upToTime, std::function<void(typename TSpline::SegmentIterator it)> apply = std::function<void(typename TSpline::SegmentIterator it)>())
	{
		unsigned int fixFirstVertices = 0;
		for (auto it = spline.getFirstRelevantSegmentByLast(spline.getSegmentIterator(startTime));it->getKnot() < upToTime;fixFirstVertices++, it++) {
			if(apply) apply(it);
		}
		return fixFirstVertices;
	}
}

	_TEMPLATE
	void _CLASS::extendAndFitSpline(TSpline & spline, KnotGenerator<time_t> & knotGenerator, const std::vector<time_t> & times, const std::vector<point_t> & points, double lambda, unsigned char honorCurrentValuePercentage, FittingBackend fittingBackend, const bool calculateControlVertexOffsets){
		const size_t numPoints = points.size();
		SM_ASSERT_TRUE(Exception, knotGenerator.supportsAppending(), "The knot generator must support appending!");
		SM_ASSERT_GE(Exception, numPoints, 1, "The must be at least one time point pair!");
		SM_ASSERT_EQ(Exception, times.size(), numPoints, "The number of times and the number of points must be equal");
		SM_ASSERT_LE(Exception, honorCurrentValuePercentage, 100, "honorCurrentValuePercentage must be in [0, 100]");

		// first we get the time to which to extend the spline.
		const time_t lastTime = times[times.size() - 1];

		const time_t oldMaxTime = spline.getMaxTime();
		if(lastTime > oldMaxTime){
			knotGenerator.extendBeyondTime(lastTime);
			spline.appendSegments(knotGenerator, -1, nullptr);
		}

		if(honorCurrentValuePercentage == 0 || honorCurrentValuePercentage == 100){
			unsigned int fixFirstVertices;
			if(honorCurrentValuePercentage) {
				if(lastTime > oldMaxTime){
					fixFirstVertices = internal::getNumberOfRelevantControlVertices(spline, times[0], oldMaxTime);
				} else {
					// no new vertices but old should not be touched
					return;
				}
			}
			else{// ignore former values
				fixFirstVertices = 0;
			}
			fitSpline(spline, times, points, lambda, (int)fixFirstVertices, std::function<scalar_t(int i) >(), fittingBackend, calculateControlVertexOffsets);
		}else{
			std::vector<scalar_t> weights;
			scalar_t baseWeight = scalar_t(honorCurrentValuePercentage) / scalar_t(100 - honorCurrentValuePercentage);

			std::vector<time_t> newTimes;
			std::vector<point_t> newPoints;
			const time_t minTime = spline.getMinTime();

			const unsigned int relevantOldVCs = internal::getNumberOfRelevantControlVertices<TSpline>(spline, times[0], oldMaxTime,
				[&](typename TSpline::SegmentIterator it)
				{
					it++;
					time_t t(it->getKnot());
					if(t >= minTime){
						newTimes.push_back(t);
						newPoints.push_back(spline.template getEvaluatorAt<0>(t).eval());
						weights.push_back(baseWeight); //TODO improve: normalize the weight base on the amount of points given per segment.
					}
				}
			);
			//TODO optimize: use better ways to concatenate times and points.
			SM_ASSERT_EQ_DBG(std::runtime_error, newTimes.size(), relevantOldVCs, "BUG in "  << __FILE__ << ":" << __LINE__);
			SM_ASSERT_EQ_DBG(std::runtime_error, newPoints.size(), relevantOldVCs, "BUG in "  << __FILE__ << ":" << __LINE__);

			newTimes.resize(times.size() + relevantOldVCs);
			copy(times.begin(), times.end(), newTimes.begin() + relevantOldVCs);
			newPoints.resize(points.size() + relevantOldVCs);
			copy(points.begin(), points.end(), newPoints.begin() + relevantOldVCs);

			fitSpline(spline, newTimes, newPoints, lambda, 0, [relevantOldVCs, &weights](int i){ return i < (int)relevantOldVCs ? weights[i]: scalar_t(1.0);}, fittingBackend, calculateControlVertexOffsets);
		}
	}

	_TEMPLATE
	void _CLASS::calcFittedControlVertices(TSpline & spline, const KnotIndexResolver<time_t> & knotIndexResolver, const std::vector<time_t> & times, const std::vector<point_t> & points, std::function<scalar_t(int i) > weights, double lambda, int fixNFirstRelevantControlVertices, FittingBackend fittingBackend, const bool calculateControlVertexOffsets){
		switch(fittingBackend){
			case FittingBackend::DENSE:
				calcFittedControlVertices<FittingBackend::DENSE>(spline, knotIndexResolver, times, points, weights, lambda, fixNFirstRelevantControlVertices, calculateControlVertexOffsets);
				break;
			case FittingBackend::SPARSE:
				calcFittedControlVertices<FittingBackend::SPARSE>(spline, knotIndexResolver, times, points, weights, lambda, fixNFirstRelevantControlVertices, calculateControlVertexOffsets);
				break;
		}
	}

	namespace internal {
		template< enum FittingBackend FittingBackend_> struct FittingBackendFunctions {
		};

		template<> struct FittingBackendFunctions<FittingBackend::DENSE> {
			typedef typename FittingBackendTraits<FittingBackend::DENSE>::Matrix Matrix;
			typedef typename FittingBackendTraits<FittingBackend::DENSE>::Vector Vector;

			inline static Matrix createA(unsigned int constraintSize, unsigned int coefficientDim, unsigned int /* D */) {
				return Matrix::Zero(constraintSize, coefficientDim);
			}
			inline static Vector createB(int constraintSize) {
				return Vector::Zero(constraintSize);
			}

			inline static auto blockA(Matrix & A, int row, int col, int D, int size = 1) -> decltype(A.block(0, 0, 0, 0)){
				return A.block(row * D, col * D, size * D, size * D);
			}
			inline static typename Vector::SegmentReturnType segmentB(Vector & b, int row, int D){
				return b.segment(row * D, D);
			}
			template <typename DERIVED, typename DERIVED2>
			inline static void addBlockToBlock(Eigen::MatrixBase<DERIVED> & A, int aRow, int aCol, const Eigen::MatrixBase<DERIVED2> & Q, int qRow, int qCol, int rows, int cols, int D, bool segmentInRowVector){
				rows *= D;
				if(segmentInRowVector){
					cols = 1;
				}
				else{
					cols *= D;
				}
				A.block(aRow * D, aCol * D, rows, cols) += Q.block(qRow * D, qCol * D, rows, cols);
			}

			inline static Eigen::VectorXd solve(Matrix & A, Vector & b){
				return A.ldlt().solve(b);
			}
		};

		template<> struct FittingBackendFunctions<FittingBackend::SPARSE> {
			typedef typename FittingBackendTraits<FittingBackend::SPARSE>::Matrix Matrix;
			typedef typename FittingBackendTraits<FittingBackend::SPARSE>::Vector Vector;

			std::vector<int> rows;
			std::vector<int> cols;

			const static bool allocateBlock = true;

			inline Matrix createA(unsigned int constraintSize, unsigned int coefficientDim, unsigned int D) {
				for(unsigned int i = D; i <= constraintSize; i+=D) rows.push_back(i);
				for(unsigned int i = D; i <= coefficientDim; i+=D) cols.push_back(i);
				return Matrix(rows,cols, true);
			}

			inline Vector createB(int /* constraintSize */) {
				std::vector<int> bcols(1);
				bcols[0] = 1;
				return Vector(rows,bcols, true);
			}

			inline static typename Matrix::SparseMatrixBlock & blockA(Matrix & A, int row, int col, int /* D */, int size = 1){
				SM_ASSERT_EQ(std::runtime_error, 1, size, "This function should does not support size != 1");
				return *A.block(row, col, allocateBlock);
			}
			inline static typename Vector::SparseMatrixBlock & segmentB(Vector & b, int row, int /* D */){
				return *b.block(row, 0, allocateBlock);
			}
			template <typename DERIVED>
			inline static void addBlockToBlock(Matrix & A, int aRow, int aCol, const Eigen::MatrixBase<DERIVED> & Q, int qRow, int qCol, int rows, int cols, int D, bool segmentInRowVector){
				for(int i = 0; i < rows; ++i){
					for(int j = 0; j < cols; ++j){
						*A.block(aRow + i, aCol + j, true) += Q.block((qRow + i) * D, (qCol + j) * D, D, segmentInRowVector ? 1 : D);
					}
				}
			}

			inline static Eigen::VectorXd solve(Matrix & A, Vector & b){
				// solve:
				sparse_block_matrix::LinearSolverCholmod<Eigen::MatrixXd> solver;
				solver.init();

				Eigen::VectorXd c(A.rows());
				c.setZero();
				Eigen::VectorXd b_dense = b.toDense();
				bool result = solver.solve(A,&c[0],&b_dense[0]);
				if(!result) {
					c.setZero();
					// fallback => use nonsparse solver:
					std::cout << "Fallback to Dense Solver" << std::endl;
					Eigen::MatrixXd Adense = A.toDense();
					c = Adense.ldlt().solve(b_dense);
				}
				return c;
			}
		};
	}

	_TEMPLATE
	template <enum FittingBackend FittingBackend_>
	void _CLASS::calcFittedControlVertices(TSpline & spline, const KnotIndexResolver<time_t> & knotIndexResolver, const std::vector<time_t> & times, const std::vector<point_t> & points, std::function<scalar_t(int i) > weights, double lambda, int fixNFirstRelevantControlVertices, const bool calculateControlVertexOffsets)
	{
		if(calculateControlVertexOffsets){
			//TODO implement : support calculateControlVertexOffsets in addCurveQuadraticIntegralDiagTo and remove this check!
			SM_ASSERT_EQ(Exception, 0.0, lambda, "control vertex offsets aren't supported together with lambda != 0, yet!");
		}

		SM_ASSERT_GE_DBG(Exception, fixNFirstRelevantControlVertices, 0, "fixNFirstRelevantControlVertices must be nonnegative.");
		const time_t splineMaxTime = spline.getMaxTime(), lastTime = times[times.size() - 1];
		const int splineOrder = spline.getSplineOrder();
		const int minKnotIndex = knotIndexResolver.getKnotIndexAtTime(times[0]) - (splineOrder - 1) - (times[0] == splineMaxTime ? 1 : 0); // get minimal relevant knot's index //TODO improve: remove this irregularity
		const int controlVertexIndexOffset = minKnotIndex + fixNFirstRelevantControlVertices;
		const int numToFitControlVertices = knotIndexResolver.getKnotIndexAtTime(lastTime) - controlVertexIndexOffset + 1 - (lastTime == splineMaxTime ? 1 : 0);//TODO improve: remove this irregularity

		if(numToFitControlVertices <= 0){
			SM_THROW_DBG(Exception, "No control vertices are left over to be optimized.")
			return;
		}

		const int capturedCurrentControlVertices = calculateControlVertexOffsets ? fixNFirstRelevantControlVertices + numToFitControlVertices : fixNFirstRelevantControlVertices;
		const typename TSpline::point_t* currentControlVertices[capturedCurrentControlVertices];
		if(capturedCurrentControlVertices){
			auto it = spline.getFirstRelevantSegmentByLast(spline.getSegmentIterator(times[0]));
			for(int i = 0; i < capturedCurrentControlVertices; ++i){
				currentControlVertices[i] = &it->getControlVertex();
				++it;
			}
		}

		const size_t numPoints = points.size();

		// What is the vector coefficient dimension
		const size_t D = spline.getPointSize();

		// Now we have to solve an Ax = b linear system to determine the correct coefficient vectors.
		size_t coefficientDim = (size_t) numToFitControlVertices * D;

		const int constraintSize = numPoints * D;

		internal::FittingBackendFunctions<FittingBackend_> backend;

		auto A = backend.createA(constraintSize, coefficientDim, D);
		auto b = backend.createB(constraintSize);

		int brow = 0;

		// Add the position constraints.
		for(size_t i = 0; i < numPoints; i++)
		{
			time_t time = times[i];
			scalar_t weight = weights? weights(i) : scalar_t(1.0);
			int knotIndex = knotIndexResolver.getKnotIndexAtTime(time) - controlVertexIndexOffset;
			//TODO optimize : a uniform time spline evaluator could be much faster here!
			typename TSpline::SplineOrderVector bi = spline.template getEvaluatorAt<1>(times[i]).getLocalBi();

			if(weight != scalar_t(1.0)){
				bi *= weight;
			}

			knotIndex -= splineOrder - 1; // get first relevant knot's index
			if(time == splineMaxTime){ //TODO improve: remove this irregularity
				knotIndex --;
			}

			for(int j = 0; j < splineOrder; j ++){
				const int col = knotIndex + j;
				if(col >= 0){ // this control vertex is not fixed
					if(bi[j] != 0.0)
						backend.blockA(A, brow, col, D).diagonal().setConstant(bi[j]);
				}
				if(col < 0 || calculateControlVertexOffsets){
					const int vertexIndex = fixNFirstRelevantControlVertices + col;
					SM_ASSERT_GE_DBG(std::runtime_error, vertexIndex, 0, "BUG in BSplineFitter");
					SM_ASSERT_LT_DBG(std::runtime_error, vertexIndex, capturedCurrentControlVertices, "BUG in BSplineFitter");
					backend.segmentB(b, brow, D) -= (*currentControlVertices[vertexIndex]) * bi[j];
				}
			}

			if(knotIndex >= 0 && ! calculateControlVertexOffsets)
				backend.segmentB(b, brow, D) = points[i] * weight;
			else{
				backend.segmentB(b, brow, D) += points[i] * weight;
			}
			++brow;
		}

		b = A.transpose() * b;
		A = A.transpose() * A;

		if(lambda != 0.0){
			// Add the motion constraint.
			point_t W = point_t::Constant(D, lambda);
			auto start = spline.getSegmentIterator(times[0]);
			auto end = spline.getSegmentIterator(lastTime);
			const int moveCount = fixNFirstRelevantControlVertices - (splineOrder - 1);
			const int movedCount = moveIterator(start, moveCount < 0 ? spline.begin() : end, moveCount);
			const int moveEndCount = splineOrder;
			moveIterator(end, spline.end(), moveEndCount);
			if(moveCount - movedCount <= 0){ // did we not reach end?
				addCurveQuadraticIntegralDiagTo<FittingBackend_>(spline, start, end, movedCount, W, 2, A, b);
			}
		}

		// Solve for the coefficient vector.
		Eigen::VectorXd c = backend.solve(A, b);

		if(calculateControlVertexOffsets){
			const auto & manifold = spline.getManifold();
			spline.manipulateControlVertices([&c, D, &manifold](int i, point_t & v) {
					v += c.block(i * D, 0, D, 1);
					manifold.projectIntoManifold(v);
				}, numToFitControlVertices, times[0], fixNFirstRelevantControlVertices);
		}
		else{
			spline.setControlVertices(c, times[0], fixNFirstRelevantControlVertices);
		}
	}

	_TEMPLATE
	template <enum FittingBackend FittingBackend_>
	void _CLASS::addCurveQuadraticIntegralDiagTo(const TSpline & spline, typename TSpline::SegmentConstIterator start, typename TSpline::SegmentConstIterator end, int startIndex, const point_t & Wdiag, int derivativeOrder, typename internal::FittingBackendTraits<FittingBackend_>::Matrix & toMatrix, typename internal::FittingBackendTraits<FittingBackend_>::Vector & toB)
	{
		SM_ASSERT_EQ(Exception,Wdiag.rows(), (int)spline.getPointSize(), "Wdiag must be of control point size");
		SM_ASSERT_EQ(Exception, toMatrix.rows(), toB.rows(), "Wdiag must be of control point size");

		typedef internal::FittingBackendFunctions<FittingBackend_> Backend;

		const int D = spline.getPointSize();
		const int blocksInQ = spline.getSplineOrder();
		SM_ASSERT_EQ(Exception,Wdiag.rows(), D, "Wdiag must be of control point size");

		const auto qiSize = spline.getSplineOrder() * spline.getPointSize();
		typedef decltype(qiSize) QiSize;
		typedef Eigen::Matrix<double, QiSize::VALUE, QiSize::VALUE> Q_T;
		Q_T Q((int)qiSize, (int)qiSize);

		int blockRow = startIndex;
		const int blockRows = toMatrix.rows() / D;
		for(typename TSpline::SegmentConstIterator sIt = start; sIt != end; sIt++)
		{
			if(FittingBackend_ == FittingBackend::DENSE && blockRow >= 0 && blockRow <= blockRows - blocksInQ){ // all splineOrder many diagonal blocks are contained in the toMatrix
				addOrSetSegmentQuadraticIntegralDiag(spline, Wdiag, sIt, derivativeOrder, Backend::blockA(toMatrix, blockRow, blockRow, D, blocksInQ), true);
			} else { // we have to assign sub blocks individually
				addOrSetSegmentQuadraticIntegralDiag<Q_T &>(spline, Wdiag, sIt, derivativeOrder, Q, false);
				/*
				 * Q and toMatrix (=:M) (both square) have shifted block index space. i_Q + brow = i_M.
				 * Lets x denote the coefficients of all non fixed control vertices
				 * first we need the overlapping region:
				 * its left most index in M is : max(0, 0 + brow).
				 * its right most index in M is : min(cols(M) - 1, cols(Q) - 1 + brow).
				 */
				const int overlappingBlocksLeftIndex = std::max(0, blockRow);
				const int overlappingBlocksRightIndex = std::min(blockRows, blocksInQ + blockRow) - 1;
				const int numOverlappingBlocks = overlappingBlocksRightIndex - overlappingBlocksLeftIndex + 1;
				if(numOverlappingBlocks > 0){
					Backend::addBlockToBlock(toMatrix, overlappingBlocksLeftIndex, overlappingBlocksLeftIndex, Q, overlappingBlocksLeftIndex - blockRow, overlappingBlocksLeftIndex - blockRow, numOverlappingBlocks, numOverlappingBlocks, D, false);
					if(numOverlappingBlocks != blocksInQ){
						/*
						 * Now we have to subtract the effect of the fixed control vertices on the acceleration term.
						 *
						 * [Q_l  Q_x  Q_r]  * [ c_l; x; c_r] = b_x
						 * <=> Q_x * x  = b_x - Q_l * c_l - Q_r * c_r
						 *
						 * We start with c_r because sIt corresponds to its bottom at the moment.
						 */
						auto constantVertexIt = sIt;
						const int Q_r_LeftColumnIndex = overlappingBlocksRightIndex - blockRow + 1;
						const int Q_l_RightColumnIndex = - blockRow - 1;
						for(int i = blocksInQ-1; i >= 0; i--) {
							if(i >= Q_r_LeftColumnIndex || i <= Q_l_RightColumnIndex){ // then vertex is fixed and we have to subtract its effect on the right hand side
								Backend::addBlockToBlock(toB, overlappingBlocksLeftIndex, 0, (-Q.block((overlappingBlocksLeftIndex - blockRow) * D, i * D, numOverlappingBlocks * D, D) * constantVertexIt->getControlVertex()).eval(), 0, 0, numOverlappingBlocks, 1, D, true);
							}
							constantVertexIt--;
						}
					}
				}
			}
			blockRow ++;
		}
	}

	_TEMPLATE
	template<typename M_T>
	inline void _CLASS::addOrSetSegmentQuadraticIntegralDiag(const TSpline & spline, const point_t & Wdiag, typename TSpline::SegmentConstIterator segmentIt, int derivativeOrder, M_T toMatrix, bool add)
	{
		const int D = spline.getPointSize();
		const int splineOrder = spline.getSplineOrder();
		SM_ASSERT_GE_LT(Exception, segmentIt.getKnot(), spline.getMinTime(), spline.getMaxTime(), "Out of range");
		SM_ASSERT_EQ(Exception, Wdiag.rows(), D, "Wdiag must be of control point size");

		typename TSpline::SplineOrderSquareMatrix Dm(splineOrder, splineOrder);
		Dm.setZero(splineOrder, splineOrder);
		spline.computeDiiInto(segmentIt, Dm);
		typename TSpline::SplineOrderSquareMatrix V(splineOrder, splineOrder);
		V.setZero(splineOrder, splineOrder);
		spline.computeViInto(segmentIt, V);

		// Calculate the appropriate derivative version of V
		// using the matrix multiplication version of the derivative.
		for(int i = 0; i < derivativeOrder; i++)
		{
			V = (Dm.transpose() * V * Dm).eval();
		}

		const auto splineOrderTimesPointSize = spline.getSplineOrder() * spline.getPointSize();
		typedef decltype(splineOrderTimesPointSize) SplineOrderTimesPointSize;
		typedef Eigen::Matrix<double, SplineOrderTimesPointSize::VALUE, SplineOrderTimesPointSize::VALUE> MType;

		MType WV((int)splineOrderTimesPointSize, (int)splineOrderTimesPointSize);

		WV.setZero(splineOrderTimesPointSize, splineOrderTimesPointSize);

		for(int d = 0; d < D; d++)
		{
			WV.block(splineOrder*d, splineOrder*d, splineOrder, splineOrder) = Wdiag(d) * V;
		}

		MType M((int)splineOrderTimesPointSize, (int)splineOrderTimesPointSize);
		computeMiInto(spline, segmentIt, M);

		if(add) toMatrix += M.transpose() * WV * M;
		else toMatrix = M.transpose() * WV * M;
	}

	_TEMPLATE
	template <typename M_T>
	void _CLASS::computeBijInto(const TSpline & spline, const typename TSpline::SegmentConstIterator & segmentIndex, int columnIndex, M_T B)
	{
		const int D = spline.getPointSize();
		const int splineOrder = spline.getSplineOrder();

		for(int i = 0; i < D; i++)
		{
			B.block(i*splineOrder,i,splineOrder,1) = segmentIndex->getBasisMatrix().col(columnIndex);
		}
	}

	_TEMPLATE
	template <typename M_T>
	void _CLASS::computeMiInto(const TSpline & spline, const typename TSpline::SegmentConstIterator & segmentIndex, M_T & M)
	{
		const int D = spline.getPointSize();
		const int splineOrder = spline.getSplineOrder();
		const int splineOrderTimesPointSize = splineOrder * D;
		M.setZero();

		for(int j = 0; j < splineOrder; j++)
		{
			computeBijInto(spline, segmentIndex, j, M.block(0, j*D, splineOrderTimesPointSize, D));
		}
	}
}

#undef _TEMPLATE
#undef _CLASS

#endif /* BSPLINEFITTERIMPL_HPP_ */


//    Copyright 2017 Rainer Gemulla
//
//    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
 *
 * Illustrates matrix factorization with parallel SGD. We first create factors and then a data matrix
 * from these factors. This process ensures that we know the best factorization of the input.
 * These matrices are distributed across a cluster. We then try to reconstruct the factors
 * using PSGD.
 *
 * Run with: psgd
 * (make sure to use a production build, otherwise it will be slow)
 */
#include <iostream>
#include <sstream>

#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>

#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/math/distributions/normal.hpp>
#include <boost/random/uniform_real.hpp>
#include <mf/matrix/io/generateDistributedMatrix.h>
#include <mpi2/mpi2.h>
#include <mf/mf.h>

log4cxx::LoggerPtr logger(log4cxx::Logger::getLogger("main"));

using namespace std;
using namespace mf;
using namespace mpi2;
using namespace rg;
using namespace boost::numeric::ublas;

// type of SGD
typedef UpdateTruncate<UpdateNzslNzl2> Update;
typedef UpdateLock<Update> UpdateL;
typedef RegularizeNone Regularize;
typedef SumLoss<NzslLoss, Nzl2Loss> Loss;
typedef NzslLoss TestLoss;

struct TrainingPoint {

	TrainingPoint(){}
	TrainingPoint (mf_size_type i , mf_size_type j, double x):i_(i), j_(j), x_(x){}

	mf_size_type i_;
	mf_size_type j_;
	double x_;

};

//global variables
mf_size_type size1;
mf_size_type size2;
mf_size_type b1;
mf_size_type b2;

mf_size_type minSize1;
mf_size_type minSize2;
mf_size_type remainder1;
mf_size_type remainder2;

int calculateBlocking(double size1, double size2, double nnz, double rank, double cache){

	int SizeOfInt = 8;
	double alpha = cache * 1024;
	double beta = (-1)*(size1+size2)*rank*SizeOfInt;
	double gamma = (-1)*nnz*3*SizeOfInt;

	double delta = beta*beta-4*alpha*gamma;

	int b = ((-1)*beta+sqrt(delta))/(2*alpha);

	return b;
}

// compare function for sort
inline bool compFunction (TrainingPoint point1, TrainingPoint point2) {

	mf_size_type block1_i = min ((mf_size_type)floor(point1.i_/minSize1), b1 -1);
	mf_size_type block2_i = min ((mf_size_type)floor(point2.i_/minSize1), b1 -1);
	mf_size_type block1_j = min ((mf_size_type)floor(point1.j_/minSize2), b2 -1);
	mf_size_type block2_j = min ((mf_size_type)floor(point2.j_/minSize2), b2 -1);
	bool result = false;
	if (block1_i < block2_i) {
		result = true;
	} else if (block1_i == block2_i) {
		if (block1_j < block2_j) {
			result = true;
		} else {
			result =false;
		}
	}
	return result;
}

int main(int argc, char* argv[]) {

	using namespace boost::program_options;
	// initialize mf library and mpi2
	boost::mpi::communicator& world = mfInit(argc, argv);

	// parameters for the factorization
	mf_size_type rank = 10;
	int tasks = 4;
	mf_size_type epochs = 10;
	SgdOrder order = SGD_ORDER_WOR;
	BalanceType balanceType = BALANCE_NONE;
	BalanceMethod balanceMethod = BALANCE_OPTIMAL;
	double lambda = 0.05;
	double eps0 = 0.0125;
	mf_size_type cache = 256;
	string traceFile,traceVar;

	// start mf library
	mfStart();

	if (world.rank() == 0)
	{
#ifndef NDEBUG
		LOG4CXX_WARN(logger, "Warning: Debug mode activated (runtimes may be slow).");
#endif

		string inputMatrixFile;
		string inputRowFacFile;
		string inputColFacFile;
		string outputRowFacFile;
		string outputColFacFile;
		string inputTestMatrixFile;
		string traceFile,traceVar;

		options_description desc("Options");
		desc.add_options()
				("help", "produce help message")
				("epochs", value<mf_size_type>(&epochs)->default_value(10), "number of epochs to run [10]")
				("rank", value<mf_size_type>(&rank)->default_value(10), "rank of factorization [10]")
				("cache", value<mf_size_type>(&cache)->default_value(256), "last level cache (in KB) available per core [256]")
				("lambda", value<double>(&lambda)->default_value(0.05), "lambda")
				("eps0", value<double>(&eps0)->default_value(0.01), "initial step size for BoldDriver")
				("tasks-per-rank", value<int>(&tasks)->default_value(1), "number of concurrent tasks [1]")
				("trace", value<string>(&traceFile)->default_value("trace.R"), "filename of trace [trace]")
				("traceVar", value<string>(&traceVar)->default_value("trace"), "variable name for trace [traceVar]")
				("input-file", value<string>(&inputMatrixFile), "input matrix")
				("input-test-file", value<string>(&inputTestMatrixFile), "input test matrix")
			    ("input-row-file", value<string>(&inputRowFacFile), "input initial row factor")
			    ("input-col-file", value<string>(&inputColFacFile), "input initial column factor")
			    ("output-row-file", value<string>(&outputRowFacFile), "output initial row factor")
			    ("output-col-file", value<string>(&outputColFacFile), "output initial column factor")
				;

		positional_options_description pdesc;
		pdesc.add("input-file", 1);
		pdesc.add("input-test-file", 2);
		pdesc.add("input-row-file", 3);
		pdesc.add("input-col-file", 4);

		variables_map vm;
		store(command_line_parser(argc, argv).options(desc).positional(pdesc).run(), vm);
		notify(vm);

		if (vm.count("help") || vm.count("input-file")==0) {
			cout << "CSGD with NZL2  [options] <input-file> " << endl << endl;
			cout << desc << endl;
			return 1;
		}

		if (vm.count("output-row-file") == 0) { outputRowFacFile = ""; }
		if (vm.count("output-col-file") == 0) { outputColFacFile = ""; }

		LOG4CXX_INFO(logger, "Using " << tasks << " parallel tasks");

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

		Random32 random; // note: this takes a default seed (not randomized!)

// 		SparseMatrix v;
// 		readMatrix(inputMatrixFile, v, MM_COORD);

// 		SparseMatrix vTest;
// 		readMatrix(inputTestMatrixFile, vTest, MM_COORD);
//
// 		DenseMatrix w;
// 		readMatrix(inputRowFacFile, w, MM_ARRAY);
// 		DenseMatrixCM h;
// 		readMatrix(inputColFacFile, h, MM_ARRAY);

		Timer t;
		t.start();
		std::vector<DistributedSparseMatrix> dataVector=getDataMatrices<SparseMatrix>(inputMatrixFile,"V",true,
				tasks, 1, 1, 1, true, false, &inputTestMatrixFile);

		SparseMatrix& v = *mpi2::env().get<SparseMatrix>(dataVector[0].blocks()(0,0).var());
		SparseMatrix& vTest = *mpi2::env().get<SparseMatrix>(dataVector[1].blocks()(0,0).var());

		std::pair<DistributedDenseMatrix, DistributedDenseMatrixCM> factorsPair= getFactors(inputRowFacFile,
			inputColFacFile,  tasks, 1, 1, 1, true);
		DenseMatrix w = *mpi2::env().get<DenseMatrix>(factorsPair.first.blocks()(0,0).var());
		DenseMatrixCM h = *mpi2::env().get<DenseMatrixCM>(factorsPair.second.blocks()(0,0).var());
		t.stop();
		LOG4CXX_INFO(logger, "Total time for loading matrices: " << t);

		mf_size_type nnz = v.nnz();
		SparseMatrix::index_array_type& vIndex1 = rowIndexData(v);
		SparseMatrix::index_array_type& vIndex2 = columnIndexData(v);
		SparseMatrix::value_array_type& vValues = v.value_data();

		// global variables
		size1 = v.size1();
		size2 = v.size2();
		int b = calculateBlocking(size1, size2, nnz, rank, cache);
		b1 = b;
		b2 = b;

		LOG4CXX_INFO(logger, "Using a " << b << " x "<< b<<" stratification");

		minSize1 = floor(size1/b1);
		minSize2 = floor(size2/b2);
		remainder1 = size1 % b1;
		remainder2 = size2 % b2;

		LOG4CXX_INFO(logger, "Data matrix: "
			<< size1 << " x " << size2 << ", " << nnz << " nonzeros");
		LOG4CXX_INFO(logger, "minSize1: " << minSize1 << " remainder1: " << remainder1);
		LOG4CXX_INFO(logger, "minSize2: " << minSize2 << " remainder2: " << remainder2);

		// sort v
		LOG4CXX_INFO(logger, "Sorting v ... ");
		std::vector<TrainingPoint> points;
		for (mf_size_type i = 0; i < v.nnz(); i++) {
			TrainingPoint point = TrainingPoint(vIndex1[i], vIndex2[i], vValues[i]);
			points.push_back(point);
		}
		sort(points.begin(), points.end(), compFunction);

		// compute nnz per block
		LOG4CXX_INFO(logger, "Computing nnz per block ... ");
		std::vector<mf_size_type> nnzPerBlock;
		for (mf_size_type i = 0; i < (b1*b2); i++) nnzPerBlock.push_back(0);
		for (mf_size_type i = 0; i < points.size(); i++) {
			vIndex1[i] = points.at(i).i_;
			vIndex2[i] = points.at(i).j_;
			vValues[i] = points.at(i).x_;

			mf_size_type block_i = min ((mf_size_type)floor(points.at(i).i_/minSize1), b1 -1);
			mf_size_type block_j = min ((mf_size_type)floor(points.at(i).j_/minSize2), b2 -1);
			mf_size_type blockId = (b2 * block_i) + block_j;

			if ((block_i > b1) || (block_j > b2)) {
				LOG4CXX_INFO(logger, "block_i " << block_i << " = floor( " << points.at(i).i_ << " / " << minSize1);
				LOG4CXX_INFO(logger, "block_j " << block_j << " = floor( " << points.at(i).j_ << " / " << minSize2);
			}
			nnzPerBlock[blockId]++;
		}
		LOG4CXX_INFO(logger, "nnzPerBlock of size: " << nnzPerBlock.size());

		// compute permutation
		LOG4CXX_INFO(logger, "Computing permutation ... ");
		std::vector<mf_size_type> permutation;
		for (mf_size_type i = 0; i < v.nnz(); i++) {
			permutation.push_back(i);
		}

		// compute offsets
		LOG4CXX_INFO(logger, "Computing offsets ... ");
		std::vector<mf_size_type> offsets((b1*b2)+1); offsets.clear();
		mf_size_type offset = 0;
		offsets.push_back(offset);
		for (mf_size_type i = 1; i < (b1*b2)+1; i++) {
			offset += nnzPerBlock[i-1];
			offsets.push_back(offset);
		}

		Update update = Update(UpdateNzslNzl2(lambda), -100, 100); // truncate for numerical stability
		Regularize regularize;
		Loss loss((NzslLoss()), Nzl2Loss(lambda));
		TestLoss testLoss;

		FactorizationData<> testData(vTest, w, h);
		LOG4CXX_INFO(logger, "Initial test loss: " << testLoss(testData));

		// initialize

		StratifiedPsgdRunner stratifiedPsgdRunner (random, permutation, offsets);
		StratifiedPsgdJob<Update,Regularize> stratifiedPsgdJob(v, w, h, update, regularize, order, tasks);
		BoldDriver decay(eps0);
		Trace trace;

		t.start();
		stratifiedPsgdRunner.run(stratifiedPsgdJob, loss, epochs, decay, trace, balanceType, balanceMethod, &testData, &testLoss);
		t.stop();
		LOG4CXX_INFO(logger, "Total time: " << t);

		// write trace to an R file
		LOG4CXX_INFO(logger, "Writing trace to " << traceFile);
		trace.toRfile(traceFile, traceVar);

					// write computed factors to file
				if (outputRowFacFile.length() > 0) {
					LOG4CXX_INFO(logger, "Writing row factors to " << outputRowFacFile);
					//DenseMatrix w0;
					//unblock(dw, w0);
					writeMatrix(outputRowFacFile, w);
				}
				if (outputColFacFile.length() > 0) {
					LOG4CXX_INFO(logger, "Writing column factors to " << outputColFacFile);
					//DenseMatrixCM h0;
					//unblock(dh, h0);
					writeMatrix(outputColFacFile, h);
				}
		/**/
	}

	mfStop();
	mfFinalize();

	return 0;
}


/* buckets.cc
   Jeremy Barnes, 12 September 2011
   Copyright (c) 2011 Jeremy Barnes.  All rights reserved.
   This file is part of MLDB. Copyright 2015 mldb.ai inc. All rights reserved.

   Refactored out buckets code.
*/

#include "mldb/utils/buckets.h"
#include <iostream>
#include <boost/utility.hpp>
#include "mldb/arch/format.h"
#include "mldb/utils/floating_point.h"
#include "mldb/utils/pair_utils.h"
#include "mldb/arch/exception.h"

using namespace std;

namespace MLDB {

void
bucket_dist_full(vector<float> & result, const BucketFreqs & freqs)
{
    bool debug = false;

    result.clear();
    result.reserve(std::max<unsigned>(freqs.size() - 1, 0));

    /* We want the split points, which are half way in between the
       current value and the next one. */
    for (int i = 0;  i < (int)freqs.size() - 1;  ++i)
        result.push_back(((freqs.begin() + i)->first
                          + (freqs.begin() + i + 1)->first) * 0.5);

    if (debug) {
        for (BucketFreqs::const_iterator it = freqs.begin();  it != freqs.end();
             ++it)
            cerr << "freqs[" << it->first << "] = " << it->second << endl;

        for (unsigned i = 0;  i < result.size();  ++i)
            cerr << "  bucket " << i << " has value " << result[i]
                 << ", \t" << (freqs.begin() + i)->second << " values" << endl;
    }
}

void
bucket_dist_reduced(vector<float> & result,
                    const BucketFreqs & freqs,
                    size_t num_buckets)
{
    bool debug = false;
    //debug = (feature.type() == 22);

    /* NOTE: this algorithm tends to overcluster well-distributed data.
       It should probably be re-done to produce a number of buckets
       closer to the required number.
    */

    /* Find good split points. */
    double total = freqs.total();

    vector<float> bucket_sizes;
    result.clear();
    result.reserve(num_buckets);

    double num = 0;
    float per_bucket = total / num_buckets;

    /* Go through and look for the really big buckets.  We take these
       out of the bucket size calculation as otherwise we end up with
       too few buckets.  Note that this calculation is not exact
       (since the bucket size becomes smaller after), but it should
       fix up the most problematic cases.
    */
    int num_big_buckets = 0;
    double num_in_big_buckets = 0.0;
    for (BucketFreqs::const_iterator it = freqs.begin();
         it != freqs.end();  ++it) {
        float n = it->second;
        if (n >= per_bucket) {
            num_in_big_buckets += n;
            num_big_buckets += 1;
        }
    }

    if (num_buckets != num_big_buckets)
        per_bucket = (total - num_in_big_buckets) /
            (num_buckets - num_big_buckets);

    if (debug) {
        cerr << "  per_bucket before = "
             << total / num_buckets << endl;
        cerr << "  per_bucket = " << per_bucket << endl;
        cerr << "  num_buckets = " << num_buckets << endl;
        cerr << "  num_big_buckets = " << num_big_buckets << endl;
        cerr << "  num_in_big_buckets = " << num_in_big_buckets << endl;
    }

    //cerr << "freqs.size() = " << freqs.size() << endl;
    int i = 0;
    for (BucketFreqs::const_iterator it = freqs.begin();
         it != freqs.end();  ++it, ++i) {
        float n = it->second;
        float v = it->first;
        num += n;

        if (debug && false)
            cerr << "i = " << i << " n = " << n << " v = " << v
                 << " num = " << num << endl;

        /* We split if we have enough in our bucket, or if the next one
           would have enough by itself. */
        if (n >= per_bucket && num > n) {
            /* Extra split, since this one had enough by itself. */
            float val = (v + boost::prior(it)->first) / 2.0;
            result.push_back(val);
            bucket_sizes.push_back(num - n);
            if (debug)
                cerr << "i = " << i << ": split [1] at "
                     << format("val: %16.9f 0x%08x ", val,
                               reinterpret_as_int(val))
                     << format("v: %16.9f 0x%08x ", v,
                               reinterpret_as_int(v))
                     << format("prior: %16.9f 0x%08x ", boost::prior(it)->first,
                               reinterpret_as_int(boost::prior(it)->first))
                     << " with " << num - n
                     << " examples." << endl;
            num -= n;
        }
        if (i < (int)freqs.size() - 1
            && (num >= per_bucket || boost::next(it)->second >= per_bucket)) {
            float val = (v + boost::next(it)->first) * 0.5f;

            // If the two values are one ulp apart, then we need to make sure
            // that it gets rounded up otherwise we will have two buckets
            // with the same split value
            if (val == v)
                val = boost::next(it)->first;

            result.push_back(val);
            bucket_sizes.push_back(num);
            if (debug)
                cerr << "i = " << i << ": split [2] at "
                     << format("val: %16.9f 0x%08x ", val,
                               reinterpret_as_int(val))
                     << format("v: %16.9f 0x%08x ", v,
                               reinterpret_as_int(v))
                     << format("next: %16.9f 0x%08x ", boost::next(it)->first,
                               reinterpret_as_int(boost::next(it)->first))
                     << " with " << num
                     << " examples." << endl;
            num = 0;
        }
    }

    for (unsigned i = 1;  i < result.size();  ++i)
        if (result[i] == result[i - 1])
            throw Exception("two buckets with the same split point");

    if (debug) {
        for (unsigned i = 0;  i < result.size();  ++i)
            cerr << format("  bucket %5d has value %16.9f 0x%08x, %8.1f values",
                           i, result[i], reinterpret_as_int(result[i]),
                           bucket_sizes[i]) << endl;
    }
}

#if 0
const Dataset_Index::Index_Entry::Bucket_Info &
Dataset_Index::Index_Entry::
create_buckets(size_t num_buckets)
{
    check_used();
    Guard guard(lock);

    //cerr << "create_buckets(" << num_buckets << ")" << endl;

}
#endif

void get_freqs(BucketFreqs & result, std::vector<float> values)
{
    std::sort(values.begin(), values.end(), MLDB::safe_less<float>());

    vector<pair<float, float> > freqs2;
    freqs2.reserve(values.size());

    float last = -INFINITY;
    int count = 0;

    for (unsigned i = 0;  i < values.size();  ++i) {
        if (isnan(values[i]))
            throw Exception("NaN in values");

        if (i == 0) {
            last = values[i];
            count += 1;
        }
        else if (bit_equal(last, values[i]))
            count += 1;
        else {
            freqs2.push_back(make_pair(last, count));
            count = 1;
            last = values[i];
        }
    }
    freqs2.push_back(make_pair(last, count));

    std::sort(freqs2.begin(), freqs2.end());

    result = BucketFreqs(freqs2.begin(), freqs2.end());
} // namespace MLDB

void bucket_dist(std::vector<float> & result,
                 const BucketFreqs & freqs,
                 size_t num_buckets)
{
    if (num_buckets >= freqs.size())
        bucket_dist_full(result, freqs);
    else
        bucket_dist_reduced(result, freqs, num_buckets);
}

Bucket_Info create_buckets(const std::vector<float> & values,
                           size_t num_buckets)
{
    BucketFreqs freqs;
    get_freqs(freqs, values);

    Bucket_Info result;
    bucket_dist(result.splits, freqs, num_buckets);

    vector<int> bucket_count(result.splits.size() + 1);

    for (unsigned i = 0;  i < values.size();  ++i) {
        float value = values[i];
        int bucket = std::upper_bound(result.splits.begin(),
                                      result.splits.end(), value)
            - result.splits.begin();

        bucket_count[bucket] += 1;

        result.buckets.push_back(bucket);
    }

    return result;
}

} // namespace MLDB


/** em.cc                                                          -*- C++ -*-
    Mathieu Marquis Bolduc, October 28th, 2015
    Copyright (c) 2015 mldb.ai inc.  All rights reserved.

    This file is part of MLDB. Copyright 2015 mldb.ai inc. All rights reserved.

    Guassian clustering procedure and functions.
*/

#include "em_interface.h"
#include "mldb/plugins/jml/em.h"
#include "mldb/builtin/matrix.h"
#include "mldb/core/mldb_engine.h"
#include "mldb/engine/procedure_collection.h"
#include "mldb/engine/function_collection.h"
#include "mldb/utils/distribution.h"
#include <boost/multi_array.hpp>
#include "base/parallel.h"
#include "mldb/utils/pair_utils.h"
#include "mldb/arch/timers.h"
#include "mldb/types/optional_description.h"
#include "mldb/utils/vector_utils.h"
#include "mldb/types/basic_value_descriptions.h"
#include "mldb/sql/sql_expression.h"
#include "mldb/core/analytics.h"
#include "mldb/types/any_impl.h"
#include "mldb/utils/smart_ptr_utils.h"
#include "mldb/vfs/fs_utils.h"
#include "mldb/builtin/sql_config_validator.h"
#include "mldb/utils/log.h"

using namespace std;

namespace MLDB {

std::vector<double> tovector(boost::multi_array<double, 2>& m)
{
    std::vector<double> embedding;
    for(int i = 0; i < m.shape()[0]; i++)
    {
        for(int j = 0; j < m.shape()[1]; j++) {
             embedding.push_back(m[i][j]); // multiply by elements on diagonal
        }
    }

    return embedding;
}

DEFINE_STRUCTURE_DESCRIPTION(EMConfig);

EMConfigDescription::
EMConfigDescription()
{
    Optional<PolyConfigT<Dataset> > optional;
    optional.emplace(PolyConfigT<Dataset>().
                     withType(EMConfig::defaultOutputDatasetType));

    addField("trainingData", &EMConfig::trainingData,
             "Specification of the data for input to the procedure.  This should be "
             "organized as an embedding, with each selected row containing the same "
             "set of columns with numeric values to be used as coordinates.  The select statement "
             "does not support groupby and having clauses.");
    addField("outputDataset", &EMConfig::output,
             "Dataset for cluster assignment.  This dataset will contain the same "
             "row names as the input dataset, but the coordinates will be replaced "
             "by a single column giving the cluster number that the row was assigned to.",
              optional);
    addField("centroidsDataset", &EMConfig::centroids,
             "Dataset in which the centroids will be recorded.  This dataset will "
             "have the same coordinates (columns) as those selected from the input "
             "dataset, but will have one row per cluster, providing the centroid of "
             "the cluster.",
             PolyConfigT<Dataset>().withType("embedding"));
    addField("modelFileUrl", &EMConfig::modelFileUrl,
             "URL where the model file (with extension '.gs') should be saved. "
             "This file can be loaded by a function of type 'gaussianclustering' to apply "
             "the trained model to new data. "
             "If someone is only interested in how the training input is clustered "
             "then the parameter can be omitted and the outputDataset param can "
             "be provided instead.");
    addField("numInputDimensions", &EMConfig::numInputDimensions,
             "Number of dimensions from the input to use (-1 = all).  This limits "
             "the number of columns used.  Columns will be ordered alphabetically "
             "and the lowest ones kept.",
             -1);
    addField("numClusters", &EMConfig::numClusters,
             "Number of clusters to create.  This will provide the total number of "
             "centroids created.  There must be at least as many rows selected as "
             "clusters.", 10);
    addField("maxIterations", &EMConfig::maxIterations,
             "Maximum number of iterations to perform.  If no convergance is "
             "reached within this number of iterations, the current clustering "
             "will be returned.", 100);
    addField("functionName", &EMConfig::functionName,
             "If specified, a function of this name will be created using "
             "the training result.");
    addParent<ProcedureConfig>();

    onPostValidate = validateFunction<EMConfig>();
}

/*****************************************************************************/
/* EM PROCEDURE                                                           */
/*****************************************************************************/

EMProcedure::
EMProcedure(MldbEngine * owner,
               PolyConfig config,
               const std::function<bool (const Json::Value &)> & onProgress)
    : Procedure(owner)
{
    this->emConfig = config.params.convert<EMConfig>();
}

Any
EMProcedure::
getStatus() const
{
    return Any();

}

RunOutput
EMProcedure::
run(const ProcedureRunConfig & run,
    const std::function<bool (const Json::Value &)> & onProgress) const
{
    auto runProcConf = applyRunConfOverProcConf(emConfig, run);

    auto onProgress2 = [&] (const Json::Value & progress)
        {
            Json::Value value;
            value["dataset"] = progress;
            return onProgress(value);
        };

    if (!runProcConf.modelFileUrl.empty()) {
        checkWritability(runProcConf.modelFileUrl.toDecodedString(),
                         "modelFileUrl");
    }

    SqlExpressionMldbScope context(engine);

    ConvertProgressToJson convertProgressToJson(onProgress);
    auto embeddingOutput = getEmbedding(*runProcConf.trainingData.stm,
                                        context,
                                        runProcConf.numInputDimensions,
                                        convertProgressToJson);

    auto rows = embeddingOutput.first;
    std::vector<KnownColumn> & vars = embeddingOutput.second;

    std::vector<distribution<double> > vecs;

    for (unsigned i = 0;  i < rows.size();  ++i) {
        vecs.emplace_back(distribution<double>(std::get<2>(rows[i]).begin(),
                                                   std::get<2>(rows[i]).end()));
    }

    if (vecs.size() == 0)
        throw AnnotatedException(400, "Gaussian clustering training requires at least 1 datapoint. "
                                  "Make sure your dataset is not empty and that your WHERE expression "
                                  "does not filter all the rows");

    ML::EstimationMaximisation em;
    vector<int> inCluster;

    int numClusters = emConfig.numClusters;
    int numIterations = emConfig.maxIterations;

    DEBUG_MSG(logger) << "EM training start";
    em.train(vecs, inCluster, numClusters, numIterations, 0);
    DEBUG_MSG(logger) << "EM training end";

    // Let the model know about its column names
    std::vector<ColumnPath> columnNames;
    for (auto & v: vars) {
        columnNames.push_back(v.columnName);
        em.columnNames.push_back(v.columnName.toUtf8String());
    }

    // output

    bool saved = false;
    if (!runProcConf.modelFileUrl.empty()) {
        try {
            makeUriDirectory(
                runProcConf.modelFileUrl.toDecodedString());
            em.save(runProcConf.modelFileUrl.toDecodedString());
            saved = true;
        }
        catch (const std::exception & exc) {
            throw AnnotatedException(400, "Error saving gaussian clustering model at location'" +
                                      runProcConf.modelFileUrl.toString() + "': " +
                                      exc.what());
        }
    }

    if (runProcConf.output.get()) {

        PolyConfigT<Dataset> outputDataset = *runProcConf.output;
        if (outputDataset.type.empty())
            outputDataset.type = EMConfig::defaultOutputDatasetType;

        auto output = createDataset(engine, outputDataset, onProgress2, true /*overwrite*/);

        Date applyDate = Date::now();

        for (unsigned i = 0;  i < rows.size();  ++i) {
            std::vector<std::tuple<ColumnPath, CellValue, Date> > cols;
            cols.emplace_back(ColumnPath("cluster"), inCluster[i], applyDate);
            output->recordRow(std::get<1>(rows[i]), cols);
        }

        output->commit();
    }

    if (runProcConf.centroids.type != "" || emConfig.centroids.id != "") {

        auto centroids = createDataset(engine, runProcConf.centroids, onProgress2, true /*overwrite*/);

        Date applyDate = Date::now();

        for (unsigned i = 0;  i < em.clusters.size();  ++i) {
            auto & cluster = em.clusters[i];

            std::vector<std::tuple<ColumnPath, CellValue, Date> > cols;

            for (unsigned j = 0;  j < cluster.centroid.size();  ++j) {
                cols.emplace_back(columnNames[j], cluster.centroid[j], applyDate);
            }

            auto flatmatrix = tovector(cluster.covarianceMatrix);

            for (unsigned j = 0;  j < flatmatrix.size();  ++j) {
                cols.emplace_back(ColumnPath(MLDB::format("c%02d", j)), flatmatrix[j], applyDate);
            }

            centroids->recordRow(RowPath(MLDB::format("%i", i)), cols);
        }

        centroids->commit();
    }

    if (!runProcConf.functionName.empty()) {
        if (saved) {
            EMFunctionConfig funcConf;
            funcConf.modelFileUrl = runProcConf.modelFileUrl;

            PolyConfig emPC;
            emPC.type = "gaussianclustering";
            emPC.id = runProcConf.functionName;
            emPC.params = funcConf;
            obtainFunction(engine, emPC, onProgress);
        } else {
            throw AnnotatedException(400, "Can't create gaussian clustering function '" +
                                      runProcConf.functionName.rawString() +
                                      "'. Have you provided a valid modelFileUrl?",
                                      "modelFileUrl", runProcConf.modelFileUrl.toString());
        }
    }

    return Any();
}

DEFINE_STRUCTURE_DESCRIPTION(EMFunctionConfig);

EMFunctionConfigDescription::
EMFunctionConfigDescription()
{
    addField("modelFileUrl", &EMFunctionConfig::modelFileUrl,
             "URL of the model file (with extension '.gs') to load. "
             "This file is created by a procedure of type 'gaussianclustering.train'.");

    onPostValidate = [] (EMFunctionConfig * cfg,
                         JsonParsingContext & context) {
        // this includes empty url
        if(!cfg->modelFileUrl.valid()) {
            throw MLDB::Exception("modelFileUrl \"" + cfg->modelFileUrl.toString()
                                + "\" is not valid");
        }
    };
}

/*****************************************************************************/
/* EM FUNCTION                                                               */
/*****************************************************************************/

DEFINE_STRUCTURE_DESCRIPTION(EMInput);

EMInputDescription::EMInputDescription()
{
    addField("embedding", &EMInput::embedding,
             "Values to be assigned to a cluster.");
}

DEFINE_STRUCTURE_DESCRIPTION(EMOutput);

EMOutputDescription::EMOutputDescription()
{
    addField("cluster", &EMOutput::cluster,
             "Cluster corresponding to the input values.");
}

struct EMFunction::Impl {
    ML::EstimationMaximisation em;
    std::vector<ColumnPath> columnNames;

    Impl(const Url & modelFileUrl) {
        em.load(modelFileUrl.toDecodedString());
        for (auto & c: em.columnNames)
            this->columnNames.push_back(PathElement(c));
    }
};

EMFunction::
EMFunction(MldbEngine * owner,
            PolyConfig config,
            const std::function<bool (const Json::Value &)> & onProgress)
    : BaseT(owner, config)
{

    functionConfig = config.params.convert<EMFunctionConfig>();

    impl.reset(new Impl(functionConfig.modelFileUrl));

    dimension = impl->em.clusters[0].centroid.size();

    auto logger = MLDB::getMldbLog<EMFunction>();
    DEBUG_MSG(logger) << "got " << impl->em.clusters.size()
         << " clusters with " << dimension
         << "values";
}

struct EMFunctionApplier: public FunctionApplierT<EMInput, EMOutput> {
    EMFunctionApplier(const EMFunction * owner,
                      const std::shared_ptr<ExpressionValueInfo> & input)
        : FunctionApplierT<EMInput, EMOutput>(owner)
    {
        info = owner->getFunctionInfo();
        extract = input->extractDoubleEmbedding(owner->impl->columnNames);
    }

    ExpressionValueInfo::ExtractDoubleEmbeddingFunction extract;
};

std::unique_ptr<FunctionApplierT<EMInput, EMOutput> >
EMFunction::
bindT(SqlBindingScope & outerContext,
      const std::vector<std::shared_ptr<ExpressionValueInfo> > & input) const
{
    if (input.size() != 1)
        throw AnnotatedException(400, "EM function requires a single input");
    return std::unique_ptr<EMFunctionApplier>
        (new EMFunctionApplier(this, input.at(0)));
}

EMOutput
EMFunction::
applyT(const ApplierT & applier_, EMInput input_) const
{
    // Extract an embedding with the given column names
    ExpressionValue storage;

    const auto * downcast
            = dynamic_cast<const EMFunctionApplier *>(&applier_);

    distribution<double> input = downcast->extract(input_.embedding);
    Date ts = input_.embedding.getEffectiveTimestamp();

    int bestCluster = impl->em.assign(input);

    return {ExpressionValue(bestCluster, ts)};
}

namespace {

RegisterProcedureType<EMProcedure, EMConfig>
regEM(builtinPackage(),
      "Gaussian clustering algorithm using Estimation Maximization on Gaussian Mixture Models",
      "procedures/EMProcedure.md.html",
      nullptr /* static route */,
      { MldbEntity::INTERNAL_ENTITY });

RegisterFunctionType<EMFunction, EMFunctionConfig>
regEMFunction(builtinPackage(), "gaussianclustering",
              "Apply an gaussian clustering to new data",
              "functions/EM.md.html",
              nullptr /* static route */,
              { MldbEntity::INTERNAL_ENTITY });

} // file scope

} // namespace MLDB


// ====================================================================
// This file is part of FlexibleSUSY.
//
// FlexibleSUSY 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.
//
// FlexibleSUSY 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 FlexibleSUSY.  If not, see
// <http://www.gnu.org/licenses/>.
// ====================================================================

#include "slha_io.hpp"
#include "wrappers.hpp"
#include "lowe.h"
#include "linalg.h"
#include "ew_input.hpp"
#include "physical_input.hpp"
#include "spectrum_generator_settings.hpp"

#include <fstream>
#include <algorithm>
#include <string>
#include <boost/bind.hpp>

namespace flexiblesusy {

SLHA_io::SLHA_io()
   : data()
   , modsel()
{
}

void SLHA_io::clear()
{
   data.clear();
   modsel.clear();
}

void SLHA_io::convert_symmetric_fermion_mixings_to_slha(double&,
                                                        Eigen::Matrix<double, 1, 1>&)
{
}

/**
 * @param m mass
 * @param z 1x1 mixing matrix
 */
void SLHA_io::convert_symmetric_fermion_mixings_to_slha(double& m,
                                                        Eigen::Matrix<std::complex<double>, 1, 1>& z)
{
   // check if 1st row contains non-zero imaginary parts
   if (!is_zero(Abs(Im(z(0,0))))) {
      z(0,0) *= std::complex<double>(0.0,1.0);
      m *= -1;
#ifdef ENABLE_DEBUG
      if (!is_zero(Abs(Im(z(0,0))))) {
         WARNING("Element (0,0) of the following fermion mixing matrix"
                 " contains entries which have non-zero real and imaginary"
                 " parts:\nZ = " << z);
      }
#endif
   }
}

void SLHA_io::convert_symmetric_fermion_mixings_to_hk(double&,
                                                      Eigen::Matrix<double, 1, 1>&)
{
}

/**
 * @param m mass
 * @param z 1x1 mixing matrix
 */
void SLHA_io::convert_symmetric_fermion_mixings_to_hk(double& m,
                                                      Eigen::Matrix<std::complex<double>, 1, 1>& z)
{
   if (m < 0.) {
      z(0,0) *= std::complex<double>(0.0,1.0);
      m *= -1;
   }
}

bool SLHA_io::block_exists(const std::string& block_name) const
{
   return data.find(block_name) != data.cend();
}

std::string SLHA_io::to_lower(const std::string& str)
{
   std::string lower(str.size(), ' ');
   std::transform(str.begin(), str.end(), lower.begin(), ::tolower);
   return lower;
}

/**
 * @brief reads from source
 *
 * If source is "-", then read_from_stream() is called.  Otherwise,
 * read_from_file() is called.
 *
 * @param source string that specifies the source
 */
void SLHA_io::read_from_source(const std::string& source)
{
   if (source == "-")
      read_from_stream(std::cin);
   else
      read_from_file(source);
}

/**
 * @brief opens SLHA input file and reads the content
 * @param file_name SLHA input file name
 */
void SLHA_io::read_from_file(const std::string& file_name)
{
   std::ifstream ifs(file_name);
   if (ifs.good()) {
      data.clear();
      data.read(ifs);
   } else {
      std::ostringstream msg;
      msg << "cannot read SLHA file: \"" << file_name << "\"";
      throw ReadError(msg.str());
   }
}

/**
 * @brief reads SLHA data from a stream
 * @param istr input stream
 */
void SLHA_io::read_from_stream(std::istream& istr)
{
   data.read(istr);
}

void SLHA_io::read_modsel()
{
   SLHA_io::Tuple_processor modsel_processor
      = boost::bind(&SLHA_io::process_modsel_tuple, boost::ref(modsel), _1, _2);

   read_block("MODSEL", modsel_processor);
}

void SLHA_io::fill(softsusy::QedQcd& qedqcd) const
{
   CKM_wolfenstein ckm_wolfenstein;
   PMNS_parameters pmns_parameters;

   SLHA_io::Tuple_processor sminputs_processor
      = boost::bind(&SLHA_io::process_sminputs_tuple, boost::ref(qedqcd), _1, _2);

   read_block("SMINPUTS", sminputs_processor);

   if (modsel.quark_flavour_violated) {
      SLHA_io::Tuple_processor vckmin_processor
         = boost::bind(&SLHA_io::process_vckmin_tuple, boost::ref(ckm_wolfenstein), _1, _2);

      read_block("VCKMIN", vckmin_processor);
   }

   if (modsel.lepton_flavour_violated) {
      SLHA_io::Tuple_processor upmnsin_processor
         = boost::bind(&SLHA_io::process_upmnsin_tuple, boost::ref(pmns_parameters), _1, _2);

      read_block("UPMNSIN", upmnsin_processor);
   }

   // fill CKM parameters in qedqcd
   CKM_parameters ckm_parameters;
   ckm_parameters.set_from_wolfenstein(
      ckm_wolfenstein.lambdaW,
      ckm_wolfenstein.aCkm,
      ckm_wolfenstein.rhobar,
      ckm_wolfenstein.etabar);
   qedqcd.setCKM(ckm_parameters);

   // fill PMNS parameters in qedqcd
   qedqcd.setPMNS(pmns_parameters);
}

/**
 * Fill struct of extra physical input parameters from SLHA object
 * (FlexibleSUSYInput block)
 *
 * @param input struct of physical input parameters
 */
void SLHA_io::fill(Physical_input& input) const
{
   SLHA_io::Tuple_processor processor
      = boost::bind(&SLHA_io::process_flexiblesusyinput_tuple, boost::ref(input), _1, _2);

   read_block("FlexibleSUSYInput", processor);
}

/**
 * Fill struct of spectrum generator settings from SLHA object
 * (FlexibleSUSY block)
 *
 * @param settings struct of spectrum generator settings
 */
void SLHA_io::fill(Spectrum_generator_settings& settings) const
{
   SLHA_io::Tuple_processor flexiblesusy_processor
      = boost::bind(&SLHA_io::process_flexiblesusy_tuple, boost::ref(settings), _1, _2);

   read_block("FlexibleSUSY", flexiblesusy_processor);
}

/**
 * Applies processor to each (key, value) pair of a SLHA block.
 * Non-data lines are ignored.
 *
 * @param block_name block name
 * @param processor tuple processor to be applied
 *
 * @return scale (or 0 if no scale is defined)
 */
double SLHA_io::read_block(const std::string& block_name, const Tuple_processor& processor) const
{
   SLHAea::Coll::const_iterator block =
      data.find(data.cbegin(), data.cend(), block_name);

   double scale = 0.;

   while (block != data.cend()) {
      for (SLHAea::Block::const_iterator line = block->cbegin(),
              end = block->cend(); line != end; ++line) {
         if (!line->is_data_line()) {
            // read scale from block definition
            if (line->size() > 3 &&
                to_lower((*line)[0]) == "block" && (*line)[2] == "Q=")
               scale = convert_to<double>((*line)[3]);
            continue;
         }

         if (line->size() >= 2) {
            const int key = convert_to<int>((*line)[0]);
            const double value = convert_to<double>((*line)[1]);
            processor(key, value);
         }
      }

      ++block;
      block = data.find(block, data.cend(), block_name);
   }

   return scale;
}

/**
 * Fills an entry from a SLHA block
 *
 * @param block_name block name
 * @param entry entry to be filled
 *
 * @return scale (or 0 if no scale is defined)
 */
double SLHA_io::read_block(const std::string& block_name, double& entry) const
{
   SLHAea::Coll::const_iterator block =
      data.find(data.cbegin(), data.cend(), block_name);

   double scale = 0.;

   while (block != data.cend()) {
      for (SLHAea::Block::const_iterator line = block->cbegin(),
              end = block->cend(); line != end; ++line) {
         if (!line->is_data_line()) {
            // read scale from block definition
            if (line->size() > 3 &&
                to_lower((*line)[0]) == "block" && (*line)[2] == "Q=")
               scale = convert_to<double>((*line)[3]);
            continue;
         }

         if (line->size() >= 1)
            entry = convert_to<double>((*line)[0]);
      }

      ++block;
      block = data.find(block, data.cend(), block_name);
   }

   return scale;
}

double SLHA_io::read_entry(const std::string& block_name, int key) const
{
   SLHAea::Coll::const_iterator block =
      data.find(data.cbegin(), data.cend(), block_name);

   double entry = 0.;
   const SLHAea::Block::key_type keys(1, ToString(key));
   SLHAea::Block::const_iterator line;

   while (block != data.cend()) {
      line = block->find(keys);

      if (line != block->end() && line->is_data_line() && line->size() > 1)
         entry = convert_to<double>(line->at(1));

      ++block;
      block = data.find(block, data.cend(), block_name);
   }

   return entry;
}

/**
 * Reads scale definition from SLHA block.
 *
 * @param block_name block name
 *
 * @return scale (or 0 if no scale is defined)
 */
double SLHA_io::read_scale(const std::string& block_name) const
{
   if (!block_exists(block_name))
      return 0.;

   double scale = 0.;

   for (SLHAea::Block::const_iterator line = data.at(block_name).cbegin(),
        end = data.at(block_name).cend(); line != end; ++line) {
      if (!line->is_data_line()) {
         if (line->size() > 3 &&
             to_lower((*line)[0]) == "block" && (*line)[2] == "Q=")
            scale = convert_to<double>((*line)[3]);
         break;
      }
   }

   return scale;
}

void SLHA_io::set_block(const std::ostringstream& lines, Position position)
{
   SLHAea::Block block;
   block.str(lines.str());
   data.erase(block.name());
   if (position == front)
      data.push_front(block);
   else
      data.push_back(block);
}

void SLHA_io::set_block(const std::string& lines, Position position)
{
   set_block(std::ostringstream(lines), position);
}

void SLHA_io::set_blocks(const std::vector<std::string>& blocks, Position position)
{
   for (std::vector<std::string>::const_iterator it = blocks.begin(),
           end = blocks.end(); it != end; it++)
      set_block(*it, position);
}

/**
 * This function treats a given scalar as 1x1 matrix.  Such a case is
 * not defined in the SLHA standard, but we still handle it to avoid
 * problems.
 */
void SLHA_io::set_block(const std::string& name, double value,
                        const std::string& symbol, double scale)
{
   std::ostringstream ss;
   ss << "Block " << name;
   if (scale != 0.)
      ss << " Q= " << FORMAT_SCALE(scale);
   ss << '\n'
      << boost::format(mixing_matrix_formatter) % 1 % 1 % value % symbol;

   set_block(ss);
}

void SLHA_io::set_block(const std::string& name, const softsusy::DoubleMatrix& matrix,
                        const std::string& symbol, double scale)
{
   std::ostringstream ss;
   ss << "Block " << name;
   if (scale != 0.)
      ss << " Q= " << FORMAT_SCALE(scale);
   ss << '\n';

   for (int i = 1; i <= matrix.displayRows(); ++i)
      for (int k = 1; k <= matrix.displayCols(); ++k) {
         ss << boost::format(mixing_matrix_formatter) % i % k % matrix(i,k)
            % (symbol + "(" + ToString(i) + "," + ToString(k) + ")");
      }

   set_block(ss);
}

void SLHA_io::set_block(const std::string& name, const softsusy::ComplexMatrix& matrix,
                        const std::string& symbol, double scale)
{
   std::ostringstream ss;
   ss << "Block " << name;
   if (scale != 0.)
      ss << " Q= " << FORMAT_SCALE(scale);
   ss << '\n';

   for (int i = 1; i <= matrix.displayRows(); ++i)
      for (int k = 1; k <= matrix.displayCols(); ++k) {
         ss << boost::format(mixing_matrix_formatter) % i % k
            % Re(matrix(i,k))
            % ("Re(" + symbol + "(" + ToString(i) + "," + ToString(k) + "))");
      }

   set_block(ss);
}

void SLHA_io::set_sminputs(const softsusy::QedQcd& qedqcd_)
{
   using namespace softsusy;
   softsusy::QedQcd qedqcd(qedqcd_);
   std::ostringstream ss;

   const double alphaEmInv = 1./qedqcd.displayAlpha(ALPHA);

   ss << "Block SMINPUTS\n";
   ss << FORMAT_ELEMENT( 1, alphaEmInv                   , "alpha^(-1) SM MSbar(MZ)");
   ss << FORMAT_ELEMENT( 2, qedqcd.displayFermiConstant(), "G_Fermi");
   ss << FORMAT_ELEMENT( 3, qedqcd.displayAlpha(ALPHAS)  , "alpha_s(MZ) SM MSbar");
   ss << FORMAT_ELEMENT( 4, qedqcd.displayPoleMZ()       , "MZ(pole)");
   ss << FORMAT_ELEMENT( 5, qedqcd.displayMbMb()         , "mb(mb) SM MSbar");
   ss << FORMAT_ELEMENT( 6, qedqcd.displayPoleMt()       , "mtop(pole)");
   ss << FORMAT_ELEMENT( 7, qedqcd.displayPoleMtau()     , "mtau(pole)");
   ss << FORMAT_ELEMENT( 8, qedqcd.displayNeutrinoPoleMass(3), "mnu3(pole)");
   ss << FORMAT_ELEMENT( 9, qedqcd.displayPoleMW()       , "MW(pole)");
   ss << FORMAT_ELEMENT(11, qedqcd.displayPoleMel()      , "melectron(pole)");
   ss << FORMAT_ELEMENT(12, qedqcd.displayNeutrinoPoleMass(1), "mnu1(pole)");
   ss << FORMAT_ELEMENT(13, qedqcd.displayPoleMmuon()    , "mmuon(pole)");
   ss << FORMAT_ELEMENT(14, qedqcd.displayNeutrinoPoleMass(2), "mnu2(pole)");

   // recalculate mc(mc)^MS-bar
   double mc = qedqcd.displayMass(mCharm);
   qedqcd.runto(mc);
   mc = qedqcd.displayMass(mCharm);

   // recalculate mu(2 GeV)^MS-bar, md(2 GeV)^MS-bar, ms^MS-bar(2 GeV)
   qedqcd.runto(2.0);
   ss << FORMAT_ELEMENT(21, qedqcd.displayMass(mDown)    , "md");
   ss << FORMAT_ELEMENT(22, qedqcd.displayMass(mUp)      , "mu");
   ss << FORMAT_ELEMENT(23, qedqcd.displayMass(mStrange) , "ms");
   ss << FORMAT_ELEMENT(24, mc                           , "mc");

   set_block(ss);
}

void SLHA_io::write_to_file(const std::string& file_name)
{
   std::ofstream ofs(file_name);
   write_to_stream(ofs);
}

void SLHA_io::write_to_stream(std::ostream& ostr)
{
   if (ostr.good())
      ostr << data;
   else
      ERROR("cannot write SLHA file");
}

/**
 * fill Modsel struct from given key - value pair
 *
 * @param modsel MODSEL data
 * @param key SLHA key in MODSEL
 * @param value value corresponding to key
 */
void SLHA_io::process_modsel_tuple(Modsel& modsel, int key, double value)
{
   switch (key) {
   case 1: // SUSY breaking model (defined in FlexibleSUSY model file)
   case 3: // SUSY model (defined in SARAH model file)
   case 4: // R-parity violation (defined in SARAH model file)
   case 5: // CP-parity violation (defined in SARAH model file)
   case 11:
   case 21:
      WARNING("Key " << key << " in Block MODSEL currently not supported");
      break;
   case 6: // Flavour violation (defined in SARAH model file)
   {
      const int ivalue = Round(value);

      if (ivalue < 0 || ivalue > 3)
         WARNING("Value " << ivalue << " in MODSEL block entry 6 out of range");

      modsel.quark_flavour_violated = ivalue & 0x1;
      modsel.lepton_flavour_violated = ivalue & 0x2;
   }
      break;
   case 12:
      modsel.parameter_output_scale = value;
      break;
   default:
      WARNING("Unrecognized entry in block MODSEL: " << key);
      break;
   }
}

/**
 * fill qedqcd from given key - value pair
 *
 * @param qedqcd low-energy data set
 * @param key SLHA key in SMINPUTS
 * @param value value corresponding to key
 */
void SLHA_io::process_sminputs_tuple(softsusy::QedQcd& qedqcd, int key, double value)
{
   using namespace softsusy;

   switch (key) {
   case 1:
      qedqcd.setAlpha(ALPHA, 1.0 / value);
      qedqcd.setAlphaEmInput(1.0 / value);
      break;
   case 2:
      qedqcd.setFermiConstant(value);
      break;
   case 3:
      qedqcd.setAlpha(ALPHAS, value);
      qedqcd.setAlphaSInput(value);
      break;
   case 4:
      qedqcd.setPoleMZ(value);
      qedqcd.setMu(value);
      softsusy::MZ = value;
      break;
   case 5:
      qedqcd.setMass(mBottom, value);
      qedqcd.setMbMb(value);
      break;
   case 6:
      qedqcd.setPoleMt(value);
      break;
   case 7:
      qedqcd.setMass(mTau, value);
      qedqcd.setPoleMtau(value);
      break;
   case 8:
      qedqcd.setNeutrinoPoleMass(3, value);
      break;
   case 9:
      qedqcd.setPoleMW(value);
      break;
   case 11:
      qedqcd.setMass(mElectron, value);
      qedqcd.setPoleMel(value);
      break;
   case 12:
      qedqcd.setNeutrinoPoleMass(1, value);
      break;
   case 13:
      qedqcd.setMass(mMuon, value);
      qedqcd.setPoleMmuon(value);
      break;
   case 14:
      qedqcd.setNeutrinoPoleMass(2, value);
      break;
   case 21:
      qedqcd.setMass(mDown, value);
      qedqcd.setMd2GeV(value);
      break;
   case 22:
      qedqcd.setMass(mUp, value);
      qedqcd.setMu2GeV(value);
      break;
   case 23:
      qedqcd.setMass(mStrange, value);
      qedqcd.setMs2GeV(value);
      break;
   case 24:
      qedqcd.setMass(mCharm, value);
      qedqcd.setMcMc(value);
      break;
   default:
      WARNING("Unrecognized entry in block SMINPUTS: " << key);
      break;
   }
}

void SLHA_io::process_flexiblesusy_tuple(Spectrum_generator_settings& settings,
                                         int key, double value)
{
   if (0 <= key && key < static_cast<int>(Spectrum_generator_settings::NUMBER_OF_OPTIONS)) {
      settings.set((Spectrum_generator_settings::Settings)key, value);
   } else {
      WARNING("Unrecognized entry in block FlexibleSUSY: " << key);
   }
}

void SLHA_io::process_flexiblesusyinput_tuple(
   Physical_input& input,
   int key, double value)
{
   if (0 <= key && key < static_cast<int>(Physical_input::NUMBER_OF_INPUT_PARAMETERS)) {
      input.set((Physical_input::Input)key, value);
   } else {
      WARNING("Unrecognized entry in block FlexibleSUSYInput: " << key);
   }
}

/**
 * fill CKM_wolfenstein from given key - value pair
 *
 * @param ckm_wolfenstein Wolfenstein parameters
 * @param key SLHA key in SMINPUTS
 * @param value value corresponding to key
 */
void SLHA_io::process_vckmin_tuple(CKM_wolfenstein& ckm_wolfenstein, int key, double value)
{
   switch (key) {
   case 1:
      ckm_wolfenstein.lambdaW = value;
      break;
   case 2:
      ckm_wolfenstein.aCkm = value;
      break;
   case 3:
      ckm_wolfenstein.rhobar = value;
      break;
   case 4:
      ckm_wolfenstein.etabar = value;
      break;
   default:
      WARNING("Unrecognized entry in block VCKMIN: " << key);
      break;
   }
}

/**
 * fill PMNS_parameters from given key - value pair
 *
 * @param pmns_parameters PMNS matrix parameters
 * @param key SLHA key in SMINPUTS
 * @param value value corresponding to key
 */
void SLHA_io::process_upmnsin_tuple(PMNS_parameters& pmns_parameters, int key, double value)
{
   switch (key) {
   case 1:
      pmns_parameters.theta_12 = value;
      break;
   case 2:
      pmns_parameters.theta_23 = value;
      break;
   case 3:
      pmns_parameters.theta_13 = value;
      break;
   case 4:
      pmns_parameters.delta = value;
      break;
   case 5:
      pmns_parameters.alpha_1 = value;
   case 6:
      pmns_parameters.alpha_2 = value;
      break;
   default:
      WARNING("Unrecognized entry in block UPMNSIN: " << key);
      break;
   }
}

} // namespace flexiblesusy


/*
    Copyright 2016 Emanuele Vespa, Imperial College London
    Redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions are met:

    1. Redistributions of source code must retain the above copyright notice,
   this list of conditions and the following disclaimer.

    2. 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.

    3. Neither the name of the copyright holder 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 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.

*/
#ifndef OCTANT_OPS_HPP
#define OCTANT_OPS_HPP
#include "octree_defines.h"
#include "utils/math_utils.h"
#include "utils/morton_utils.hpp"

#include <supereight/shared/commons.h>

#include <Eigen/Dense>
#include <bitset>
#include <iostream>

namespace se {
namespace keyops {

SE_DEVICE_FUNC
inline se::key_t code(const se::key_t key) { return key & ~SCALE_MASK; }

SE_DEVICE_FUNC
inline int level(const se::key_t key) { return key & SCALE_MASK; }

SE_DEVICE_ONLY_FUNC
inline se::key_t encode(const int x, const int y, const int z, const int level,
    const int max_depth) {
    const int offset = MAX_BITS - max_depth + level - 1;
    return (compute_morton(x, y, z) & MASK[offset] & ~SCALE_MASK) | level;
}

SE_DEVICE_FUNC
inline Eigen::Vector3i decode(const se::key_t key) {
    return unpack_morton(key & ~SCALE_MASK);
}
} // namespace keyops

/*
 * Algorithm 5 of p4est paper: https://epubs.siam.org/doi/abs/10.1137/100791634
 */
SE_DEVICE_FUNC
inline Eigen::Vector3i face_neighbour(const se::key_t o,
    const unsigned int face, const unsigned int l,
    const unsigned int max_depth) {
    Eigen::Vector3i coords  = se::keyops::decode(o);
    const unsigned int side = 1 << (max_depth - l);
    coords(0) = coords(0) + ((face == 0) ? -side : (face == 1) ? side : 0);
    coords(1) = coords(1) + ((face == 2) ? -side : (face == 3) ? side : 0);
    coords(2) = coords(2) + ((face == 4) ? -side : (face == 5) ? side : 0);
    return {coords(0), coords(1), coords(2)};
}

/*
 * \brief Return true if octant is a descendant of ancestor
 * \param octant
 * \param ancestor
 * \param max_depth max depth of the tree on which the octant lives
 */
SE_DEVICE_ONLY_FUNC
inline bool descendant(
    se::key_t octant, se::key_t ancestor, const int max_depth) {
    const int level = se::keyops::level(ancestor);
    const int idx   = MAX_BITS - max_depth + level - 1;
    ancestor        = se::keyops::code(ancestor);
    octant          = se::keyops::code(octant) & MASK[idx];
    return (ancestor ^ octant) == 0;
}

/*
 * \brief Computes the parent's morton code of a given octant
 * \param octant
 * \param max_depth max depth of the tree on which the octant lives
 */
SE_DEVICE_ONLY_FUNC
inline se::key_t parent(const se::key_t& octant, const int max_depth) {
    const int level = se::keyops::level(octant) - 1;
    const int idx   = MAX_BITS - max_depth + level - 1;
    return (octant & MASK[idx]) | level;
}

/*
 * \brief Computes the octants's id in its local brotherhood
 * \param octant
 * \param level of octant
 * \param max_depth max depth of the tree on which the octant lives
 */
SE_DEVICE_FUNC
inline int child_id(se::key_t octant, const int level, const int max_depth) {
    int shift = max_depth - level;
    octant    = se::keyops::code(octant) >> shift * 3;
    int idx   = (octant & 0x01) | (octant & 0x02) | (octant & 0x04);
    return idx;
}

/*
 * \brief Computes the octants's corner which is not shared with its siblings
 * \param octant
 * \param level of octant
 * \param max_depth max depth of the tree on which the octant lives
 */
SE_DEVICE_FUNC
inline Eigen::Vector3i far_corner(
    const se::key_t octant, const int level, const int max_depth) {
    const unsigned int side           = 1 << (max_depth - level);
    const int idx                     = child_id(octant, level, max_depth);
    const Eigen::Vector3i coordinates = se::keyops::decode(octant);
    return Eigen::Vector3i(coordinates(0) + (idx & 1) * side,
        coordinates(1) + ((idx & 2) >> 1) * side,
        coordinates(2) + ((idx & 4) >> 2) * side);
}

/*
 * \brief Computes the non-sibling neighbourhood around an octants. In the
 * special case in which the octant lies on an edge, neighbour are duplicated
 * as movement outside the enclosing cube is forbidden.
 * \param result 7-vector containing the neighbours
 * \param octant
 * \param level of octant
 * \param max_depth max depth of the tree on which the octant lives
 */
SE_DEVICE_ONLY_FUNC
inline void exterior_neighbours(se::key_t result[7], const se::key_t octant,
    const int level, const int max_depth) {
    const int idx       = child_id(octant, level, max_depth);
    Eigen::Vector3i dir = Eigen::Vector3i(
        (idx & 1) ? 1 : -1, (idx & 2) ? 1 : -1, (idx & 4) ? 1 : -1);
    Eigen::Vector3i base = far_corner(octant, level, max_depth);
    dir(0) =
        se::math::in(base(0) + dir(0), 0, (1 << max_depth) - 1) ? dir(0) : 0;
    dir(1) =
        se::math::in(base(1) + dir(1), 0, (1 << max_depth) - 1) ? dir(1) : 0;
    dir(2) =
        se::math::in(base(2) + dir(2), 0, (1 << max_depth) - 1) ? dir(2) : 0;

    result[0] = se::keyops::encode(
        base(0) + dir(0), base(1) + 0, base(2) + 0, level, max_depth);
    result[1] = se::keyops::encode(
        base(0) + 0, base(1) + dir(1), base(2) + 0, level, max_depth);
    result[2] = se::keyops::encode(
        base(0) + dir(0), base(1) + dir(1), base(2) + 0, level, max_depth);
    result[3] = se::keyops::encode(
        base(0) + 0, base(1) + 0, base(2) + dir(2), level, max_depth);
    result[4] = se::keyops::encode(
        base(0) + dir(0), base(1) + 0, base(2) + dir(2), level, max_depth);
    result[5] = se::keyops::encode(
        base(0) + 0, base(1) + dir(1), base(2) + dir(2), level, max_depth);
    result[6] = se::keyops::encode(
        base(0) + dir(0), base(1) + dir(1), base(2) + dir(2), level, max_depth);
}

/*
 * \brief Computes the six face neighbours of an octant. These are stored in an
 * 4x6 matrix in which each column represents the homogeneous coordinates of a
 * neighbouring octant. The neighbours along the x axis come first, followed by
 * neighbours along the y axis and finally along the z axis. All coordinates are
 * clamped to be in the range between [0, max_size] where max size is given
 * by pow(2, max_depth).
 * \param res 4x6 matrix containing the neighbours
 * \param octant octant coordinates
 * \param level level of the octant
 * \param max_depth max depth of the tree on which the octant lives
 */
SE_DEVICE_FUNC
static inline void one_neighbourhood(Eigen::Ref<Eigen::Matrix<int, 4, 6>> res,
    const Eigen::Vector3i& octant, const int level, const int max_depth) {
    const Eigen::Vector3i base = octant;
    const int size             = 1 << max_depth;
    const int step             = 1 << (max_depth - level);
    Eigen::Matrix<int, 4, 6> cross;
    res << -step, step, 0, 0, 0, 0, 0, 0, -step, step, 0, 0, 0, 0, 0, 0, -step,
        step, 0, 0, 0, 0, 0, 0;
    res.colwise() += base.homogeneous();
    res = res.unaryExpr(
        [size](const int a) { return std::max(std::min(a, size - 1), 0); });
}

/*
 * \brief Computes the six face neighbours of an octant. These are stored in an
 * 4x6 matrix in which each column represents the homogeneous coordinates of a
 * neighbouring octant. The neighbours along the x axis come first, followed by
 * neighbours along the y axis and finally along the z axis. All coordinates are
 * clamped to be in the range between [0, max_size] where max size is given
 * by pow(2, max_depth).
 * \param res 4x6 matrix containing the neighbours
 * \param octant octant key
 * \param max_depth max depth of the tree on which the octant lives
 */
SE_DEVICE_FUNC
static inline void one_neighbourhood(Eigen::Ref<Eigen::Matrix<int, 4, 6>> res,
    const se::key_t octant, const int max_depth) {
    one_neighbourhood(
        res, se::keyops::decode(octant), se::keyops::level(octant), max_depth);
}

/*
 * \brief Computes the morton number of all siblings around an octant,
 * including itself.
 * \param result 8-vector containing the neighbours
 * \param octant
 * \param max_depth max depth of the tree on which the octant lives
 */
SE_DEVICE_ONLY_FUNC
inline void siblings(
    se::key_t result[8], const se::key_t octant, const int max_depth) {
    const int level   = (octant & SCALE_MASK);
    const int shift   = 3 * (max_depth - level);
    const se::key_t p = parent(octant, max_depth) + 1; // set-up next level
    for (int i = 0; i < 8; ++i) { result[i] = p | (i << shift); }
}
} // namespace se
#endif