/* Copyright (c) 2015, Julian Straub <jstraub@csail.mit.edu>, Randi Cabezas <rcabezas@csail.mit.edu>
 * Licensed under the MIT license. See the license file LICENSE.
 */

#pragma once
#include <iostream>
#include <stdint.h>
#include <vector>
#include <Eigen/Dense>

#include <boost/shared_ptr.hpp>

#include <dpMM/dpMM.hpp>
#include <dpMM/cat.hpp>
#include <dpMM/dir.hpp>
#include <dpMM/niw.hpp>
#include <dpMM/sampler.hpp>
#include <dpMM/basemeasure.hpp>
#include <dpMM/niwBaseMeasure.hpp>
#include <dpMM/niwSphere.hpp>
#include <dpMM/dirBaseMeasure.hpp>
#include <mmf/mfBaseMeasure.hpp>

using namespace Eigen;
using std::cout;
using std::endl;
using boost::shared_ptr;
using std::vector;

template<typename T=double>
class DirMultiNaiveBayes : public DpMM<T>{

public:
  DirMultiNaiveBayes(std::ifstream &in, boost::mt19937 *rng);
  DirMultiNaiveBayes(const Dir<Cat<T>, T>& alpha, const vector<boost::shared_ptr<BaseMeasure<T> > >&thetas);
  DirMultiNaiveBayes(const Dir<Cat<T>, T>& alpha, const vector< vector<boost::shared_ptr<BaseMeasure<T> > > >&thetas);
  virtual ~DirMultiNaiveBayes();

  virtual void loadData(const vector<vector<Matrix<T,Dynamic,Dynamic> > > &x);//does nothing other than load data
  virtual void initialize(const vector<vector< Matrix<T,Dynamic,Dynamic> > >&x);
  virtual void initializeNoParamSampling(const vector<vector< Matrix<T,Dynamic,Dynamic> > >&x);
  virtual void initialize(const vector<vector< Matrix<T,Dynamic,Dynamic> > >&x, VectorXu &z);
  virtual void initialize(const boost::shared_ptr<ClGMMData<T> >&cld)
    {cout<<"not supported"<<endl; assert(false);};

  virtual void sampleLabels();
  virtual void MAPLabel();
  virtual void sampleParameters();

  virtual T logJoint(bool verbose=false);
  virtual const VectorXu& labels(){return z_;};
  virtual const VectorXu& getLabels(){return z_;};
  virtual uint32_t getK() const { return K_;};
  virtual uint32_t getM() const { return M_;};
  virtual uint32_t getN() const { return Nd_;};

//  virtual MatrixXu mostLikelyInds(uint32_t n);
  virtual MatrixXu mostLikelyInds(uint32_t n, Matrix<T,Dynamic,Dynamic>& logLikes);

  Matrix<T,Dynamic,1> getCounts();

  virtual void inferAll(uint32_t nIter, bool verbose=false);

  virtual void dump(std::ofstream& fOutMeans, std::ofstream& fOutCovs);
  virtual void dump_clean(std::ofstream &out);

  virtual vector<boost::shared_ptr<BaseMeasure<T> > > getThetas(uint32_t m) {
	  return(thetas_[m]);
  };
  virtual boost::shared_ptr<BaseMeasure<T> > getThetas(uint32_t m, uint32_t k) {
	  return(thetas_[m][k]);
  };

  virtual void setTheta(uint32_t m, uint32_t k, boost::shared_ptr<BaseMeasure<T> > newTheta ) {
	  thetas_[m][k] = newTheta;
  };

  virtual vector<T> evalLogLik(const vector<Matrix<T,Dynamic,1> > xnew, const vector<uint32_t> clusterInd,
							   const vector<uint32_t> comp2eval =vector<uint32_t>());

  virtual uint32_t sampleLabels(const vector<Matrix<T,Dynamic,1> > xnew,
								const vector<uint32_t> comp2eval =vector<uint32_t>());
  virtual vector<uint32_t> sampleLabels(const vector<vector<Matrix<T,Dynamic,1> > > xnew,
										const vector<uint32_t> comp2eval =vector<uint32_t>());
  virtual uint32_t MAPLabels(const vector<Matrix<T,Dynamic,1> > xnew,
							 const vector<uint32_t> comp2eval =vector<uint32_t>());
  virtual vector<uint32_t> MAPLabels(const vector<vector<Matrix<T,Dynamic,1> > > xnew,
									 const vector<uint32_t> comp2eval =vector<uint32_t>());
  virtual void updatePDF();

  vector<uint32_t> getLogEvalItersHist() {return logJointIterEval;}
  vector<T> getLogJointHist() {return logJointHist;}

protected:
  virtual T evalLogLik(const vector<Matrix<T,Dynamic,1> > xnew, const uint32_t clusterInd,
					   const vector<uint32_t> comp2eval);
  virtual uint32_t labels_sample_max(const vector<Matrix<T,Dynamic,1> > xnew,
									 const vector<uint32_t> comp2eval, const bool return_MAP_labels=false);

  uint32_t Nd_;
  uint32_t K_; //num cluseters
  uint32_t M_; //num data sources
  Dir<Cat<T>, T> dir_;
  Cat<T> pi_;
#ifdef CUDA
  SamplerGpu<T>* sampler_;
#else
  Sampler<T>* sampler_;
#endif
  virtual void initialize_sampler();
  Matrix<T,Dynamic,Dynamic> pdfs_;
//  Cat cat_;
  vector<vector<boost::shared_ptr<BaseMeasure<T> > > > thetas_;  // theta_[M][K]

  //suffiecient stats
  vector<vector<Matrix<T,Dynamic,Dynamic> > > x_; //x_[M][doc](:,word)
  VectorXu z_;

  virtual void helper_setDims();
  vector<VectorXu> dataDim;

  vector<uint32_t> logJointIterEval;
  vector<T> logJointHist;
};

// --------------------------------------- impl -------------------------------

template<typename T>
void DirMultiNaiveBayes<T>::initialize_sampler() {
	if (sampler_ != NULL) {
		delete sampler_;
		sampler_ = NULL;
	}
	//initialize sampler
	#ifdef CUDA
	  sampler_ = new SamplerGpu<T>(uint32_t(Nd_),K_,dir_.pRndGen_);
	#else
	  sampler_ = new Sampler<T>(dir_.pRndGen_);
	#endif
}

