/**
 * @file    GaussianFactorGraph.cpp
 * @brief   Linear Factor Graph where all factors are Gaussians
 * @author  Kai Ni
 * @author  Christian Potthast
 */

#include <boost/foreach.hpp>
#include <boost/tuple/tuple.hpp>
#include <boost/numeric/ublas/lu.hpp>
#include <boost/numeric/ublas/io.hpp>

#include <gtsam/linear/GaussianFactorGraph.h>
#include <gtsam/linear/GaussianFactorSet.h>
#include <gtsam/inference/FactorGraph-inl.h>
#include <gtsam/inference/inference-inl.h>
#include <gtsam/linear/iterative.h>
//#include <gtsam/linear/GaussianJunctionTree.h>

using namespace std;
using namespace gtsam;
using namespace boost::assign;

// trick from some reading group
#define FOREACH_PAIR( KEY, VAL, COL) BOOST_FOREACH (boost::tie(KEY,VAL),COL)

namespace gtsam {

// Explicitly instantiate so we don't have to include everywhere
template class FactorGraph<GaussianFactor>;

/* ************************************************************************* */
GaussianFactorGraph::GaussianFactorGraph(const GaussianBayesNet& CBN) :
	FactorGraph<GaussianFactor> (CBN) {
}

/* ************************************************************************* */
std::set<varid_t, std::less<varid_t>, boost::fast_pool_allocator<varid_t> >
GaussianFactorGraph::keys() const {
  std::set<varid_t, std::less<varid_t>, boost::fast_pool_allocator<varid_t> > keys;
  BOOST_FOREACH(const sharedFactor& factor, *this) {
    if(factor) keys.insert(factor->begin(), factor->end()); }
  return keys;
}

/* ************************************************************************* */
void GaussianFactorGraph::permuteWithInverse(const Permutation& inversePermutation) {
  BOOST_FOREACH(const sharedFactor& factor, factors_) {
    factor->permuteWithInverse(inversePermutation);
  }
}

/* ************************************************************************* */
double GaussianFactorGraph::error(const VectorValues& x) const {
	double total_error = 0.;
	BOOST_FOREACH(sharedFactor factor,factors_)
		total_error += factor->error(x);
	return total_error;
}

/* ************************************************************************* */
Errors GaussianFactorGraph::errors(const VectorValues& x) const {
	return *errors_(x);
}

/* ************************************************************************* */
boost::shared_ptr<Errors> GaussianFactorGraph::errors_(const VectorValues& x) const {
	boost::shared_ptr<Errors> e(new Errors);
	BOOST_FOREACH(const sharedFactor& factor,factors_)
		e->push_back(factor->error_vector(x));
	return e;
}

/* ************************************************************************* */
Errors GaussianFactorGraph::operator*(const VectorValues& x) const {
	Errors e;
	BOOST_FOREACH(const sharedFactor& Ai,factors_)
		e.push_back((*Ai)*x);
	return e;
}

/* ************************************************************************* */
void GaussianFactorGraph::multiplyInPlace(const VectorValues& x, Errors& e) const {
	multiplyInPlace(x,e.begin());
}

/* ************************************************************************* */
void GaussianFactorGraph::multiplyInPlace(const VectorValues& x,
		const Errors::iterator& e) const {
	Errors::iterator ei = e;
	BOOST_FOREACH(const sharedFactor& Ai,factors_) {
		*ei = (*Ai)*x;
		ei++;
	}
}

///* ************************************************************************* */
//VectorValues GaussianFactorGraph::operator^(const Errors& e) const {
//	VectorValues x;
//	// For each factor add the gradient contribution
//	Errors::const_iterator it = e.begin();
//	BOOST_FOREACH(const sharedFactor& Ai,factors_) {
//		VectorValues xi = (*Ai)^(*(it++));
//		x.insertAdd(xi);
//	}
//	return x;
//}

/* ************************************************************************* */
// x += alpha*A'*e
void GaussianFactorGraph::transposeMultiplyAdd(double alpha, const Errors& e,
		VectorValues& x) const {
	// For each factor add the gradient contribution
	Errors::const_iterator ei = e.begin();
	BOOST_FOREACH(const sharedFactor& Ai,factors_)
		Ai->transposeMultiplyAdd(alpha,*(ei++),x);
}

///* ************************************************************************* */
//VectorValues GaussianFactorGraph::gradient(const VectorValues& x) const {
//	// It is crucial for performance to make a zero-valued clone of x
//	VectorValues g = VectorValues::zero(x);
//	transposeMultiplyAdd(1.0, errors(x), g);
//	return g;
//}

///* ************************************************************************* */
//set<varid_t> GaussianFactorGraph::find_separator(varid_t key) const
//{
//	set<varid_t> separator;
//	BOOST_FOREACH(const sharedFactor& factor,factors_)
//		factor->tally_separator(key,separator);
//
//	return separator;
//}

///* ************************************************************************* */
//template <class Factors>
//std::pair<Matrix, SharedDiagonal> combineFactorsAndCreateMatrix(
//		const Factors& factors,
//		const Ordering& order, const Dimensions& dimensions) {
//	// find the size of Ab
//	size_t m = 0, n = 1;
//
//	// number of rows
//	BOOST_FOREACH(GaussianFactor::shared_ptr factor, factors) {
//		m += factor->numberOfRows();
//	}
//
//	// find the number of columns
//	BOOST_FOREACH(varid_t key, order) {
//		n += dimensions.at(key);
//	}
//
//	// Allocate the new matrix
//	Matrix Ab = zeros(m,n);
//
//	// Allocate a sigmas vector to make into a full noisemodel
//	Vector sigmas = ones(m);
//
//	// copy data over
//	size_t cur_m = 0;
//	bool constrained = false;
//	bool unit = true;
//	BOOST_FOREACH(GaussianFactor::shared_ptr factor, factors) {
//		// loop through ordering
//		size_t cur_n = 0;
//		BOOST_FOREACH(varid_t key, order) {
//			// copy over matrix if it exists
//			if (factor->involves(key)) {
//				insertSub(Ab, factor->get_A(key), cur_m, cur_n);
//			}
//			// move onto next element
//			cur_n += dimensions.at(key);
//		}
//		// copy over the RHS
//		insertColumn(Ab, factor->get_b(), cur_m, n-1);
//
//		// check if the model is unit already
//		if (!boost::shared_dynamic_cast<noiseModel::Unit>(factor->get_model())) {
//			unit = false;
//			const Vector& subsigmas = factor->get_model()->sigmas();
//			subInsert(sigmas, subsigmas, cur_m);
//
//			// check for constraint
//			if (boost::shared_dynamic_cast<noiseModel::Constrained>(factor->get_model()))
//				constrained = true;
//		}
//
//		// move to next row
//		cur_m += factor->numberOfRows();
//	}
//
//	// combine the noisemodels
//	SharedDiagonal model;
//	if (unit) {
//		model = noiseModel::Unit::Create(m);
//	} else if (constrained) {
//		model = noiseModel::Constrained::MixedSigmas(sigmas);
//	} else {
//		model = noiseModel::Diagonal::Sigmas(sigmas);
//	}
//	return make_pair(Ab, model);
//}

///* ************************************************************************* */
//GaussianConditional::shared_ptr
//GaussianFactorGraph::eliminateOneMatrixJoin(varid_t key) {
//	// find and remove all factors connected to key
//	vector<GaussianFactor::shared_ptr> factors = findAndRemoveFactors(key);
//
//	// Collect all dimensions as well as the set of separator keys
//	set<varid_t> separator;
//	Dimensions dimensions;
//	BOOST_FOREACH(const sharedFactor& factor, factors) {
//		Dimensions factor_dim = factor->dimensions();
//		dimensions.insert(factor_dim.begin(), factor_dim.end());
//		BOOST_FOREACH(varid_t k, factor->keys())
//			if (!k == key)
//				separator.insert(k);
//	}
//
//	// add the keys to the rendering
//	Ordering render; render += key;
//	BOOST_FOREACH(varid_t k, separator)
//			if (k != key) render += k;
//
//	// combine the factors to get a noisemodel and a combined matrix
//	Matrix Ab; SharedDiagonal model;
//
//	boost::tie(Ab, model) =	combineFactorsAndCreateMatrix(factors,render,dimensions);
//
//	// eliminate that joint factor
//	GaussianFactor::shared_ptr factor;
//	GaussianConditional::shared_ptr conditional;
//	render.pop_front();
//	boost::tie(conditional, factor) =
//			GaussianFactor::eliminateMatrix(Ab, model, key, render, dimensions);
//
//	// add new factor on separator back into the graph
//	if (!factor->empty()) push_back(factor);
//
//	// return the conditional Gaussian
//	return conditional;
//}
//
///* ************************************************************************* */
//GaussianBayesNet
//GaussianFactorGraph::eliminateFrontals(const Ordering& frontals)
//{
//	// find the factors that contain at least one of the frontal variables
//	Dimensions dimensions = this->dimensions();
//
//	// collect separator
//	Ordering separator;
//	set<varid_t> frontal_set(frontals.begin(), frontals.end());
//	BOOST_FOREACH(varid_t key, this->keys()) {
//		if (frontal_set.find(key) == frontal_set.end())
//			separator.push_back(key);
//	}
//
//	Matrix Ab; SharedDiagonal model;
//	Ordering ord = frontals;
//	ord.insert(ord.end(), separator.begin(), separator.end());
//	boost::tie(Ab, model) = combineFactorsAndCreateMatrix(*this, ord, dimensions);
//
//	// eliminate that joint factor
//	GaussianFactor::shared_ptr factor;
//	GaussianBayesNet bn;
//	boost::tie(bn, factor) =
//			GaussianFactor::eliminateMatrix(Ab, model, frontals, separator, dimensions);
//
//	// add new factor on separator back into the graph
//	*this = GaussianFactorGraph();
//	if (!factor->empty()) push_back(factor);
//
//	return bn;
//}


///* ************************************************************************* */
//VectorValues GaussianFactorGraph::optimize(const Ordering& ordering, bool old)
//{
//	// eliminate all nodes in the given ordering -> chordal Bayes net
//	GaussianBayesNet chordalBayesNet = eliminate(ordering, old);
//
//	// calculate new values structure (using backsubstitution)
//	VectorValues delta = ::optimize(chordalBayesNet);
//	return delta;
//}
//
///* ************************************************************************* */
//VectorValues GaussianFactorGraph::optimizeMultiFrontals(const Ordering& ordering)
//{
//	GaussianJunctionTree junctionTree(*this, ordering);
//
//	// calculate new values structure (using backsubstitution)
//	VectorValues delta = junctionTree.optimize();
//	return delta;
//}

///* ************************************************************************* */
//boost::shared_ptr<GaussianBayesNet>
//GaussianFactorGraph::eliminate_(const Ordering& ordering)
//{
//	boost::shared_ptr<GaussianBayesNet> chordalBayesNet(new GaussianBayesNet); // empty
//	BOOST_FOREACH(varid_t key, ordering) {
//		GaussianConditional::shared_ptr cg = eliminateOne(key);
//		chordalBayesNet->push_back(cg);
//	}
//	return chordalBayesNet;
//}
//
///* ************************************************************************* */
//boost::shared_ptr<VectorValues>
//GaussianFactorGraph::optimize_(const Ordering& ordering) {
//	return boost::shared_ptr<VectorValues>(new VectorValues(optimize(ordering)));
//}

/* ************************************************************************* */
void GaussianFactorGraph::combine(const GaussianFactorGraph &lfg){
	for(const_iterator factor=lfg.factors_.begin(); factor!=lfg.factors_.end(); factor++){
		push_back(*factor);
	}
}

/* ************************************************************************* */
GaussianFactorGraph GaussianFactorGraph::combine2(const GaussianFactorGraph& lfg1,
		const GaussianFactorGraph& lfg2) {

	// create new linear factor graph equal to the first one
	GaussianFactorGraph fg = lfg1;

	// add the second factors_ in the graph
	for (const_iterator factor = lfg2.factors_.begin(); factor
			!= lfg2.factors_.end(); factor++) {
		fg.push_back(*factor);
	}
	return fg;
}

//
///* ************************************************************************* */
//Dimensions GaussianFactorGraph::dimensions() const {
//	Dimensions result;
//	BOOST_FOREACH(const sharedFactor& factor,factors_) {
//		Dimensions vs = factor->dimensions();
//		varid_t key; size_t dim;
//		FOREACH_PAIR(key,dim,vs) result.insert(make_pair(key,dim));
//	}
//	return result;
//}

/* ************************************************************************* */  
GaussianFactorGraph GaussianFactorGraph::add_priors(double sigma, const GaussianVariableIndex<>& variableIndex) const {

	// start with this factor graph
	GaussianFactorGraph result = *this;

	// for each of the variables, add a prior

	for(varid_t var=0; var<variableIndex.size(); ++var) {
	  const GaussianVariableIndex<>::mapped_factor_type& factor_pos(variableIndex[var].front());
	  const GaussianFactor& factor(*operator[](factor_pos.factorIndex));
	  size_t dim = factor.getDim(factor.begin() + factor_pos.variablePosition);
		Matrix A = eye(dim);
		Vector b = zero(dim);
		SharedDiagonal model = noiseModel::Isotropic::Sigma(dim,sigma);
		sharedFactor prior(new GaussianFactor(var,A,b, model));
		result.push_back(prior);
	}
	return result;
}

///* ************************************************************************* */
//Errors GaussianFactorGraph::rhs() const {
//	Errors e;
//	BOOST_FOREACH(const sharedFactor& factor,factors_)
//		e.push_back(ediv(factor->get_b(),factor->get_sigmas()));
//	return e;
//}
//
///* ************************************************************************* */
//Vector GaussianFactorGraph::rhsVector() const {
//	Errors e = rhs();
//	return concatVectors(e);
//}
//
///* ************************************************************************* */
//pair<Matrix,Vector> GaussianFactorGraph::matrix(const Ordering& ordering) const {
//
//	// get all factors
//	GaussianFactorSet found;
//	BOOST_FOREACH(const sharedFactor& factor,factors_)
//		found.push_back(factor);
//
//	// combine them
//	GaussianFactor lf(found);
//
//	// Return Matrix and Vector
//	return lf.matrix(ordering);
//}

///* ************************************************************************* */
//VectorValues GaussianFactorGraph::assembleValues(const Vector& vs, const Ordering& ordering) const {
//	Dimensions dims = dimensions();
//	VectorValues config;
//	Vector::const_iterator itSrc = vs.begin();
//	Vector::iterator itDst;
//	BOOST_FOREACH(varid_t key, ordering){
//		int dim = dims.find(key)->second;
//		Vector v(dim);
//		for (itDst=v.begin(); itDst!=v.end(); itDst++, itSrc++)
//			*itDst = *itSrc;
//		config.insert(key, v);
//	}
//	if (itSrc != vs.end())
//		throw runtime_error("assembleValues: input vector and ordering are not compatible with the graph");
//	return config;
//}
//
///* ************************************************************************* */
//pair<Dimensions, size_t> GaussianFactorGraph::columnIndices_(const Ordering& ordering) const {
//
//	// get the dimensions for all variables
//	Dimensions variableSet = dimensions();
//
//	// Find the starting index and dimensions for all variables given the order
//	size_t j = 1;
//	Dimensions result;
//	BOOST_FOREACH(varid_t key, ordering) {
//		// associate key with first column index
//		result.insert(make_pair(key,j));
//		// find dimension for this key
//		Dimensions::const_iterator it = variableSet.find(key);
//		if (it==variableSet.end()) // key not found, now what ?
//				throw invalid_argument("ColumnIndices: this ordering contains keys not in the graph");
//		// advance column index to next block by adding dim(key)
//		j += it->second;
//	}
//
//	return make_pair(result, j-1);
//}
//
///* ************************************************************************* */
//Dimensions GaussianFactorGraph::columnIndices(const Ordering& ordering) const {
//	return columnIndices_(ordering).first;
//}
//
///* ************************************************************************* */
//pair<size_t, size_t> GaussianFactorGraph::sizeOfA() const {
//	size_t m = 0, n = 0;
//	Dimensions variableSet = dimensions();
//	BOOST_FOREACH(const Dimensions::value_type value, variableSet)
//		n += value.second;
//	BOOST_FOREACH(const sharedFactor& factor,factors_)
//		m += factor->numberOfRows();
//	return make_pair(m, n);
//}

///* ************************************************************************* */
//Matrix GaussianFactorGraph::sparse(const Ordering& ordering) const {
//
//	// get the starting column indices for all variables
//	Dimensions indices = columnIndices(ordering);
//
//	return sparse(indices);
//}
//
///* ************************************************************************* */
//Matrix GaussianFactorGraph::sparse(const Dimensions& indices) const {
//
//	// return values
//	list<int> I,J;
//	list<double> S;
//
//	// Collect the I,J,S lists for all factors
//	int row_index = 0;
//	BOOST_FOREACH(const sharedFactor& factor,factors_) {
//
//		// get sparse lists for the factor
//		list<int> i1,j1;
//		list<double> s1;
//		boost::tie(i1,j1,s1) = factor->sparse(indices);
//
//		// add row_start to every row index
//		transform(i1.begin(), i1.end(), i1.begin(), bind2nd(plus<int>(), row_index));
//
//		// splice lists from factor to the end of the global lists
//		I.splice(I.end(), i1);
//		J.splice(J.end(), j1);
//		S.splice(S.end(), s1);
//
//		// advance row start
//		row_index += factor->numberOfRows();
//	}
//
//	// Convert them to vectors for MATLAB
//	// TODO: just create a sparse matrix class already
//	size_t nzmax = S.size();
//	Matrix ijs(3,nzmax);
//	copy(I.begin(),I.end(),ijs.begin2());
//	copy(J.begin(),J.end(),ijs.begin2()+nzmax);
//	copy(S.begin(),S.end(),ijs.begin2()+2*nzmax);
//
//	// return the result
//	return ijs;
//}

///* ************************************************************************* */
//VectorValues GaussianFactorGraph::steepestDescent(const VectorValues& x0,
//		bool verbose, double epsilon, size_t maxIterations) const {
//	return gtsam::steepestDescent(*this, x0, verbose, epsilon, maxIterations);
//}
//
///* ************************************************************************* */
//boost::shared_ptr<VectorValues> GaussianFactorGraph::steepestDescent_(
//		const VectorValues& x0, bool verbose, double epsilon, size_t maxIterations) const {
//	return boost::shared_ptr<VectorValues>(new VectorValues(
//			gtsam::conjugateGradientDescent(*this, x0, verbose, epsilon,
//					maxIterations)));
//}
//
///* ************************************************************************* */
//VectorValues GaussianFactorGraph::conjugateGradientDescent(
//		const VectorValues& x0, bool verbose, double epsilon, size_t maxIterations) const {
//	return gtsam::conjugateGradientDescent(*this, x0, verbose, epsilon,
//			maxIterations);
//}
//
///* ************************************************************************* */
//boost::shared_ptr<VectorValues> GaussianFactorGraph::conjugateGradientDescent_(
//		const VectorValues& x0, bool verbose, double epsilon, size_t maxIterations) const {
//	return boost::shared_ptr<VectorValues>(new VectorValues(
//			gtsam::conjugateGradientDescent(*this, x0, verbose, epsilon,
//					maxIterations)));
//}



/* ************************************************************************* */

//template std::pair<Matrix, SharedDiagonal> combineFactorsAndCreateMatrix<vector<GaussianFactor::shared_ptr> >(
//		const vector<GaussianFactor::shared_ptr>& factors,	const Ordering& order, const Dimensions& dimensions);
//
//template std::pair<Matrix, SharedDiagonal> combineFactorsAndCreateMatrix<GaussianFactorGraph>(
//		const GaussianFactorGraph& factors,	const Ordering& order, const Dimensions& dimensions);

} // namespace gtsam