template<typename T>
DirMultiNaiveBayes<T>::DirMultiNaiveBayes(std::ifstream &in, boost::mt19937 *rng) :
 dir_(Matrix<T,2,1>::Ones(),rng), pi_(dir_.sample()),sampler_(NULL)
{
	//initialize the class from the file pointer given
	in >> M_;
	in >> K_;
	in >> Nd_;

	vector<uint32_t> dim;
	vector<baseMeasureType> type;
	Matrix<T,Dynamic,1> alpha(K_), pi(K_);

	for(uint32_t m = 0; m<M_; ++m){
		uint32_t temp;
		in >> temp;
		dim.push_back(temp);
	}

	for(uint32_t m = 0; m<M_; ++m){
		uint32_t temp;
		in >> temp;
		type.push_back(baseMeasureType(temp));
	}

	z_ = VectorXu(Nd_);
	for(uint32_t n=0; n<Nd_; ++n)
		in >> z_(n);

	for(uint32_t k=0; k<K_; ++k)
		in >> alpha(k);

	for(uint32_t k=0; k<K_; ++k)
		in >> pi(k);

	pdfs_ = Matrix<T,Dynamic,Dynamic>(Nd_,K_);
	//for(uint32_t n=0; n<Nd_*K_; ++n)
	//	in >> pdfs_(n%Nd_, (n-(n%Nd_))/Nd_);
	//	//in >> pdfs_((n-(n%Nd_))/Nd_, n%Nd_);
	////pdfs_ = pdfs_.transpose();

	//get parameters
	for(uint32_t m=0; m<M_; ++m) {
		baseMeasureType typeIter = baseMeasureType(type[m]);
		vector<boost::shared_ptr<BaseMeasure<T> > > thetaM;
		if(typeIter==NIW_SAMPLED) {
			uint32_t Diter = dim[m];
			T nu, kappa;
			Matrix<T,Dynamic,Dynamic> scatter(Diter, Diter), sigma(Diter,Diter);
			Matrix<T,Dynamic,1> theta(Diter), mu(Diter);
			for(uint32_t k=0; k<K_; ++k) {
				//get nu and kappa
					in>> nu;
					in>> kappa;
				//get theta
					for(uint32_t n=0; n<Diter; ++n)
						in >> theta(n);
					theta = theta.transpose();
				//get scatter
					for(uint32_t n=0; n<Diter*Diter; ++n)
						in >> scatter((n-(n%Diter))/Diter, n%Diter);
				//get mean
					for(uint32_t n=0; n<Diter; ++n)
						in >> mu(n);
					mu = mu.transpose();
				//get sigma
					for(uint32_t n=0; n<Diter*Diter; ++n)
						in >> sigma((n-(n%Diter))/Diter, n%Diter);

				//build theta[m][k]
				NIW<T> niw(scatter,theta,nu,kappa,rng);
				Normal<T> normal(mu,sigma,rng);
				boost::shared_ptr<NiwSampled<T> > baseIter( new NiwSampled<T>(niw, normal));

				//set
				thetaM.push_back(boost::shared_ptr<BaseMeasure<T> >(baseIter));
			}
		} else if(typeIter==NIW_SPHERE) {
			uint32_t Diter = dim[m]-1;
			T nu;
			T counts;
			Matrix<T,Dynamic,Dynamic> scatter(Diter, Diter), sigma(Diter,Diter),delta(Diter,Diter);
			Matrix<T,Dynamic,1> mu(Diter+1), mu_prior(Diter), north(Diter+1);
			for(uint32_t k=0; k<K_; ++k) {
				//get nu
					in>> nu;
					in>> counts;
				//get prior mean
					for(uint32_t n=0; n<Diter; ++n)
						in>> mu_prior(n);
				//get scatter
					for(uint32_t n=0; n<Diter*Diter; ++n)
						in >> scatter((n-(n%Diter))/Diter, n%Diter);
				//get delta
					for(uint32_t n=0; n<Diter*Diter; ++n)
						in >> delta((n-(n%Diter))/Diter, n%Diter);
				//get mean
					for(uint32_t n=0; n<(Diter+1); ++n)
						in >> mu(n);
					mu = mu.transpose();
				//get sigma
					for(uint32_t n=0; n<Diter*Diter; ++n)
						in >> sigma((n-(n%Diter))/Diter, n%Diter);
				//get north
					for(uint32_t n=0; n<(Diter+1); ++n)
						in >> north(n);
					north = north.transpose();

				//build theta[m][k]
				IW<T> iw(delta,nu, scatter, mu_prior, counts, rng);
				//IW<T> iw(delta,nu, rng);
				NiwSphere<T> niwSp(iw,rng);
				niwSp.S_ = Sphere<T>(north);
				niwSp.normalS_ = NormalSphere<T>(mu,sigma,rng);

				//set
				thetaM.push_back(boost::shared_ptr<BaseMeasure<T> >(niwSp.copy()));
				//boost::shared_ptr<NiwSphere<T> > baseIter( new NiwSphere<T>(iw,rng));
				//thetaM.push_back(boost::shared_ptr<BaseMeasure<T> >(baseIter));

			}

		} else if(typeIter==DIR_SAMPLED) {
			uint32_t Diter = dim[m];
			T localCount;
			Matrix<T,Dynamic,1> post_alpha(Diter), counts(Diter), pdf(Diter);

			for(uint32_t k=0; k<K_; ++k) {
				//get dir alpha
				for(uint32_t n=0; n<Diter; ++n)
					in >> post_alpha(n);

				//get dir counts
				for(uint32_t n=0; n<Diter; ++n)
					in >> counts(n);

				//get disc pdf
				for(uint32_t n=0; n<Diter; ++n)
					in >> pdf(n);

				in >> localCount;

				//build theta[m][k]
				Dir<Cat<T>,T> dirBase(post_alpha,counts, rng);
				DirSampled<Cat<T>,T> dirSamp(dirBase);
				Cat<T> disc(pdf,rng);
				dirSamp.disc_ = disc;
				dirSamp.count_	= localCount;

				//set
				thetaM.push_back(boost::shared_ptr<BaseMeasure<T> >(dirSamp.copy()));
			}

		} else if(typeIter==MF_T) {
			//			uint32_t Diter = dim[m];
			//			T localCount;
			Matrix<T,Dynamic,1> post_alpha(6), counts(6), pi_pdf(6);
			post_alpha.fill(1);
			counts.fill(0);
			pi_pdf.fill(0);

			Matrix<T,3,3>  R;
			Dir<Cat<T>, T> alpha(post_alpha,rng);
			Cat<T> pi(pi_pdf,rng);
			std::vector<shared_ptr<BaseMeasure<T> > > iwTs;
			std::vector<NormalSphere<T> > TGs;
			uint32_t nIter;

			for(uint32_t k=0; k<K_; ++k)
			{
        in >> nIter;
        //get dir alpha
        for(uint32_t n=0; n<6; ++n) in >> post_alpha(n);
        //get dir counts
        for(uint32_t n=0; n<6; ++n) in >> counts(n);
        //get pi pdf
        for(uint32_t n=0; n<6; ++n) in >> pi_pdf(n);
        alpha.alpha_ = post_alpha;
        alpha.setCounts(counts);
        pi.pdf(pi_pdf);
        // get rotation
        for(uint32_t n=0; n<9; ++n) in >> R(n/3,n%3);
        for(uint32_t j=0; j<6; ++j)
        {
          // load IW in tangent space
          Matrix<T,Dynamic,Dynamic> Delta(2,2);
          Matrix<T,Dynamic,Dynamic> Scatter(2,2);
          Matrix<T,Dynamic,1> mean(2);
          T nu,count;
          in >> nu;
          for(uint32_t n=0; n<4; ++n) in >> Delta(n/2,n%2);
          for(uint32_t n=0; n<4; ++n) in >> Scatter(n/2,n%2);
          for(uint32_t n=0; n<2; ++n) in >> mean(n);
          in >> count;
          IW<T> iw(Delta,nu,Scatter,mean,count,rng);
          iwTs.push_back(shared_ptr<IwTangent<T> >(
                new IwTangent<T>(iw,rng)));
          // load tangent space gaussians
          Matrix<T,Dynamic,Dynamic> Sigma(2,2);
          Matrix<T,Dynamic,1> mu(3);
          for(uint32_t n=0; n<4; ++n) in >> Sigma(n/2,n%2);
          for(uint32_t n=0; n<3; ++n) in >> mu(n);
          TGs.push_back(NormalSphere<T>(mu,Sigma,rng));
          reinterpret_cast<IwTangent<T>* >(
              iwTs[j].get())->normalS_=TGs[j];
        };
        // labels -- not set
        uint32_t Nk = 0;
        in >> Nk;
        VectorXu z(Nk);
        for(uint32_t n=0; n<Nk; ++n) in >> z(n);
        // build model
        DirMM<T> dirMM(alpha,iwTs);
        MfPrior<T> mfPrior(dirMM, nIter);
        MF<T> mf(R,pi,TGs);
				//set
				thetaM.push_back(boost::shared_ptr<BaseMeasure<T> >(
              new MfBase<T>(mfPrior, mf)));
			}

		} else {
				std::cerr << "[DirMultiNaiveBayes::dump_clean] error saving...returning" << endl;
				return;
		}
		thetas_.push_back(thetaM);
	}

	dir_ =  Dir<Cat<T>, T>(alpha,rng);
	pi_ = Cat<T>(pi,rng);

	this->initialize_sampler();

	//cout << "M:" << M_ << " K: " << K_ << " nd: " << Nd_ << endl;
	//for(uint32_t m = 0; m<M_; ++m){
	//	cout << m << ": d="<< dim[m] << ", type=" << type[m] << endl;
	//}
	//cout << "z_ " << z_.transpose() << endl;
	//cout << "alpha " << alpha.transpose() << endl;
}

template<typename T>
DirMultiNaiveBayes<T>::DirMultiNaiveBayes(const Dir<Cat<T>,T>& alpha,
    const vector<boost::shared_ptr<BaseMeasure<T> > >& thetas) :
  sampler_(NULL), K_(alpha.K_), dir_(alpha), pi_(dir_.sample()), M_(uint32_t(thetas.size()))
{
	for (uint32_t m=0; m<M_; ++m)
	{
		vector<boost::shared_ptr<BaseMeasure<T> > > temp;
		for (uint32_t k=0; k<K_; ++k)
		{
      		temp.push_back(boost::shared_ptr<BaseMeasure<T> >(thetas[m]->copy()));
      	}
      thetas_.push_back(temp);
    }

#ifndef NDEBUG
	for(uint32_t m=0; m<int(M_); ++m) {
		for(int k=0; k<int(K_); ++k) {
				thetas_[m][k]->print();
		}
	}
#endif

};

template<typename T>
DirMultiNaiveBayes<T>::DirMultiNaiveBayes(const Dir<Cat<T>,T>& alpha,
    const vector< vector<boost::shared_ptr<BaseMeasure<T> > > >& theta) :
 K_(alpha.K_), M_(uint32_t(theta.size())),dir_(alpha),
  pi_(dir_.sample()), sampler_(NULL),thetas_(theta)
{ };

template<typename T>
DirMultiNaiveBayes<T>::~DirMultiNaiveBayes()
{
  if (sampler_ != NULL) {
	delete sampler_;
	sampler_ = NULL;
  }
};

template <typename T>
Matrix<T,Dynamic,1> DirMultiNaiveBayes<T>::getCounts()
{
  return counts<T,uint32_t>(z_,K_);
};

template<typename T>
void DirMultiNaiveBayes<T>::loadData(const vector<vector<Matrix<T,Dynamic,Dynamic> > > &x){
  x_ = x;
  this->helper_setDims();
}

template<typename T>
void DirMultiNaiveBayes<T>::initializeNoParamSampling(const vector< vector< Matrix<T,Dynamic,Dynamic> > > &x)
{
  // randomly init labels from prior
  Nd_= uint32_t(x.front().size());

  z_ = VectorXu::Zero(Nd_);
  //init data and labels from given
  pi_.sample(z_);

  x_ = x;

  pdfs_.setZero(Nd_,K_);

  this->initialize_sampler();
  this->helper_setDims();
};

template<typename T>
void DirMultiNaiveBayes<T>::initialize(const vector< vector< Matrix<T,Dynamic,Dynamic> > > &x)
{
  uint32_t Nd= uint32_t(x.front().size());

  // randomly init labels from prior
  VectorXu z;
  z.setZero(Nd);
  Cat<T> pi = dir_.sample();
  pi.sample(z);

  //delegate the initialization to the main intitialization function
  this->initialize(x,z);
};

template<typename T>
void DirMultiNaiveBayes<T>::initialize(const vector< vector< Matrix<T,Dynamic,Dynamic> > > &x, VectorXu &z)
{
  Nd_= uint32_t(x.front().size());

  //init data and labels from given
  x_ = x;
  z_ = z;

  pi_ = dir_.sample();

  pdfs_.setZero(Nd_,K_);

  this->initialize_sampler();
  this->helper_setDims();
  this->sampleParameters();
};

template<typename T>
void DirMultiNaiveBayes<T>::helper_setDims()
{
	dataDim.clear();
	dataDim.reserve(M_);
	for(uint32_t m=0; m<M_; ++m) {
		VectorXu temp(x_[m].size());
		for(uint32_t t=0; t<x_[m].size(); ++t)
			temp(t) = uint32_t(x_[m][t].cols());
		dataDim.push_back(temp);
	}
}

template<typename T>
void DirMultiNaiveBayes<T>::sampleLabels()
{
	// obtain posterior categorical under labels
	pi_ = dir_.posterior(z_).sample();
	//  cout<<pi_.pdf().transpose()<<endl;
	this->updatePDF();
	// sample z_i
	sampler_->sampleDiscPdf(pdfs_,z_);
};

template<typename T>
void DirMultiNaiveBayes<T>::updatePDF() {

// compute categorical distribution over label z_i
// no need to re-compute the array and log every iteration)
VectorXd logPdf_z_value = pi_.pdf().array().log();

#pragma omp parallel for
  for(int32_t d=0; d<int32_t(Nd_); ++d)
  {
	VectorXd logPdf_z = logPdf_z_value;
	for(uint32_t m=0; m<uint32_t(M_); ++m)
	{
		for(uint32_t k=0; k<K_; ++k)
		{
			//updated to SS
			logPdf_z[k] += thetas_[m][k]->logLikelihoodFromSS(x_[m][d]);
		}
	}
//    cout<<endl;
    // make pdf sum to 1. and exponentiate
    pdfs_.row(d) = (logPdf_z.array()-logSumExp(logPdf_z)).exp().matrix().transpose();
//    cout<<pi_.pdf().transpose()<<endl;
//    cout<<pdf.transpose()<<" |.|="<<pdf.sum();
//    cout<<" z_i="<<z_[d]<<endl;
  }

}

template<typename T>
void DirMultiNaiveBayes<T>::MAPLabel()
{
	/* it chooses the MAP label rather than sampling */
	this->updatePDF();

	#pragma omp parallel for
	 for(int32_t d=0; d<int32_t(Nd_); ++d) {
		 int r,c;
		 pdfs_.row(d).maxCoeff(&r, &c);
		 z_(d) = c;
	 }
};

template<typename T>
void DirMultiNaiveBayes<T>::sampleParameters()
{
//unpacks the contains here vector<vector<Matrix>> into what the posterior expects Matrix
	MatrixXu dim(M_,K_);
	for(uint32_t m=0; m<M_; ++m) {
		#pragma omp parallel for
		for(int32_t k=0; k<int32_t(K_); ++k) {
			VectorXu temp = (z_.array()==k).select(dataDim[m],0);
			dim(m,k) = temp.sum();
		}
	}

	for(int32_t m=0; m<M_; ++m) {
	#ifdef _WINDOWS
		//#pragma omp parallel for
		for(int32_t k=0; k<int32_t(K_); ++k) {
	#else
		//#pragma omp parallel for
		#pragma omp parallel for schedule(dynamic)
		for(int32_t k=0; k<int32_t(K_); ++k) {
	#endif

			if(dim(m,k)!=0) {

				//Matrix<T,Dynamic,1> ssIn = Matrix<T,Dynamic,1>::Zero(x_[m].front().rows());
				vector< Matrix<T,Dynamic,1> >dataIn;
				dataIn.reserve(dim(m,k));

				uint32_t count=0;

				for(int32_t d=0; d<Nd_; ++d) {
					if(z_[d]==k) {
						int add_size =int(x_[m][d].cols());
						//ssIn += x_[m][d]; //update iteration SS
						dataIn.push_back(x_[m][d]);
						count+=add_size;

						if(count==dim(m,k))
							break; //early out if you found them all
					}
				}
				//update values
				//thetas_[m][k]->posteriorFromSS(ssIn);

//        if(thetas_[m][k]->getBaseMeasureType() == MF_T)
//        {
//          cout<<"MF "<<k<<" ---------------- "<<endl;
//        }
				//always sends zeros and look for zeros
				thetas_[m][k]->posteriorFromSS(dataIn,VectorXu::Zero(dim(m,k)),0);

			} else {
				//Matrix<T,Dynamic,1> ssIn = Matrix<T,Dynamic,1>::Zero(x_[m].front().rows());
				//ssIn[0]=1; //set counts to 1 to avoid inf
				//the posterior needs to reset
				//thetas_[m][k]->posteriorFromSS(ssIn);
				//passing in one data point (all zeros, with 1 index value=0 and looking for 1)

				vector<Matrix<T,Dynamic,1> >dataIn;
				dataIn.push_back(Matrix<T,Dynamic,1>::Zero(x_[m].front().rows(),1));
				//the posterior needs to reset
        VectorXu zz = VectorXu::Zero(1);
				thetas_[m][k]->posteriorFromSS(dataIn,zz,1);
			}
		}
	}
};

template<typename T>
T DirMultiNaiveBayes<T>::logJoint(bool verbose)
{
  T logJoint = dir_.logPdf(pi_);
  if(verbose)
  	cout<<"\tlog p(pi)="<< logJoint << endl;

  for(int32_t m=0; m<int32_t(M_); ++m) {
	  T logPriorM = 0;
	  #pragma omp parallel for reduction(+:logPriorM)
	  for (int32_t k=0; k<int32_t(K_); ++k) {
		logPriorM = logPriorM + thetas_[m][k]->logPdfUnderPrior();
	  }
	  logJoint+=logPriorM;
	  if(verbose)
	  	cout<<"\tlog p(theta_" << m << ")="<< logPriorM << endl;

  }

	for (int32_t m=0; m<int32_t(M_); ++m) {
		T logThetaM=0;
		#pragma omp parallel for reduction(+:logThetaM)
		for (int32_t d=0; d<int32_t(Nd_); ++d) {
			logThetaM = logThetaM + thetas_[m][z_[d]]->logLikelihoodFromSS(x_[m][d]);
		}
		logJoint += logThetaM;
		if(verbose)
			cout<<"\tlog p(x|z,theta_" << m << ")=" << logThetaM << endl;

	}
    if(verbose)
		cout<<"log p(pi)*p(theta)*p(x|z,theta)=" << logJoint << endl;

  return logJoint;
};

template<typename T>
void DirMultiNaiveBayes<T>::inferAll(uint32_t nIter, bool verbose)
{
  if(verbose){
  	cout<<"[DirMultiNaiveBayes::inferALL] ------ inferingALL (nIter=" << nIter << ") ------"<<endl;
  	if(Nd_<=100) {
		cout <<"initial labels:"<< endl;
  		cout<<this->labels().transpose()<<endl;
	}
  }

  logJointIterEval.clear(); logJointHist.clear();
  if(verbose) {
	//all iterations stored
	logJointIterEval.reserve(nIter); logJointHist.reserve(nIter);
  } else {
	  //only mod 100 stored
	  logJointIterEval.reserve(int((nIter/100) + 1)); logJointHist.reserve(int((nIter/100) + 1));
  }

  for(uint32_t t=0; t<nIter; ++t)
  {
    this->sampleLabels();
    this->sampleParameters();
    if(verbose)
    {
      for(int m=0; m<int(M_); ++m) {
        for(int k=0; k<int(K_); ++k) {
          thetas_[m][k]->print();
        }
      }
    }
    if(verbose || t%int(ceil(nIter/100.))==0)
    {
      VectorXu Ns = counts<uint32_t,uint32_t>(
          this->labels(),K_).transpose();
      uint32_t K = K_;
      for(uint32_t k = 0; k<K_; ++k)
        if (Ns(k) == 0) --K;
      cout<<"@i "<<t<<": # "
        <<K<<" "<<std::setw(1)<<Ns.transpose() <<endl;

      T iterLogJoint = this->logJoint(true) ;
      //log iterJoint Prob
      logJointIterEval.push_back(t);
      logJointHist.push_back(iterLogJoint);

      if(Nd_<=10) {
        cout << "[" << std::setw(3)<< std::setfill('0')
          << t <<"] label: "
          << this->labels().transpose()
          << " [joint= " << std::setw(6) << iterLogJoint << "]"<< endl;
      } else {
        cout << "[" << std::setw(3)<< std::setfill('0')
          << t <<"] joint= "
          << std::setw(6) << iterLogJoint << endl;
      }
    }
//    if(verbose)
//    {
//      VectorXu Ns = counts<uint32_t,uint32_t>(this->labels(),K_).transpose();
//      uint32_t K = K_;
//      for(uint32_t k = 0; k<K_; ++k)
//        if (Ns(k) == 0) --K;
//      cout<<"@i "<<t<<": # "<<K<<" "<<std::setw(1) <<Ns.transpose() <<endl;
//    }
  }
  //keeps the MAP label in memory
  this->MAPLabel();
  this->sampleParameters();
}

template <typename T>
void DirMultiNaiveBayes<T>::dump(std::ofstream& fOutMeans, std::ofstream& fOutCovs)
{
	cout << "dumping MultiObs naiveBayes" << endl;
	cout << "doc index: " << endl;
	cout << this->labels().transpose() << endl;

	cout << "printing num components: " << endl;
	cout << M_ << endl;

	cout << "printing cluster params: " << endl;
	cout << K_ << endl;

	for(uint32_t m=0; m<M_; ++m) {
		cout << "component: " << m  << endl;
		for(uint32_t k=0; k<K_; ++k) {
			cout << "theta: " << k  << endl;
			thetas_[m][k]->print();
		}
	}

	cout << "printing mixture params: " << endl;
	pi_.print();
}

template <typename T>
void DirMultiNaiveBayes<T>::dump_clean(std::ofstream &out){
	//clean dump, only data with specific format
//FORMAT:
	//M 1x1
	//K	1x1
	//Nd 1x1
	//D[m] 1xM
	//Type[m] 1xM
	//labels 1xNd
	//Dir alpha 1xK
	//pi pdf 1xK
	//pdf KxNd
	// mixture parameters
	//params Loop over M then K each contains data type for specific type
	//for type 1 (NIWSampled)
		//---prior--- (NIW)
			//nu 1x1
			//kappa 1x1
			//theta 1xD
			//scatter DxD
		//---estimate (normal)
			//mu 1xD
			//Sigma DxD
	//for type 2 (NIWSphereFull)
		//----prior--- (IW)
			//nu 1x1
			//count 1x1
			//mean 1x(D-1)
			//scatter (D-1)x(D-1)
			//Delta	 (D-1)x(D-1)
		//--posterior (NormalSphere)
			//mean 1x(D-1)
			//Sigma (D-1)x(D-1)
		//--sphere--- (Sphere)
			//north 1x(D-1)
	//for type 3 (DirSampled)
		//--posterior (Dir)
			//alpha 1xK
			//counts 1xK
		//--distribution (Cat)
			//pdf 1xK
		//--counts (scalar)
			//counts 1x1
	//logJoint history
		// Niter 1x1
		// iterValue 1xNiter (iteration corresponding to the logValue)
		// logJoint	 1xNiter (logJoint)

	//this fixes issues with eigen matrices printing (eg, 00-0.7 )
	int curPres = int(out.precision());
	out.precision(10);
	IOFormat fullPresPrint(FullPrecision,DontAlignCols);

	//prints headers
	out << M_ << endl
		 << K_ << endl
		 << Nd_ << endl;

	//print dim
	for(uint32_t m=0; m<M_; ++m) {
		vector<boost::shared_ptr<BaseMeasure<T> > >  theta_base = this->getThetas(m);
		uint32_t temp = theta_base.front()->getDim();
		out << temp << " ";
		//out << x_[m].front().rows() << " ";
	}
	out << endl;
	//print type
	for(uint32_t m=0; m<M_; ++m) {
		out << thetas_[m].front()->getBaseMeasureType() << " ";
	}
	out << endl;

	//print labels
	out << this->labels().transpose() << endl;

	//print mixture parameters
	out << this->dir_.alpha_.transpose() << endl;
	out << this->pi_.pdf_.transpose().format(fullPresPrint) << endl;
	//out << this->pdfs_.transpose().format(fullPresPrint) << endl;

	//print parameters
	for(uint32_t m=0; m<M_; ++m) {
		vector<boost::shared_ptr<BaseMeasure<T> > >  theta_base = this->getThetas(m);
		for(uint32_t k=0; k<K_; ++k) {
			baseMeasureType type = theta_base[k]->getBaseMeasureType();
			if(type==NIW_SAMPLED) {
				boost::shared_ptr<NiwSampled<T> >  *theta_iter =
						reinterpret_cast<boost::shared_ptr<NiwSampled<T> >* >( &theta_base[k]);
					//printing prior
					NIW<T> prior = theta_iter->get()->niw0_;
					out << prior.nu_				 << endl <<
						    prior.kappa_			 << endl <<
						    prior.theta_.transpose() << endl <<
						    prior.Delta_.format(fullPresPrint)	<<	endl;

					//printing posterior
					Normal<T> norm = theta_iter->get()->normal_;
					out << norm.mu_.transpose() << endl;
					out << norm.Sigma().format(fullPresPrint) << endl;
			} else if(type==NIW_SPHERE) {
        boost::shared_ptr<NiwSphere<T> >  *theta_iter =
          reinterpret_cast<boost::shared_ptr<NiwSphere<T> >* >(
              &theta_base[k]);
				//prior
				IW<T> prior = theta_iter->get()->iw0_;
				out << prior.nu_ 				 << endl
					 << prior.count()			 << endl
				 	 << prior.mean().transpose() << endl
					 << prior.scatter().format(fullPresPrint)			 << endl
					 << prior.Delta_.format(fullPresPrint)			 << endl;

				//posterior
				NormalSphere<T> norm = theta_iter->get()->normalS_;
				out << norm.getMean().transpose().format(fullPresPrint) << endl;
				out << norm.Sigma().format(fullPresPrint) << endl;
				//sphere
				Sphere<T> sp = theta_iter->get()->S_;
				out << sp.north().transpose() << endl;
			} else if(type==DIR_SAMPLED) {
				boost::shared_ptr<DirSampled<Cat<T>,T> >  *theta_iter =
						reinterpret_cast<boost::shared_ptr<DirSampled<Cat<T>,T> >* >( &theta_base[k]);
				//posterior
				Dir<Catd,T>  post = theta_iter->get()->dir0_;
				out <<	post.alpha_.transpose()		<< endl;
				out <<	post.counts().transpose()	<< endl;

				//distribution
				Catd dist = theta_iter->get()->disc_;
				out <<	dist.pdf_.transpose() << endl;

				//counts
				T counts = theta_iter->get()->count_;
				out << counts << endl;

			} else if(type==MF_T) {
        boost::shared_ptr<MfBase<T> > *theta_iter =
          reinterpret_cast<boost::shared_ptr<MfBase<T> >* >(
              &theta_base[k]);

        out<< theta_iter->get()->mf0_.T_<<endl;

				//posterior
				Dir<Cat<T>, T>  post = theta_iter->get()->mf0_.dirMM().Alpha();
				out <<	post.alpha_.transpose()		<< endl;
				out <<	post.counts().transpose()	<< endl;
				Cat<T> pi = theta_iter->get()->mf0_.dirMM().Pi();
				out <<	pi.pdf().transpose() << endl;

        out << theta_iter->get()->mf_.R().format(fullPresPrint)<<endl;
        for(uint32_t j=0; j<6; ++j)
        {
          out<< theta_iter->get()->mf0_.theta(j)->iw0_.nu_ << endl;
          out<< theta_iter->get()->mf0_.theta(j)->iw0_.Delta_ << endl;
          out<< theta_iter->get()->mf0_.theta(j)->iw0_.scatter() << endl;
          out<< theta_iter->get()->mf0_.theta(j)->iw0_.mean() << endl;
          out<< theta_iter->get()->mf0_.theta(j)->iw0_.count() << endl;
          out<< theta_iter->get()->mf0_.theta(j)->normalS_.Sigma() << endl;
          out<< theta_iter->get()->mf0_.theta(j)->normalS_.getMean() << endl;
        }

        // output the labeling
        out<<theta_iter->get()->mf0_.dirMM().labels().size()<<endl;
	      out<<theta_iter->get()->mf0_.dirMM().labels().transpose()<<endl;

			} else {
					std::cerr << "[DirMultiNaiveBayes::dump_clean] error saving...returning" << endl;
					return;
			}
		}
	}

	//print logHistory
	out << int(logJointHist.size()) << endl;
	for(int i=0; i<logJointIterEval.size(); ++i)
		out << logJointIterEval[i] << " ";
	out << endl;

	for(int i=0; i<logJointHist.size(); ++i)
		out << logJointHist[i] << " ";
	out << endl;

	out.precision(curPres);

}

//template <typename T>
//void DirMultiNaiveBayes<T>::dump_clean(std::ofstream &out){
//	streambuf *coutbuf = std::cout.rdbuf(); //save old cout buffer
//	cout.rdbuf(out.rdbuf()); //redirect std::cout to fout1 buffer
//	this->dump_clean(); //write using cout to the specified buffer
//	std::cout.rdbuf(coutbuf); //reset to standard output again
//}

template <typename T>
T DirMultiNaiveBayes<T>::evalLogLik(const vector<Matrix<T,Dynamic,1> > xnew,
							  const uint32_t clusterInd, const vector<uint32_t> comp2eval)
{
	//T logJoint = pi_.pdf_(clusterInd);
	T logJoint  = 0;
	for (int32_t m=0; m<int32_t(comp2eval.size()); ++m)
	{
		logJoint += thetas_[comp2eval[m]][clusterInd]->logLikelihoodFromSS(xnew[m]);
	}

  return logJoint;
}

template <typename T>
vector<T> DirMultiNaiveBayes<T>::evalLogLik(const vector<Matrix<T,Dynamic,1> > xnew,
											const vector<uint32_t> clusterInd,
											const vector<uint32_t> comp2eval) {

	vector<uint32_t> comp2evalLocal = comp2eval;
	if(comp2evalLocal.empty()) {
		for(uint32_t m=0; m<M_; ++m)
			comp2evalLocal.push_back(m);
	}

	vector<T> out;
	for(uint32_t k=0; k<uint32_t(clusterInd.size()); ++k) {
		out.push_back(this->evalLogLik(xnew,clusterInd[k],comp2evalLocal));
	}
	return(out);
}

  template <typename T>
uint32_t DirMultiNaiveBayes<T>::labels_sample_max(const
    vector<Matrix<T,Dynamic,1> > xnew, const vector<uint32_t>
    comp2eval, const bool return_MAP_labels)
{
  /* xnew in the form x[docs][m][SS] */
  VectorXd logPdf_z = pi_.pdf().array().log();

  for(int32_t m=0; m<comp2eval.size(); ++m)
  {
    for(int32_t k=0; k<int32_t(K_); ++k)
    {
      logPdf_z[k] += thetas_[comp2eval[m]][k]->logLikelihoodFromSS(
          xnew[m]);
    }
  }

  // make pdf sum to 1. and exponentiate
  Matrix<T,Dynamic,Dynamic> pdfLocal =  Matrix<T,Dynamic,Dynamic>(1,K_);
  pdfLocal = (logPdf_z.array()-logSumExp(logPdf_z)).exp().matrix().transpose();

  VectorXu zout = VectorXu(1);

  if(return_MAP_labels) {
    // return MAP label
    int r,c;
    pdfLocal.maxCoeff(&r, &c);
    zout(0) = c;
  } else {
    // sample z_i
    sampler_->sampleDiscPdf(pdfLocal,zout);
  }

  return(zout(0));
};

template <typename T>
uint32_t DirMultiNaiveBayes<T>::sampleLabels(const vector<Matrix<T,Dynamic,1> > xnew,
											 const vector<uint32_t> comp2eval) {
	return(this->labels_sample_max(xnew, comp2eval, false));
}

template <typename T>
vector<uint32_t> DirMultiNaiveBayes<T>::sampleLabels(const vector<vector<Matrix<T,Dynamic,1> > > xnew,
													 const vector<uint32_t> comp2eval)
{
	/* xnew in the form x[doc][m][SS] */
	vector<uint32_t> out;
	for(uint32_t d=0; d<xnew.size(); ++d) {
		out.push_back(this->sampleLabels(xnew[d],comp2eval));
	}
	return(out);
}

template <typename T>
uint32_t DirMultiNaiveBayes<T>::MAPLabels(const vector<Matrix<T,Dynamic,1> > xnew,
										  const vector<uint32_t> comp2eval) {
	return(this->labels_sample_max(xnew, comp2eval, true));
}

template <typename T>
vector<uint32_t> DirMultiNaiveBayes<T>::MAPLabels(const vector<vector<Matrix<T,Dynamic,1> > > xnew,
												  const vector<uint32_t> comp2eval) {
	/* xnew in the form x[docs][m][SS] */
	vector<uint32_t> out;
	for(uint32_t d=0; d<xnew.size(); ++d) {
		out.push_back(this->MAPLabels(xnew[d],comp2eval));
	}
	return(out);
}

template<typename T>
MatrixXu DirMultiNaiveBayes<T>::mostLikelyInds(uint32_t n, Matrix<T,Dynamic,Dynamic>& logLikes)
{
  MatrixXu inds = MatrixXu::Zero(n,K_);
  logLikes = Matrix<T,Dynamic,Dynamic>::Ones(n,K_);

#pragma omp parallel for
  for (int32_t k=0; k<K_; ++k)
  {
    for (uint32_t i=0; i<z_.size(); ++i)
      if(z_(i) == k)
      {
        T logLike = 0.;
        // iterate over datasources and sum up their logLikes
        for(uint32_t m=0; m<uint32_t(M_); ++m)
        {
          logLike += thetas_[m][z_[i]]->logLikelihoodFromSS(x_[m][i]);
        }
        // keep only the top n and sorted
        for (uint32_t j=0; j<n; ++j)
          if(logLikes(j,k) < logLike)
          {
            for(uint32_t l=n-1; l>j; --l)
            {
              logLikes(l,k) = logLikes(l-1,k);
              inds(l,k) = inds(l-1,k);
            }
            logLikes(j,k) = logLike;
            inds(j,k) = i;
//            cout<<"after update "<<logLike<<endl;
//            Matrix<T,Dynamic,Dynamic> out(n,K_*2);
//            out<<logLikes.cast<T>(),inds.cast<T>();
//            cout<<out<<endl;
            break;
          }
      }
  }
  cout<<"::mostLikelyInds: logLikes"<<endl;
  cout<<logLikes<<endl;
  cout<<"::mostLikelyInds: inds"<<endl;
  cout<<inds<<endl;
  return inds;
};


/**
 * \copyright
 * Copyright (c) 2012-2017, OpenGeoSys Community (http://www.opengeosys.org)
 *            Distributed under a Modified BSD License.
 *              See accompanying file LICENSE.txt or
 *              http://www.opengeosys.org/project/license
 */

#include "Utils.h"

#include <cmath>

#ifndef Q_MOC_RUN // to avoid Qt4 bug, https://bugreports.qt.io/browse/QTBUG-22829
#include <boost/math/constants/constants.hpp>
#endif

#include "MeshLib/Elements/Element.h"
#include "MeshLib/Node.h"

namespace ProcessLib
{
namespace LIE
{

namespace
{

// Computed normal vector is oriented in the left direction of the given line element
// such that computeRotationMatrix2D() returns the indentity matrix for line elements
// parallel to a vector (1,0,0)
void computeNormalVector2D(MeshLib::Element const& e, Eigen::Vector3d & normal_vector)
{
    assert(e.getGeomType() == MeshLib::MeshElemType::LINE);

    auto v1 = (*e.getNode(1)) - (*e.getNode(0));
    normal_vector[0] = -v1[1];
    normal_vector[1] = v1[0];
    normal_vector[2] = 0.0;
    normal_vector.normalize();
}

// Compute a rotation matrix from global coordinates to local coordinates whose y axis
// should be same as the given normal vector
void computeRotationMatrix2D(Eigen::Vector3d const& n, Eigen::MatrixXd& matR)
{
    matR(0,0) = n[1];
    matR(0,1) = -n[0];
    matR(1,0) = n[0];
    matR(1,1) = n[1];
}

} // no named namespace

void computeNormalVector(MeshLib::Element const& e, Eigen::Vector3d & normal_vector)
{
    if (e.getDimension() == 1)
        computeNormalVector2D(e, normal_vector);
    else
        OGS_FATAL("2D elements are not supported in computeNormalVector()");
}

void computeRotationMatrix(Eigen::Vector3d const& normal_vector, int dim, Eigen::MatrixXd &matR)
{
    if (dim==2)
        computeRotationMatrix2D(normal_vector, matR);
    else
        OGS_FATAL("2D elements are not supported in computeNormalVector()");
}

}  // namespace LIE
}  // namespace ProcessLib


//
// Copyright 2012 Christian Henning, Andreas Pokorny, Lubomir Bourdev
//
// Distributed under the Boost Software License, Version 1.0
// See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt
//
#ifndef BOOST_GIL_EXTENSION_TOOLBOX_METAFUNCTIONS_PIXEL_BIT_SIZE_HPP
#define BOOST_GIL_EXTENSION_TOOLBOX_METAFUNCTIONS_PIXEL_BIT_SIZE_HPP

#include <boost/gil/bit_aligned_pixel_reference.hpp>
#include <boost/gil/packed_pixel.hpp>

#include <boost/mpl/accumulate.hpp>
#include <boost/mpl/int.hpp>

namespace boost{ namespace gil {

/// pixel_bit_size metafunctions
/// \brief Accumulates the all channel size.
///
/// \code
/// using image_t = bit_aligned_image5_type<16, 16, 16, 8, 8, devicen_layout_t<5>>::type;
/// const int size = pixel_bit_size<image_t::view_t::reference>::value;
/// \endcode
template< typename PixelRef >
struct pixel_bit_size : mpl::int_<0> {};

template <typename B, typename C, typename L, bool M>
struct pixel_bit_size<bit_aligned_pixel_reference<B,C,L,M> > : mpl::int_< mpl::accumulate< C, mpl::int_<0>, mpl::plus<mpl::_1, mpl::_2> >::type::value >{};

template <typename B, typename C, typename L, bool M>
struct pixel_bit_size<const bit_aligned_pixel_reference<B,C,L,M> > : mpl::int_< mpl::accumulate< C, mpl::int_<0>, mpl::plus<mpl::_1, mpl::_2> >::type::value >{};

template <typename B, typename C, typename L>
struct pixel_bit_size<packed_pixel<B,C,L> > : mpl::int_< mpl::accumulate< C, mpl::int_<0>, mpl::plus<mpl::_1, mpl::_2> >::type::value >{};

template <typename B, typename C, typename L>
struct pixel_bit_size<const packed_pixel<B,C,L> > : mpl::int_< mpl::accumulate< C, mpl::int_<0>, mpl::plus<mpl::_1, mpl::_2> >::type::value >{};

} // namespace gil
} // namespace boost

#endif


#include <Ziran/CS/Util/Debug.h>
#include <Eigen/SparseLU>
#include "EigenSparseLU.h"

namespace ZIRAN {

template <class T>
void eigenSparseLUSolve(const Eigen::SparseMatrix<T, Eigen::RowMajor>& mat, Eigen::Map<const Vector<T, Eigen::Dynamic>> rhs, Eigen::Map<Vector<T, Eigen::Dynamic>> to_solve)
{
    using SpMatCol = Eigen::SparseMatrix<T, Eigen::ColMajor>;
    SpMatCol copy = mat;
    Eigen::SparseLU<SpMatCol> solver;
    solver.analyzePattern(copy);
    solver.factorize(copy);
    to_solve = solver.solve(rhs);
    ZIRAN_VERB_IF(solver.info() == Eigen::NumericalIssue, "NumericalIssue");
    ZIRAN_VERB_IF(solver.info() == Eigen::InvalidInput, "InvalidInput");
    ZIRAN_VERB_IF(solver.info() == Eigen::NoConvergence, "NoConvergence");
    ZIRAN_ASSERT(solver.info() == Eigen::Success);
}

template void eigenSparseLUSolve<double>(const Eigen::SparseMatrix<double, Eigen::RowMajor>&, Eigen::Map<const Vector<double, Eigen::Dynamic>>, Eigen::Map<Vector<double, Eigen::Dynamic>>);
template void eigenSparseLUSolve<float>(const Eigen::SparseMatrix<float, Eigen::RowMajor>&, Eigen::Map<const Vector<float, Eigen::Dynamic>>, Eigen::Map<Vector<float, Eigen::Dynamic>>);
} // namespace ZIRAN


#include <tiny_type_traits.h>
#include <tiny_constants.h>

#define BOOST_AUTO_TEST_MAIN
#include <boost/test/auto_unit_test.hpp>
#include <boost/test/unit_test_suite.hpp>
#include <boost/test/floating_point_comparison.hpp>
#include <boost/test/test_tools.hpp>

template<size_t stride, typename type_traits, typename op_type, typename real_type>
void Verify (op_type _result, real_type value)
{
	real_type result[stride];
	type_traits::store_op_type(result,_result);
	real_type tol = real_type(0.1); // 1/10 percent
	for (size_t i=0; i<stride;++i)
	  {
	    BOOST_CHECK_CLOSE(result[i],value,tol);
	  }
}

template <typename type_traits>
void test_traits()
{
  using namespace tiny;

	typedef typename type_traits::real_type	real_type;
	typedef typename type_traits::op_type 	op_type;
	enum {stride = type_traits::stride};

	BOOST_REQUIRE(stride>0); // Types needs to contain at least one element.

	real_type const zero 	= detail::zero<real_type>();
	real_type const one  	= detail::one<real_type>();
	real_type const two  	= detail::two<real_type>();
	real_type const four 	= detail::four<real_type>();
  //real_type const eight 	= detail::eight<real_type>();
  //real_type const sixteen	= eight+eight;
	real_type const quarter = one/four;
	real_type const half 	= one/two;

	real_type const _true  	= detail::one<real_type>();
	real_type const _false 	= detail::zero<real_type>();

	op_type const _zero 	= type_traits::set_op_type(zero);
	op_type const _one  	= type_traits::set_op_type(one);
	op_type const _two  	= type_traits::set_op_type(two);
	op_type const _four 	= type_traits::set_op_type(four);
	op_type const _quarter 	= type_traits::set_op_type(quarter);
	op_type const _half 	= type_traits::set_op_type(half);

	op_type const _mone 	= type_traits::set_op_type(-one);
	op_type const _mtwo 	= type_traits::set_op_type(-two);
	op_type const _mfour 	= type_traits::set_op_type(-four);
	op_type const  _mquarter	= type_traits::set_op_type(-quarter);
	op_type const  _mhalf 	= type_traits::set_op_type(-half);

	// Verify correctness of load & store
	Verify<stride,type_traits>(_one,one);

	// Verify correctness of assign
	op_type a = _one;
	Verify<stride,type_traits>(a,one);

  op_type x;
	// Verify correctness of set_pad_value
	// TODO

	// Verify correctness of equal
	Verify<stride,type_traits>(type_traits::equal(_one,_one),_true);
	Verify<stride,type_traits>(type_traits::equal(_one,_zero),_false);

	// Verify correctness of not_equal
	Verify<stride,type_traits>(type_traits::not_equal(_one,_zero),_true);
	Verify<stride,type_traits>(type_traits::not_equal(_one,_one),_false);

	// Verify correctness of less_than
	Verify<stride,type_traits>(type_traits::less_than(_one,_zero),_false);
	Verify<stride,type_traits>(type_traits::less_than(_one,_one),_false);
	Verify<stride,type_traits>(type_traits::less_than(_zero,_one),_true);

	// Verify correctness of greater_than
	Verify<stride,type_traits>(type_traits::greater_than(_one,_zero),_true);
	Verify<stride,type_traits>(type_traits::greater_than(_one,_one),_false);
	Verify<stride,type_traits>(type_traits::greater_than(_zero,_one),_false);

	// Verify correctness of less_than_or_equal
	Verify<stride,type_traits>(type_traits::less_than_or_equal(_one,_zero),_false);
	Verify<stride,type_traits>(type_traits::less_than_or_equal(_one,_one),_true);
	Verify<stride,type_traits>(type_traits::less_than_or_equal(_zero,_one),_true);

	// Verify correctness of greater_than_or_equal
	Verify<stride,type_traits>(type_traits::greater_than_or_equal(_one,_zero),_true);
	Verify<stride,type_traits>(type_traits::greater_than_or_equal(_one,_one),_true);
	Verify<stride,type_traits>(type_traits::greater_than_or_equal(_zero,_one),_false);

	// Verify correctness of add_assign
	x = _one;
	Verify<stride,type_traits>(type_traits::add_assign(x,_one),two);
	Verify<stride,type_traits>(x,two); // Verify that x has been updated

	// Verify correctness of sub_assign
	x = _two;
	Verify<stride,type_traits>(type_traits::sub_assign(x,_one),one);
	Verify<stride,type_traits>(x,one);// Verify that x has been updated

	// Verify correctness of mul_assign
	x = _two;
	Verify<stride,type_traits>(type_traits::mul_assign(x,_two),four);
	Verify<stride,type_traits>(x,four);// Verify that x has been updated

	// Verify correctness of div_assign
	x = _four;
	Verify<stride,type_traits>(type_traits::div_assign(x,_two),two);
	Verify<stride,type_traits>(x,two);// Verify that x has been updated

	// Verify correctness of add
	Verify<stride,type_traits>(type_traits::add(_one,_zero),one);

	// Verify correctness of sub
	Verify<stride,type_traits>(type_traits::sub(_one,_one),zero);

	// Verify correctness of mul
	Verify<stride,type_traits>(type_traits::mul(_two,_two),four);

	// Verify correctness of div
	Verify<stride,type_traits>(type_traits::div(_four,_two),two);

	// Verify correctness of negate
	Verify<stride,type_traits>(type_traits::negate(_one),-one);
	Verify<stride,type_traits>(type_traits::negate(_mone),one);

	// Verify correctness of sign
	Verify<stride,type_traits>(type_traits::sign(_mtwo),-one);
	Verify<stride,type_traits>(type_traits::sign(_zero),zero);
	Verify<stride,type_traits>(type_traits::sign(_two),one);

	// Verify correctness of round
	Verify<stride,type_traits>(type_traits::round(_half),one);
	Verify<stride,type_traits>(type_traits::round(_quarter),zero);
	Verify<stride,type_traits>(type_traits::round(_mhalf),-one);
	Verify<stride,type_traits>(type_traits::round(_mquarter),zero);

	// Verify correctness of max_assign
	x = _zero;
	Verify<stride,type_traits>(type_traits::max_assign(x,_half),half);
	Verify<stride,type_traits>(x,half);
	x = _mfour;
	Verify<stride,type_traits>(type_traits::max_assign(x,_mtwo),-two);
	Verify<stride,type_traits>(x,-two);
	// Verify correctness of min_assign
	x = _zero;
	Verify<stride,type_traits>(type_traits::min_assign(x,_half),zero);
	Verify<stride,type_traits>(x,zero);
	x = _mfour;
	Verify<stride,type_traits>(type_traits::min_assign(x,_mtwo),-four);
	Verify<stride,type_traits>(x,-four);

	// Verify correctness of max
	Verify<stride,type_traits>(type_traits::max(_half,_zero),half);
	Verify<stride,type_traits>(type_traits::max(_mtwo,_mfour),-two);

	// Verify correctness of min
	Verify<stride,type_traits>(type_traits::min(_half,_zero),zero);
	Verify<stride,type_traits>(type_traits::min(_mtwo,_mfour),-four);

	// Verify correctness of abs
	Verify<stride,type_traits>(type_traits::abs(_one),one);
	Verify<stride,type_traits>(type_traits::abs(_mone),one);

	// Verify correctness of floor
	Verify<stride,type_traits>(type_traits::floor(_half),zero);
	Verify<stride,type_traits>(type_traits::floor(_mhalf),-one);

	// Verify correctness of ceil
	Verify<stride,type_traits>(type_traits::ceil(_half),one);
	Verify<stride,type_traits>(type_traits::ceil(_mhalf),zero);

	// Verify correctness of sqrt
	Verify<stride,type_traits>(type_traits::sqrt(_four),two);

	// Verify correctness of rsqrt
	Verify<stride,type_traits>(type_traits::rsqrt(_four),half);

	// TODO: For horizontal functions - consider cases where padding has been used

	real_type const tol = real_type(0.1);

	// Verify correctness of horizontal sum
	switch (stride)
	{
//	case 4:  // 2009-07-14 Kenny:  warning: case label value exceeds maximum value for type (for ScalarTrait)
//	  BOOST_CHECK_CLOSE(type_traits::sum(_one),four,tol);
//	  break;
	case 1:
	  BOOST_CHECK_CLOSE(type_traits::sum(_one),one,tol);
	  break;
	default: break;
	}
	// Verify correctness of horizontal mul
	switch (stride)
	{
//	case 4:
//	  BOOST_CHECK_CLOSE(type_traits::mul(_two),sixteen,tol);
//	  break;
	case 1:
	  BOOST_CHECK_CLOSE(type_traits::mul(_two),two,tol);
	  break;
	default: break;
	}

	// Verify correctness of horizontal min and max
	// Prepare variables

	size_t const size = stride;
	real_type scalars[size];
	for (size_t i = 0 ; i<size;++i)
	{
		scalars[i]=tiny::detail::numeric_cast<real_type>(i+1);
	}

	x = type_traits::load_op_type(scalars);

	BOOST_CHECK(type_traits::min(x)==one);
	BOOST_CHECK(type_traits::max(x)==tiny::detail::numeric_cast<real_type>(size));
}

BOOST_AUTO_TEST_SUITE(tiny_traits);

BOOST_AUTO_TEST_CASE(default_traits_testing)
{

	typedef tiny::ScalarTraits<float>  float_traits;
	typedef tiny::ScalarTraits<double> double_traits;

	test_traits<float_traits>();
	test_traits<double_traits>();

}

BOOST_AUTO_TEST_SUITE_END();


#define BOOST_TEST_MODULE bwm_queries

#include "succinct/test_common.hpp"
#include <boost/test/floating_point_comparison.hpp>

#include "ds2i_config.hpp"
#include "index_types.hpp"
#include "queries.hpp"

namespace ds2i {
    namespace test {

        struct index_initialization {

            typedef opt_index index_type;
            typedef wand_data<bm25, wand_data_compressed<bm25, uniform_score_compressor>> WandTypeUniform;
            typedef wand_data<bm25, wand_data_raw<bm25>> WandTypePlain;

            index_initialization()
                    : collection(DS2I_SOURCE_DIR "/test/test_data/test_collection"),
                      document_sizes(DS2I_SOURCE_DIR "/test/test_data/test_collection.sizes"),
                      wdata(document_sizes.begin()->begin(), collection.num_docs(), collection, partition_type::variable_blocks),
                      wdata_fixed(document_sizes.begin()->begin(), collection.num_docs(), collection, partition_type::fixed_blocks),
                      wdata_uniform(document_sizes.begin()->begin(), collection.num_docs(), collection, partition_type::variable_blocks) {
                index_type::builder builder(collection.num_docs(), params);
                for (auto const &plist: collection) {
                    uint64_t freqs_sum = std::accumulate(plist.freqs.begin(),
                                                         plist.freqs.end(), uint64_t(0));
                    builder.add_posting_list(plist.docs.size(), plist.docs.begin(),
                                             plist.freqs.begin(), freqs_sum);
                }
                builder.build(index);

                term_id_vec q;
                std::ifstream qfile(DS2I_SOURCE_DIR "/test/test_data/queries");
                while (read_query(q, qfile)) queries.push_back(q);
            }

            global_parameters params;
            binary_freq_collection collection;
            binary_collection document_sizes;
            index_type index;
            std::vector<term_id_vec> queries;
            WandTypePlain wdata;
            WandTypePlain wdata_fixed;
            WandTypeUniform wdata_uniform;

            template<typename QueryOp>
            void test_against_wand(QueryOp &op_q) const {
                wand_query<WandTypePlain> or_q(wdata, 10);

                for (auto const &q: queries) {
                    or_q(index, q);
                    op_q(index, q);
                    BOOST_REQUIRE_EQUAL(or_q.topk().size(), op_q.topk().size());

                    for (size_t i = 0; i < or_q.topk().size(); ++i) {
                        BOOST_REQUIRE_CLOSE(or_q.topk()[i].first, op_q.topk()[i].first, 0.01); // tolerance is % relative
                    }
                    op_q.clear_topk();
                }
            }

        };

    }
}

BOOST_FIXTURE_TEST_CASE(block_max_wand,
                        ds2i::test::index_initialization) {
    ds2i::block_max_wand_query<WandTypePlain> block_max_wand_q(wdata, 10);
    ds2i::block_max_wand_query<WandTypeUniform> block_max_wand_uniform_q(wdata_uniform, 10);
    ds2i::block_max_wand_query<WandTypePlain> block_max_wand_fixed_q(wdata_fixed, 10);
    test_against_wand(block_max_wand_uniform_q);
    test_against_wand(block_max_wand_q);
    test_against_wand(block_max_wand_fixed_q);
}


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

#pragma once

/// @file
///
/// Code related to fully analytical 1D RTA solutions.

#include <constants.hpp>
#include <structures.hpp>
#include <qpoint_grid.hpp>
#include <Eigen/Dense>

namespace alma {

namespace analytic1D {

/// Class for computing basic thermal properties
/// (kappa, Cv, diffusivity) and cumulative functions
/// (resolved for MFP, energy, etc.) along a given transport direction.
/// Calculations involve appropriate thin film corrections
/// for both cross-plane and in-plane transport.

class BasicProperties_calculator {
public:
    /// Constructor: initialise internal variables
    BasicProperties_calculator(const alma::Crystal_structure* poscar,
                               const alma::Gamma_grid* grid,
                               const Eigen::ArrayXXd* w,
                               double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a logarithmic grid of MFP bins
    /// from MFPmin to MFPmax with Nbins elements
    void setLogMFPbins(double MFPmin, double MFPmax, int Nbins);

    /// Construct a logarithmic grid of MFP bins witn Nbins elements.
    /// MFPmax is assigned automatically to twice the largest MFP.
    /// MFPmin is assigned to 1e-6*largest MFP.
    void setAutoMFPbins(int Nbins);

    /// Construct a logarithmic grid of MFP bins witn Nbins elements.
    /// MFPmax is assigned automatically to the twice the largest projected MFP.
    /// MFPmin is assigned to 1e-6*largest projected MFP.
    void setAutoProjMFPbins(int Nbins);

    /// Construct a logarithmic grid of tau bins
    /// from taumin to taumax with Nbins elements
    void setLogRTbins(double taumin, double taumax, int Nbins);

    /// Construct a logarithmic grid of tau bins witn Nbins elements.
    /// taumax is assigned automatically to the twice the largest relaxation
    /// time.
    /// taumin is assigned to 1e-6*largest relaxation time.
    void setAutoRTbins(int Nbins);

    /// Construct a linear grid of omega bins
    /// from omegamin to omegamax with Nbins elements
    void setLinOmegabins(double omegamin, double omegamax, int Nbins);

    /// Construct a linear grid of omega bins witn Nbins elements.
    /// omegamax is assigned automatically to 1.1*largest angular frequency.
    /// omegamin is automatically set to 0.
    void setAutoOmegabins(int Nbins);

    /// Retrieve the MFP bins
    Eigen::VectorXd getMFPbins();

    /// Retrieve the relaxation time bins
    Eigen::VectorXd getRTbins();

    /// Retrieve the relaxation time bins
    Eigen::VectorXd getOmegabins();

    /// Specify that the medium should be treated as
    /// infinite bulk (default option)
    void setBulk();

    /// Specify that the medium should be treated as
    /// a thin film with provided thickness.
    /// This option corrects MFPs and relaxation times
    /// in conductivity calculations.
    void setInPlaneFilm(double filmthickness,
                        const Eigen::Vector3d normal,
                        double specularity = 0.0);
    void setCrossPlaneFilm(double filmthickness);

    /// Retrieve thermal conductivity
    double getConductivity();

    /// Obtain spectrally computed thermal conductivity
    double getSpectralConductivity();

    /// Compute the heat capacity and retrieve it
    double getCapacity();

    /// Compute the Fourier diffusivity and retrieve it
    double getDiffusivity();

    /// Retrieve the dominant projected MFP as defined below
    double getDominantProjMFP();

    /// Retrieve the dominant phonon relaxation time
    double getDominantRT();

    /// Compute the anisotropy index kappa_max/kappa_min and retrieve it.
    double getAnisotropyIndex();

    /// Choose resolving cumulative curves by MFP
    void resolveByMFP();

    /// Choose resolving cumulative curves by projected MFP
    void resolveByProjMFP();

    /// Choose resolving cumulative curves by relaxation time
    void resolveByRT();

    /// Choose resolving cumulative curves by angular frequency
    void resolveByOmega();

    /// Compute the cumulative thermal conductivity curve
    /// with respect to the MFP bins and retrieve it
    Eigen::VectorXd getCumulativeConductivity();

    /// Compute the cumulative heat capacity curve
    /// with respect to the MFP bins and retrieve it
    Eigen::VectorXd getCumulativeCapacity();

    /// Compute the cumulative hydrodinamic l2
    /// using RTA and assuming isotropy
    Eigen::VectorXd getCumulativel2RTAiso();

    /// Set the number of bins to be used in DOS evaluation
    void setDOSgridsize(int nbins);

    /// Obtain the DOS energy grid (in meV)
    Eigen::VectorXd getDOSgrid();

    /// Calculate and obtain the DOS
    Eigen::VectorXd getDOS();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Unitvector describing the film normal for in-plane transport
    Eigen::Vector3d normalvector;
    /// MFP bins for cumulative curves
    Eigen::VectorXd MFPbins;
    /// relaxation time bins for cumulative curves
    Eigen::VectorXd taubins;
    /// omega bins for cumulative curves
    Eigen::VectorXd omegabins;
    /// Max MFP value present in medium
    double internal_MFPmax;
    /// Max proj MFP value present in medium
    double internal_MFPprojmax;
    /// Max relaxation time present in medium
    double internal_taumax;
    /// Max phonon frequency present in medium
    double internal_omegamax;
    /// Thin film specifiers
    bool thinfilm;
    bool crossplane;
    double filmthickness;
    double specularity; // for in-plane films
    /// Basic properties
    double kappa;
    double Cv;
    /// "Dominant" projected MFP, defined as kappa-weighted average
    /// <Lambda_proj> = sum[(kappamode/kappabulk)*Lambda_proj]
    double dominantProjMFP;
    /// "Dominant" relaxation time, defined as kappa-weighted average
    /// <RT> = sum[(kappamode/kappabulk)*RT_mode]
    double dominantRT;
    /// 3D rotation matrix that rotates the cartesian coordinate system such
    /// that
    /// the transport axis becomes (1,0,0) and film normal becomes (0,0,1)
    Eigen::Matrix3d FuchsRotation;
    /// Precalculate the 3D rotation matrix used for in-plane Fuch corrections.
    void initFuchsRotation();
    /// Precalculate thermal properties
    void updateMe();
    /// Calculated cumulative conductivity function
    Eigen::VectorXd kappacumul;
    /// Calculated cumulative capacity function
    Eigen::VectorXd Cvcumul;
    /// Calculated cumulative l**2
    Eigen::VectorXd l2cumul;

    /// Specifier for kappacumul calculation
    enum kappacumulID {
        resolve_by_MFP,
        resolve_by_ProjMFP,
        resolve_by_RT,
        resolve_by_omega
    };
    int kappacumulIdentifier;
    /// Number of bins to be used for DOS evaluation
    int DOS_Nbins;
};

/// Class for computing the RTA propagator function psi(xi) of a medium.
/// The psi function fully determines the analytical single pulse
/// response of the infinite bulk in weakly quasi-ballistic regime
/// (time scales exceeding phonon relaxation times):
/// Energy density in Fourier-Laplace domain = 1/[s + psi(xi)]
/// Energy density in Fourier-time domain = exp[-psi(xi)*t]

class psi_calculator {
public:
    /// Constructor: initialise internal variables
    psi_calculator(const alma::Crystal_structure* poscar,
                   const alma::Gamma_grid* grid,
                   const Eigen::ArrayXXd* w,
                   double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a linear grid of spatial frequencies
    /// from ximin to ximax with Nxi elements
    void setLinGrid(double ximin, double ximax, int Nxi);

    /// Construct a logarithmic grid of spatial frequencies
    /// from ximin to ximax with Nxi elements
    void setLogGrid(double ximin, double ximax, int Nxi);

    /// Manually set a grid of spatial frequencies
    void setXiGrid(const Eigen::Ref<const Eigen::VectorXd> xigrid);

    /// Retrieve the spatial frequency grid
    Eigen::VectorXd getSpatialFrequencies();

    /// Determine whether computation output should be
    /// normalised by the Fourier solution Dbulk*xi^2
    void normaliseOutput(bool norm);

    /// Obtain Fourier diffusivity of the medium
    double getDiffusivity();

    /// Compute psi function and retrieve it
    Eigen::VectorXd getPsi();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Spatial frequency grid
    Eigen::VectorXd xi;
    /// Normalisation specifier
    bool scaleoutput;
    /// Fourier diffusivity
    double Dbulk;
    void updateDiffusivity();
    /// Calculated psi function
    Eigen::VectorXd psi;
};

/////////////////// SPR_calculator_FourierLaplace ///////////////////

/// Class for computing the exact analytical RTA single pulse
/// energy density response of the infinite bulk medium
/// in Fourier-Laplace domain. [see PRB 91 085202 (2015)]

class SPR_calculator_FourierLaplace {
public:
    /// Constructor: initialise internal variables
    SPR_calculator_FourierLaplace(const alma::Crystal_structure* poscar,
                                  const alma::Gamma_grid* grid,
                                  const Eigen::ArrayXXd* w,
                                  double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a linear grid of spatial frequencies
    /// from ximin to ximax with Nxi elements
    void setLinSpatialGrid(double ximin, double ximax, int Nxi);

    /// Construct a logarithmic grid of spatial frequencies
    /// from ximin to ximax with Nxi elements
    void setLogSpatialGrid(double ximin, double ximax, int Nxi);

    /// Construct a linear grid of temporal frequencies
    /// from fmin to fmax with Nf elements.
    /// From this a Laplace grid s = 2*pi*1i*f is created.
    void setLinTemporalGrid(double fmin, double fmax, int Nf);

    /// Construct a logarithmic grid of temporal frequencies
    /// from fmin to fmax with Nf elements.
    /// From this a Laplace grid s = 2*pi*1i*f is created.
    void setLogTemporalGrid(double fmin, double fmax, int Nf);

    /// Retrieve the spatial frequency grid
    Eigen::VectorXd getSpatialFrequencies();

    /// Retrieve the temporal frequency grid
    Eigen::VectorXd getTemporalFrequencies();

    /// Compute SPR and retrieve it
    Eigen::MatrixXcd getSPR();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Spatial frequency grid
    Eigen::VectorXd xi;
    /// Temporal frequency grids
    Eigen::VectorXd f;
    Eigen::VectorXcd s;
    /// Calculated SPR
    Eigen::MatrixXcd Pxis;
};

/////////////////// SPR_calculator_RealSpace ///////////////////

/// Class for computing the approximate analytical RTA single pulse
/// energy density response in real space at a given time.
///
/// The approach is only valid in weakly quasiballistic regime
/// (time scales exceeding phonon relaxation times) so that
/// P(xi,t) \approx exp[-psi(xi)*t].
///
/// Calculations are performed by evaluating the Fourier inversion
/// (1/pi)*Integral(exp[-psi(xi)*t]*cos(xi*x),xi=0..infinity)
/// semi-analytically with 2nd order Filon-type quadrature.

class SPR_calculator_RealSpace {
public:
    /// Constructor: initialise internal variables
    SPR_calculator_RealSpace(const alma::Crystal_structure* poscar,
                             const alma::Gamma_grid* grid,
                             const Eigen::ArrayXXd* w,
                             double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a linear space grid
    /// from xmin to xmax with Nx elements
    void setLinGrid(double xmin, double xmax, int Nx);

    /// Construct a logarithmic space grid
    /// from xmin to xmax with Nx elements
    void setLogGrid(double xmin, double xmax, int Nx);

    /// Declare whether the space grid is normalised.
    /// If yes, space grid values are counted in
    /// Fourier diffusion lengths, i.e.
    /// x_actual = x_grid*sqrt(2*Dbulk*t)
    void declareGridNormalised(bool norm);

    // Optional normalisation of calculation output
    // by sqrt(4*pi*Dbulk*t)
    void normaliseOutput(bool norm);

    /// Set time value to be used for calculations
    void setTime(double t);

    /// Set MFP bins to be used for resolveSPRbyMFP()
    void setLogMFPbins(double MFPmin, double MFPmax, int Nbins);

    /// Retrieve the spatial grid
    Eigen::VectorXd getGrid();
    Eigen::VectorXd getNormalisedGrid();

    /// Retrieve time value
    double getTime();

    /// Retrieve MFP bins
    Eigen::VectorXd getMFPbins();

    /// Retrieve Fourier diffusivity along thermal transport axis
    double getDiffusivity();

    /// Compute source transient P(x=0) at the provided times
    Eigen::VectorXd getSourceTransient(
        const Eigen::Ref<Eigen::VectorXd> timegrid);

    /// Compute SPR and retrieve it
    Eigen::VectorXd getSPR();

    /// Compute SPR of individual modes resolved by MFP
    Eigen::MatrixXd resolveSPRbyMFP();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Unitvector describing the film normal for inplane transport
    Eigen::Vector3d filmnormal;
    /// Spatial grid
    Eigen::VectorXd x;
    bool gridIsNormalised;
    /// Output specifier
    bool normaliseoutput;
    /// MFP bins
    Eigen::VectorXd MFPbins;
    /// time value
    double t;
    /// diffusivity of the medium along the transport axis
    double Dbulk;
    void updateDiffusivity();
    /// Calculated SPR
    Eigen::VectorXd Pxt;
    /// Calculated SPR resolved by MFP
    Eigen::MatrixXd Pmodes;
};

/////////////////// MSD_calculator_Laplace ///////////////////

/// Class for computing the exact analytical RTA solution for
/// mean square thermal energy displacement. By definition,
/// the MSD is the variance of the macroscopic energy density:
/// MSD = Integral(x^2*P(x,s),x=-infinity..infinity).
/// Results are valid across all regimes
/// (from fully ballistic to fully diffusive transport).
/// Calculations are performed fully analytically by using
/// the moment generating properties of the energy density
/// in spatial frequency domain:
/// MSD = minus second derivative of P(xi,s) at xi = 0.

class MSD_calculator_Laplace {
public:
    /// Constructor: initialise internal variables
    MSD_calculator_Laplace(const alma::Crystal_structure* poscar,
                           const alma::Gamma_grid* grid,
                           const Eigen::ArrayXXd* w,
                           double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a linear temporal frequency grid
    /// from fmin to fmax with Nf elements
    void setLinGrid(double fmin, double fmax, int Nf);

    /// Construct a logarithmic temporal frequency grid
    /// from fmin to fmax with Nf elements
    void setLogGrid(double fmin, double fmax, int Nf);

    /// Manually set a Laplace grid
    void setLaplaceGrid(const Eigen::Ref<const Eigen::VectorXcd> sgrid);

    /// Retrieve the temporal frequency grid
    Eigen::VectorXd getGrid();
    Eigen::VectorXcd getLaplaceGrid();

    /// Compute MSD and retrieve it
    Eigen::VectorXcd getMSD();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Temporal frequency grid
    Eigen::VectorXd f;
    /// Laplace grid
    Eigen::VectorXcd s;
    /// Calculated MSD function
    Eigen::VectorXcd MSD;
};

/////////////////// MSD_calculator_RealTime ///////////////////

/// Class for computing the exact analytical RTA solution for
/// mean square thermal energy displacement in time domain.
///
/// Calculations are performed by numerically inverting
/// output from MSD_calculator_Laplace to time domain
/// with a Gaver-Stehfest scheme.

class MSD_calculator_RealTime {
public:
    /// Constructor: initialise internal variables
    MSD_calculator_RealTime(const alma::Crystal_structure* poscar,
                            const alma::Gamma_grid* grid,
                            const Eigen::ArrayXXd* w,
                            double T);

    /// Function setting the thermal transport axis
    void setDirection(const Eigen::Vector3d unitvector);

    /// Construct a linear time grid
    /// from tmin to tmax with Nt elements
    void setLinGrid(double tmin, double tmax, int Nt);

    /// Construct a logarithmic time grid
    /// from tmin to tmax with Nt elements
    void setLogGrid(double tmin, double tmax, int Nt);

    /// Manually set a time grid
    void setTimeGrid(const Eigen::Ref<const Eigen::VectorXd> tgrid);

    /// Retrieve the timegrid
    Eigen::VectorXd getGrid();

    /// Optional choice to normalise calculated results
    /// by Fourier solution 2*D*t.
    void normaliseOutput(bool norm);

    /// Retrieve Fourier diffusivity
    double getDiffusivity();

    /// Compute MSD and retrieve it
    Eigen::VectorXd getMSD();

private:
    /// Pointer to description of the unit cell
    const alma::Crystal_structure* poscar;
    /// Pointer to phonon spectrum on a regular q-point grid
    const alma::Gamma_grid* grid;
    /// Pointer to scattering rates for all modes irreducible points in the grid
    const Eigen::ArrayXXd* w;
    /// Temperature
    double T;

    /// Unitvector describing 1D thermal transport axis
    Eigen::Vector3d unitvector;
    /// Unitvector describing the film normal for inplane transport
    Eigen::Vector3d filmnormal;
    /// Time grid
    Eigen::VectorXd t;
    /// Fourier diffusivity
    double Dbulk;
    void updateDiffusivity();
    /// output specifier
    bool normaliseoutput;
    /// Calculated MSD function
    Eigen::VectorXd MSD;

    /// Gaver-Stehfest Laplace inversion coefficients
    int GS_depth;
    Eigen::VectorXd GS_coeffs;
    /// Helper function for Gaver-Stehfest Laplace inversion
    double factorial(int n);
};
} // end namespace analytic1D
} // end namespace alma


#if !defined(PAWN_MEXPR_HPP)
#define PAWN_MEXPR_HPP

#define BOOST_SPIRIT_NO_PREDEFINED_TERMINALS

#include <boost/spirit/include/qi.hpp>
#include "mast.hpp"
#include "error_handler.hpp"
#include <vector>

namespace client { namespace math { namespace parser
{
    namespace qi = boost::spirit::qi;
    namespace ascii = boost::spirit::ascii;

    ///////////////////////////////////////////////////////////////////////////////
    //  The math expression grammar
    ///////////////////////////////////////////////////////////////////////////////
    template <typename Iterator>
    struct expression : qi::grammar<Iterator, ast::expr(), ascii::space_type>
    {
        expression(error_handler<Iterator>& error_handler);

        qi::rule<Iterator, ast::expr(), ascii::space_type>
          expr, additive_expr, multiplicative_expr;
        qi::rule<Iterator, ast::operand(), ascii::space_type>
          unary_expr, primary_expr;
        qi::rule<Iterator, std::string(), ascii::space_type> identifier;
        qi::rule<Iterator, unsigned int, ascii::space_type> colIndex;
        qi::symbols<char, ast::optoken> additive_op, multiplicative_op, unary_op;
        qi::symbols<char>
            keywords
            ;
    };
}}}

#endif