#include "mtf/AM/SSDBase.h"
#include "mtf/Utilities/miscUtils.h"
#include "mtf/Utilities/spiUtils.h"

#include <boost/random/random_device.hpp>
#include <boost/random/seed_seq.hpp>

_MTF_BEGIN_NAMESPACE

SSDBase::SSDBase(const AMParams *am_params, const int _n_channels) :
AppearanceModel(am_params, _n_channels), ilm(nullptr), likelihood_alpha(1),
I_diff(0, 0){
	if(am_params){
		ilm = am_params->ilm;
		likelihood_alpha = am_params->likelihood_alpha;
	}
	if(ilm){
		ilm->setPixHessType(ILMPixHessT::Constant);
		ilm_d2f_dIt_type = ilm->getPixHessType();
		state_size = ilm->getStateSize();
		if(ilm_d2f_dIt_type == ILMPixHessT::General){
			d2f_dIt2.resize(n_pix, n_pix);
		}
	}
	pix_norm_mult = 1;
	pix_norm_add = 0;
}

void SSDBase::initializeSimilarity(){
	if(is_initialized.similarity)
		return;

	df_dI0.resize(patch_size);
	new (&I_diff) VectorXdM(df_dI0.data(), patch_size);
	I_diff.fill(0);
	f = 0;
	if(ilm){
		It_orig.resize(patch_size);
		It_orig = It;
		p_am.resize(state_size);
		ilm->initialize(p_am.data());
	}
	is_initialized.similarity = true;
}

void SSDBase::initializeGrad(){
	if(is_initialized.grad)
		return;

	df_dIt.resize(patch_size);
	df_dIt = df_dI0;

	if(ilm){
		df_dpam.resize(state_size);
		df_dgt.resize(patch_size);
		df_dg0.resize(patch_size);
		df_dg0 = df_dI0;
		df_dgt = df_dIt;
		ilm->cmptParamJacobian(df_dpam.data(), df_dgt.data(), It_orig.data(), p_am.data());
		ilm->cmptPixJacobian(df_dIt.data(), df_dgt.data(), It_orig.data(), p_am.data());
	}
	is_initialized.grad = true;
}
double SSDBase::getLikelihood() const{
	/**
	since SSD can be numerically too large for exponential function to work,
	we take the square root of the per pixel SSD instead;
	it is further normalized by dividing with the range of pixel values to avoid
	very small numbers that may lead to loss of information due to the limited precision
	of floating point operations;
	*/
	return exp(-likelihood_alpha * sqrt(-f / (static_cast<double>(patch_size))));
}

void SSDBase::updateSimilarity(bool prereq_only){
	if(ilm){
		//! I_t actually refers to g(I_t, p_am) to keep the illumination model transparent
		//! and avoid any unnecessary runtime costs if it is not used
		It_orig = It;
		ilm->apply(It.data(), It_orig.data(), p_am.data());
	}
	I_diff = It - I0;
	if(prereq_only){ return; }
#ifndef DISABLE_SPI
	if(spi_mask){
		VectorXd masked_err_vec = Map<const VectorXb>(spi_mask, n_pix).select(I_diff, 0);
		f = -masked_err_vec.squaredNorm() / 2;
	} else{
#endif
		f = -I_diff.squaredNorm() / 2;
#ifndef DISABLE_SPI
	}
#endif
	//utils::printMatrix(p_am.transpose(), "p_am");
	//printf("f: %f\n", f);
}
void SSDBase::updateState(const VectorXd& state_update){
	if(ilm){
		assert(state_update.size() == state_size);
		VectorXd p_am_old = p_am;
		ilm->update(p_am.data(), p_am_old.data(), state_update.data());
		//utils::printMatrix(state_update.transpose(), "state_update");
		//utils::printMatrix(p_am_old.transpose(), "p_am before update");
		//utils::printMatrix(p_am.transpose(), "p_am after update");
	}
}
void SSDBase::invertState(VectorXd& inv_p, const VectorXd& p){
	if(ilm){
		assert(inv_p.size() == state_size && p.size() == state_size);
		ilm->invert(inv_p.data(), p.data());
	}
}

// curr_grad is same as negative of init_grad
void SSDBase::updateCurrGrad(){
	df_dIt = -df_dI0;
	if(ilm){
		df_dgt = df_dIt;
		ilm->cmptPixJacobian(df_dIt.data(), df_dgt.data(), It_orig.data(), p_am.data());
	}
}

void SSDBase::cmptInitJacobian(RowVectorXd &df_dp,
	const MatrixXd &dI0_dpssm){
	assert(df_dp.size() == state_size + dI0_dpssm.cols());
	assert(dI0_dpssm.rows() == patch_size);
	if(ilm){
		df_dp.head(dI0_dpssm.cols()).noalias() = df_dI0 * dI0_dpssm;
		ilm->cmptParamJacobian(df_dpam.data(), df_dg0.data(),
			I0.data(), p_am.data());
		df_dp.tail(state_size) = df_dpam;
	} else{
#ifndef DISABLE_SPI
		if(spi_mask){
			getJacobian(df_dp, spi_mask, df_dI0, dI0_dpssm);
		} else{
#endif
			df_dp.noalias() = df_dI0 * dI0_dpssm;
#ifndef DISABLE_SPI
		}
#endif
	}
}
void SSDBase::cmptCurrJacobian(RowVectorXd &df_dp,
	const MatrixXd &dIt_dpssm){
	assert(df_dp.size() == state_size + dIt_dpssm.cols());
	assert(dIt_dpssm.rows() == patch_size);

	if(ilm){
		df_dp.head(dIt_dpssm.cols()).noalias() = df_dIt * dIt_dpssm;
		ilm->cmptParamJacobian(df_dpam.data(), df_dgt.data(),
			It_orig.data(), p_am.data());
		df_dp.tail(state_size) = df_dpam;
	} else{
#ifndef DISABLE_SPI
		if(spi_mask){
			getJacobian(df_dp, spi_mask, df_dIt, dIt_dpssm);
		} else{
#endif
			//printf("df_dp: %ld x %ld\n", df_dp.rows(), df_dp.cols());
			//printf("df_dIt: %ld x %ld\n", df_dIt.rows(), df_dIt.cols());
			//printf("dIt_dpssm: %ld x %ld\n", dIt_dpssm.rows(), dIt_dpssm.cols());
			df_dp.noalias() = df_dIt * dIt_dpssm;
#ifndef DISABLE_SPI
		}
#endif
	}
}
void SSDBase::cmptDifferenceOfJacobians(RowVectorXd &df_dp_diff,
	const MatrixXd &dI0_dpssm, const MatrixXd &dIt_dpssm){
	assert(df_dp_diff.size() == state_size + dIt_dpssm.cols());
	assert(dI0_dpssm.cols() == dIt_dpssm.cols());
	assert(dIt_dpssm.rows() == patch_size && dI0_dpssm.rows() == patch_size);
	if(ilm){
		df_dp_diff.head(dIt_dpssm.cols()).noalias() = df_dIt * (dI0_dpssm + dIt_dpssm);
		VectorXd df_dpam_0(state_size), df_dpam_t(state_size);
		ilm->cmptParamJacobian(df_dpam_0.data(), df_dg0.data(), I0.data(), p_am.data());
		ilm->cmptParamJacobian(df_dpam_t.data(), df_dgt.data(), It_orig.data(), p_am.data());
		df_dp_diff.tail(state_size) = df_dpam_t - df_dpam_0;
	} else{
#ifndef DISABLE_SPI
		if(spi_mask){
			getDifferenceOfJacobians(df_dp_diff, spi_mask, dI0_dpssm, dIt_dpssm);
		} else{
#endif
			df_dp_diff.noalias() = df_dIt * (dI0_dpssm + dIt_dpssm);
#ifndef DISABLE_SPI
		}
#endif
	}
}

void SSDBase::cmptILMHessian(MatrixXd &d2f_dp2, const MatrixXd &dI_dpssm,
	const double* I, const double* df_dg){
	int ssm_state_size = static_cast<int>(dI_dpssm.cols());

	assert(d2f_dp2.rows() == state_size + ssm_state_size);
	assert(d2f_dp2.cols() == state_size + ssm_state_size);

	MatrixXd d2f_dpam2(state_size, state_size);
	ilm->cmptParamHessian(d2f_dpam2.data(), nullptr, df_dg, I, p_am.data());
	d2f_dp2.bottomRightCorner(state_size, state_size) = -d2f_dpam2;
	switch(ilm_d2f_dIt_type){
	case  ILMPixHessT::Constant:
	{
		double d2f_dIt2_const;
		if(df_dg){
			ilm->cmptPixHessian(&d2f_dIt2_const, nullptr, df_dg, I, p_am.data());
		} else{
			ilm->cmptPixHessian(&d2f_dIt2_const, nullptr, I, p_am.data());
		}

		d2f_dp2.topLeftCorner(ssm_state_size, ssm_state_size).noalias() = -d2f_dIt2_const * dI_dpssm.transpose() * dI_dpssm;
		break;
	}
	case  ILMPixHessT::Diagonal:
	{
		VectorXd d2f_dIt2_diag(patch_size);
		if(df_dg){
			ilm->cmptPixHessian(d2f_dIt2_diag.data(), nullptr, df_dg, I, p_am.data());
		} else{
			ilm->cmptPixHessian(d2f_dIt2_diag.data(), nullptr, I, p_am.data());
		}
		d2f_dp2.topLeftCorner(ssm_state_size, ssm_state_size).noalias() =
			-(dI_dpssm.array().colwise() * d2f_dIt2_diag.array()).matrix().transpose() * dI_dpssm;
		break;
	}
	case  ILMPixHessT::General:
	{
		if(df_dg){
			ilm->cmptPixHessian(d2f_dIt2.data(), nullptr, df_dg, I, p_am.data());
		} else{
			ilm->cmptPixHessian(d2f_dIt2.data(), nullptr, I, p_am.data());
		}
		d2f_dp2.topLeftCorner(ssm_state_size, ssm_state_size).noalias() = -dI_dpssm.transpose() * d2f_dIt2 * dI_dpssm;
		break;
	}
	default:
		throw utils::InvalidArgument(
			cv::format("SSDBase :: ILM has invalid hessian type provided: %d", ilm_d2f_dIt_type));
	}
	if(df_dg){
		ilm->cmptCrossHessian(d2f_dpam_dIt.data(), nullptr, df_dg, I, p_am.data());
	} else{
		ilm->cmptCrossHessian(d2f_dpam_dIt.data(), nullptr, I, p_am.data());
	}
	d2f_dp2.topRightCorner(ssm_state_size, state_size).transpose() =
		d2f_dp2.bottomLeftCorner(state_size, ssm_state_size) = -d2f_dpam_dIt*dI_dpssm;
}

void SSDBase::cmptInitHessian(MatrixXd &d2f_dp2, const MatrixXd &dI0_dpssm){
	assert(d2f_dp2.rows() == dI0_dpssm.cols() + state_size && d2f_dp2.rows() == d2f_dp2.cols());
	assert(dI0_dpssm.rows() == patch_size);
	if(ilm){
		cmptILMHessian(d2f_dp2, dI0_dpssm, I0.data());
	} else{
#ifndef DISABLE_SPI
		if(spi_mask){
			getHessian(d2f_dp2, spi_mask, dI0_dpssm);
		} else{
#endif
			d2f_dp2.noalias() = -dI0_dpssm.transpose() * dI0_dpssm;
#ifndef DISABLE_SPI
		}
#endif
	}
}
void SSDBase::cmptCurrHessian(MatrixXd &d2f_dp2,
	const MatrixXd &dIt_dpssm){
	assert(d2f_dp2.rows() == dIt_dpssm.cols() + state_size && d2f_dp2.rows() == d2f_dp2.cols());
	assert(dIt_dpssm.rows() == patch_size);
	if(ilm){
		cmptILMHessian(d2f_dp2, dIt_dpssm, It_orig.data());
	} else{
#ifndef DISABLE_SPI
		if(spi_mask){
			getHessian(d2f_dp2, spi_mask, dIt_dpssm);
		} else{
#endif
			d2f_dp2.noalias() = -dIt_dpssm.transpose() * dIt_dpssm;
#ifndef DISABLE_SPI
		}
#endif
	}
}
//! analogous to cmptDifferenceOfJacobians except for computing the difference between the current and initial Hessians
void SSDBase::cmptSumOfHessians(MatrixXd &d2f_dp2_sum,
	const MatrixXd &dI0_dpssm,	const MatrixXd &dIt_dpssm){
	assert(d2f_dp2_sum.cols() == d2f_dp2_sum.rows());
	assert(dI0_dpssm.rows() == patch_size && dIt_dpssm.rows() == patch_size);

	if(ilm){
		MatrixXd d2f_dp2t(d2f_dp2_sum.rows(), d2f_dp2_sum.cols());
		cmptILMHessian(d2f_dp2t, dIt_dpssm, It_orig.data());
		MatrixXd d2f_dp20(d2f_dp2_sum.rows(), d2f_dp2_sum.cols());
		cmptILMHessian(d2f_dp20, dI0_dpssm, I0.data());
		d2f_dp2_sum.noalias() = d2f_dp2t + d2f_dp20;
	} else{
#ifndef DISABLE_SPI
		if(spi_mask){
			getSumOfHessians(d2f_dp2_sum, spi_mask,
				dI0_dpssm, dIt_dpssm);
		} else{
#endif
			d2f_dp2_sum.noalias() = -(dI0_dpssm.transpose() * dI0_dpssm
				+ dIt_dpssm.transpose() * dIt_dpssm);
#ifndef DISABLE_SPI
		}
#endif
	}
}

void SSDBase::cmptInitHessian(MatrixXd &d2f_dp2, const MatrixXd &dI0_dpssm,
	const MatrixXd &d2I0_dpssm2){
	int ssm_state_size = static_cast<int>(d2f_dp2.rows());
	assert(d2f_dp2.cols() == ssm_state_size + state_size);
	assert(d2I0_dpssm2.rows() == ssm_state_size * ssm_state_size && d2I0_dpssm2.cols() == n_channels*n_pix);

	assert(d2f_dp2.rows() == dI0_dpssm.cols() + state_size && d2f_dp2.rows() == d2f_dp2.cols());
	assert(dI0_dpssm.rows() == patch_size);
	if(ilm){
		cmptILMHessian(d2f_dp2, dI0_dpssm, I0.data(), df_dg0.data());
	} else{
#ifndef DISABLE_SPI
		if(spi_mask){
			getHessian(d2f_dp2, spi_mask, dI0_dpssm);
		} else{
#endif
			d2f_dp2.noalias() = -dI0_dpssm.transpose() * dI0_dpssm;
#ifndef DISABLE_SPI
		}
#endif
	}
	int ch_pix_id = 0;
	for(unsigned int pix_id = 0; pix_id < n_pix; ++pix_id){
		spi_pt_check_mc(spi_mask, pix_id, ch_pix_id);
		for(unsigned int channel_id = 0; channel_id < n_channels; ++channel_id){
			d2f_dp2 += Map<const MatrixXd>(d2I0_dpssm2.col(ch_pix_id).data(), ssm_state_size, ssm_state_size) * df_dI0(ch_pix_id);
			++ch_pix_id;
		}
	}
}

void SSDBase::cmptCurrHessian(MatrixXd &d2f_dp2, const MatrixXd &dIt_dpssm,
	const MatrixXd &d2It_dpssm2){
	int ssm_state_size = static_cast<int>(d2f_dp2.rows());

	assert(d2f_dp2.cols() == ssm_state_size + state_size);
	assert(d2It_dpssm2.rows() == ssm_state_size * ssm_state_size && d2It_dpssm2.cols() == n_channels * n_pix);

	assert(d2f_dp2.rows() == dIt_dpssm.cols() + state_size && d2f_dp2.rows() == d2f_dp2.cols());
	assert(dIt_dpssm.rows() == patch_size);
	if(ilm){
		cmptILMHessian(d2f_dp2, dIt_dpssm, It_orig.data(), df_dgt.data());
	} else{
#ifndef DISABLE_SPI
		if(spi_mask){
			getHessian(d2f_dp2, spi_mask, dIt_dpssm);
		} else{
#endif
			d2f_dp2.noalias() = -dIt_dpssm.transpose() * dIt_dpssm;
#ifndef DISABLE_SPI
		}
#endif
	}
	int ch_pix_id = 0;
	for(unsigned int pix_id = 0; pix_id < n_pix; ++pix_id){
		spi_pt_check_mc(spi_mask, pix_id, ch_pix_id);
		for(unsigned int channel_id = 0; channel_id < n_channels; ++channel_id){
			d2f_dp2 += Map<const MatrixXd>(d2It_dpssm2.col(ch_pix_id).data(), ssm_state_size, ssm_state_size) * df_dIt(ch_pix_id);
			++ch_pix_id;
		}
	}
}

void SSDBase::cmptSumOfHessians(MatrixXd &d2f_dp2_sum,
	const MatrixXd &dI0_dpssm, const MatrixXd &dIt_dpssm,
	const MatrixXd &d2I0_dpssm2, const MatrixXd &d2It_dpssm2){

	int ssm_state_size = static_cast<int>(d2f_dp2_sum.rows());
	assert(d2f_dp2_sum.cols() == ssm_state_size);
	assert(d2I0_dpssm2.rows() == ssm_state_size * ssm_state_size && d2I0_dpssm2.cols() == n_channels * n_pix);
	assert(d2It_dpssm2.rows() == ssm_state_size * ssm_state_size && d2It_dpssm2.cols() == n_channels * n_pix);

	if(ilm){
		MatrixXd d2f_dp2t(d2f_dp2_sum.rows(), d2f_dp2_sum.cols());
		cmptILMHessian(d2f_dp2t, dIt_dpssm, It_orig.data(), df_dgt.data());
		MatrixXd d2f_dp20(d2f_dp2_sum.rows(), d2f_dp2_sum.cols());
		cmptILMHessian(d2f_dp20, dI0_dpssm, I0.data(), df_dgt.data());
		d2f_dp2_sum.noalias() = d2f_dp2t + d2f_dp20;
	} else{
#ifndef DISABLE_SPI
		if(spi_mask){
			getSumOfHessians(d2f_dp2_sum, spi_mask,
				dI0_dpssm, dIt_dpssm);
		} else{
#endif
			d2f_dp2_sum.noalias() = -(dI0_dpssm.transpose() * dI0_dpssm
				+ dIt_dpssm.transpose() * dIt_dpssm);
#ifndef DISABLE_SPI
		}
#endif
	}

	int ch_pix_id = 0;
	for(unsigned int pix_id = 0; pix_id < n_pix; ++pix_id){
		spi_pt_check_mc(spi_mask, pix_id, ch_pix_id);
		for(unsigned int channel_id = 0; channel_id < n_channels; ++channel_id){
			d2f_dp2_sum += df_dI0(pix_id)*
				(Map<const MatrixXd>(d2I0_dpssm2.col(ch_pix_id).data(), ssm_state_size, ssm_state_size)
				+ Map<const MatrixXd>(d2It_dpssm2.col(ch_pix_id).data(), ssm_state_size, ssm_state_size));
			++ch_pix_id;
		}
	}
}

void SSDBase::getJacobian(RowVectorXd &jacobian, const bool *pix_mask,
	const RowVectorXd &df_dI, const MatrixXd &dI_dpssm){
	assert(dI_dpssm.rows() == patch_size && dI_dpssm.rows() == jacobian.size());
	jacobian.setZero();
	int ch_pix_id = 0;
	for(unsigned int pix_id = 0; pix_id < n_pix; ++pix_id){
		spi_check_mc(pix_mask, pix_id, ch_pix_id);
		for(unsigned int channel_id = 0; channel_id < n_channels; ++channel_id){
			jacobian += df_dI[ch_pix_id] * dI_dpssm.row(ch_pix_id);
			++ch_pix_id;
		}

	}
}

void SSDBase::getDifferenceOfJacobians(RowVectorXd &diff_of_jacobians, const bool *pix_mask,
	const MatrixXd &dI0_dpssm, const MatrixXd &dIt_dpssm){
	assert(dI0_dpssm.rows() == n_channels * n_pix && dIt_dpssm.rows() == n_channels * n_pix);
	assert(dI0_dpssm.rows() == diff_of_jacobians.size());

	diff_of_jacobians.setZero();
	int ch_pix_id = 0;
	for(unsigned int pix_id = 0; pix_id < n_pix; ++pix_id){
		spi_check_mc(pix_mask, pix_id, ch_pix_id);
		for(unsigned int channel_id = 0; channel_id < n_channels; ++channel_id){
			diff_of_jacobians += df_dIt[ch_pix_id] *
				(dI0_dpssm.row(ch_pix_id) + dIt_dpssm.row(ch_pix_id));
			++ch_pix_id;
		}

	}
}

void SSDBase::getHessian(MatrixXd &d2f_dp2, const bool *pix_mask, const MatrixXd &dI_dpssm){
	assert(dI_dpssm.rows() == n_channels * n_pix);

	d2f_dp2.setZero();
	int ch_pix_id = 0;
	for(unsigned int pix_id = 0; pix_id < n_pix; ++pix_id){
		spi_check_mc(pix_mask, pix_id, ch_pix_id);
		for(unsigned int channel_id = 0; channel_id < n_channels; ++channel_id){
			d2f_dp2 -= dI_dpssm.row(ch_pix_id).transpose()*dI_dpssm.row(ch_pix_id);
			++ch_pix_id;
		}
	}
}

void SSDBase::getSumOfHessians(MatrixXd &d2f_dp2, const bool *pix_mask,
	const MatrixXd &dI0_dpssm, const MatrixXd &dIt_dpssm){
	assert(dI0_dpssm.rows() == n_channels * n_pix && dIt_dpssm.rows() == n_channels * n_pix);
	assert(d2f_dp2.rows() == d2f_dp2.cols() && d2f_dp2.rows() == dI0_dpssm.cols());

	d2f_dp2.setZero();
	unsigned int ch_pix_id = 0;
	for(unsigned int pix_id = 0; pix_id < n_pix; ++pix_id){
		spi_check_mc(pix_mask, pix_id, ch_pix_id);
		for(unsigned int channel_id = 0; channel_id < n_channels; ++channel_id){
			d2f_dp2 -= dI0_dpssm.row(ch_pix_id).transpose()*dI0_dpssm.row(ch_pix_id)
				+ dIt_dpssm.row(ch_pix_id).transpose()*dIt_dpssm.row(ch_pix_id);
			++ch_pix_id;
		}
	}
}

// -------------------------------------------------------------------------- //
// --------------------------- Stochastic Sampler --------------------------- //
// -------------------------------------------------------------------------- //

void SSDBase::initializeSampler(const VectorXd &_state_sigma,
	const VectorXd &_state_mean){
	if(!ilm){ return; }

	VectorXd state_sigma(state_size), state_mean(state_size);
	ilm->parseSamplerSigma(state_sigma, _state_sigma);
	ilm->parseSamplerMean(state_mean, _state_mean);

	printf("Initializing %s sampler with sigma: ",
		ilm ? (name + "/" + ilm->name).c_str() : name.c_str());
	utils::printMatrix(state_sigma.transpose(), nullptr, "%e");

	state_perturbation.resize(state_size);
	rand_gen.resize(state_size);
	rand_dist.resize(state_size);

	boost::random_device r;
	for(int state_id = 0; state_id < state_size; state_id++) {
		boost::random::seed_seq seed{ r(), r(), r(), r(), r(), r(), r(), r() };
		rand_gen[state_id] = SampleGenT(seed);
		rand_dist[state_id] = SampleDistT(state_mean[state_id], state_sigma[state_id]);
	}
	is_initialized.sampler = true;
}

void SSDBase::setSampler(const VectorXd &_state_sigma,
	const VectorXd &_state_mean){
	if(!ilm){ return; }
	VectorXd state_sigma(state_size), state_mean(state_size);
	ilm->parseSamplerSigma(state_sigma, _state_sigma);
	ilm->parseSamplerMean(state_mean, _state_mean);
	for(int state_id = 0; state_id < state_size; state_id++){
		rand_dist[state_id].param(DistParamT(state_mean[state_id], state_sigma[state_id]));
	}
}

void SSDBase::setSamplerMean(const VectorXd &_state_mean){
	if(!ilm){ return; }
	VectorXd state_mean(state_size);
	ilm->parseSamplerMean(state_mean, _state_mean);
	for(int state_id = 0; state_id < state_size; state_id++){
		double state_sigma = rand_dist[state_id].sigma();
		rand_dist[state_id].param(DistParamT(state_mean[state_id], state_sigma));
	}
}
void SSDBase::setSamplerSigma(const VectorXd &_state_sigma){
	VectorXd state_sigma(state_size);
	ilm->parseSamplerSigma(state_sigma, _state_sigma);
	for(int state_id = 0; state_id < state_size; state_id++){
		double mean = rand_dist[state_id].mean();
		rand_dist[state_id].param(DistParamT(mean, state_sigma[state_id]));
	}
}

void SSDBase::getSamplerSigma(VectorXd &std){
	if(!ilm){ return; }
	assert(std.size() == state_size);
	for(int state_id = 0; state_id < state_size; state_id++){
		std(state_id) = rand_dist[state_id].sigma();
	}
}
void SSDBase::getSamplerMean(VectorXd &mean){
	if(!ilm){ return; }
	assert(mean.size() == state_size);
	for(int state_id = 0; state_id < state_size; state_id++){
		mean(state_id) = rand_dist[state_id].mean();
	}
}

void SSDBase::generatePerturbation(VectorXd &perturbation){
	if(!ilm){ return; }
	assert(perturbation.size() == state_size);
	for(int state_id = 0; state_id < state_size; state_id++){
		perturbation(state_id) = rand_dist[state_id](rand_gen[state_id]);
	}
}

/**
* Squared Euclidean distance functor, optimized version
*/
/**
*  Compute the squared Euclidean distance between two vectors.
*
*	This is highly optimized, with loop unrolling, as it is one
*	of the most expensive inner loops.
*
*	The computation of squared root at the end is omitted for
*	efficiency.
*/
double SSDBaseDist::operator()(const double* a, const double* b,
	size_t size, double worst_dist) const{
	double result = 0;
	double diff0, diff1, diff2, diff3;
	const double* last = a + size;
	const double* lastgroup = last - 3;

	/* Process 4 items with each loop for efficiency. */
	while(a < lastgroup){
		diff0 = (a[0] - b[0]);
		diff1 = (a[1] - b[1]);
		diff2 = (a[2] - b[2]);
		diff3 = (a[3] - b[3]);
		result += diff0 * diff0 + diff1 * diff1 + diff2 * diff2 + diff3 * diff3;
		a += 4;
		b += 4;

		if((worst_dist > 0) && (result > worst_dist)){
			return result;
		}
	}
	/* Process last 0-3 pixels.  Not needed for standard vector lengths. */
	while(a < last){
		diff0 = (*a++ - *b++);
		result += diff0 * diff0;
	}
	return result;
}

_MTF_END_NAMESPACE


//  (C) Copyright John Maddock 2006.
//  Use, modification and distribution are subject to 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)

//
// This is not a complete header file, it is included by gamma.hpp
// after it has defined it's definitions.  This inverts the incomplete
// gamma functions P and Q on the first parameter "a" using a generic
// root finding algorithm (TOMS Algorithm 748).
//

#ifndef BOOST_MATH_SP_DETAIL_GAMMA_INVA
#define BOOST_MATH_SP_DETAIL_GAMMA_INVA

#ifdef _MSC_VER
#pragma once
#endif

#include <boost/math/tools/toms748_solve.hpp>
#include <boost/cstdint.hpp>

namespace boost{ namespace math{ namespace detail{

template <class T, class Policy>
struct gamma_inva_t
{
   gamma_inva_t(T z_, T p_, bool invert_) : z(z_), p(p_), invert(invert_) {}
   T operator()(T a)
   {
      return invert ? p - boost::math::gamma_q(a, z, Policy()) : boost::math::gamma_p(a, z, Policy()) - p;
   }
private:
   T z, p;
   bool invert;
};

template <class T, class Policy>
T inverse_poisson_cornish_fisher(T lambda, T p, T q, const Policy& pol)
{
   BOOST_MATH_STD_USING
   // mean:
   T m = lambda;
   // standard deviation:
   T sigma = sqrt(lambda);
   // skewness
   T sk = 1 / sigma;
   // kurtosis:
   // T k = 1/lambda;
   // Get the inverse of a std normal distribution:
   T x = boost::math::erfc_inv(p > q ? 2 * q : 2 * p, pol) * constants::root_two<T>();
   // Set the sign:
   if(p < 0.5)
      x = -x;
   T x2 = x * x;
   // w is correction term due to skewness
   T w = x + sk * (x2 - 1) / 6;
   /*
   // Add on correction due to kurtosis.
   // Disabled for now, seems to make things worse?
   //
   if(lambda >= 10)
      w += k * x * (x2 - 3) / 24 + sk * sk * x * (2 * x2 - 5) / -36;
   */
   w = m + sigma * w;
   return w > tools::min_value<T>() ? w : tools::min_value<T>();
}

template <class T, class Policy>
T gamma_inva_imp(const T& z, const T& p, const T& q, const Policy& pol)
{
   BOOST_MATH_STD_USING  // for ADL of std lib math functions
   //
   // Special cases first:
   //
   if(p == 0)
   {
      return tools::max_value<T>();
   }
   if(q == 0)
   {
      return tools::min_value<T>();
   }
   //
   // Function object, this is the functor whose root
   // we have to solve:
   //
   gamma_inva_t<T, Policy> f(z, (p < q) ? p : q, (p < q) ? false : true);
   //
   // Tolerance: full precision.
   //
   tools::eps_tolerance<T> tol(policies::digits<T, Policy>());
   //
   // Now figure out a starting guess for what a may be,
   // we'll start out with a value that'll put p or q
   // right bang in the middle of their range, the functions
   // are quite sensitive so we should need too many steps
   // to bracket the root from there:
   //
   T guess;
   T factor = 8;
   if(z >= 1)
   {
      //
      // We can use the relationship between the incomplete
      // gamma function and the poisson distribution to
      // calculate an approximate inverse, for large z
      // this is actually pretty accurate, but it fails badly
      // when z is very small.  Also set our step-factor according
      // to how accurate we think the result is likely to be:
      //
      guess = 1 + inverse_poisson_cornish_fisher(z, q, p, pol);
      if(z > 5)
      {
         if(z > 1000)
            factor = 1.01f;
         else if(z > 50)
            factor = 1.1f;
         else if(guess > 10)
            factor = 1.25f;
         else
            factor = 2;
         if(guess < 1.1)
            factor = 8;
      }
   }
   else if(z > 0.5)
   {
      guess = z * 1.2f;
   }
   else
   {
      guess = -0.4f / log(z);
   }
   //
   // Max iterations permitted:
   //
   boost::uintmax_t max_iter = policies::get_max_root_iterations<Policy>();
   //
   // Use our generic derivative-free root finding procedure.
   // We could use Newton steps here, taking the PDF of the
   // Poisson distribution as our derivative, but that's
   // even worse performance-wise than the generic method :-(
   //
   std::pair<T, T> r = bracket_and_solve_root(f, guess, factor, false, tol, max_iter, pol);
   if(max_iter >= policies::get_max_root_iterations<Policy>())
      policies::raise_evaluation_error<T>("boost::math::gamma_p_inva<%1%>(%1%, %1%)", "Unable to locate the root within a reasonable number of iterations, closest approximation so far was %1%", r.first, pol);
   return (r.first + r.second) / 2;
}

} // namespace detail

template <class T1, class T2, class Policy>
inline typename tools::promote_args<T1, T2>::type
   gamma_p_inva(T1 x, T2 p, const Policy& pol)
{
   typedef typename tools::promote_args<T1, T2>::type result_type;
   typedef typename policies::evaluation<result_type, Policy>::type value_type;
   typedef typename policies::normalise<
      Policy,
      policies::promote_float<false>,
      policies::promote_double<false>,
      policies::discrete_quantile<>,
      policies::assert_undefined<> >::type forwarding_policy;

   if(p == 0)
   {
      return tools::max_value<result_type>();
   }
   if(p == 1)
   {
      return tools::min_value<result_type>();
   }

   return policies::checked_narrowing_cast<result_type, forwarding_policy>(
      detail::gamma_inva_imp(
         static_cast<value_type>(x),
         static_cast<value_type>(p),
         1 - static_cast<value_type>(p),
         pol), "boost::math::gamma_p_inva<%1%>(%1%, %1%)");
}

template <class T1, class T2, class Policy>
inline typename tools::promote_args<T1, T2>::type
   gamma_q_inva(T1 x, T2 q, const Policy& pol)
{
   typedef typename tools::promote_args<T1, T2>::type result_type;
   typedef typename policies::evaluation<result_type, Policy>::type value_type;
   typedef typename policies::normalise<
      Policy,
      policies::promote_float<false>,
      policies::promote_double<false>,
      policies::discrete_quantile<>,
      policies::assert_undefined<> >::type forwarding_policy;

   if(q == 1)
   {
      return tools::max_value<result_type>();
   }
   if(q == 0)
   {
      return tools::min_value<result_type>();
   }

   return policies::checked_narrowing_cast<result_type, forwarding_policy>(
      detail::gamma_inva_imp(
         static_cast<value_type>(x),
         1 - static_cast<value_type>(q),
         static_cast<value_type>(q),
         pol), "boost::math::gamma_q_inva<%1%>(%1%, %1%)");
}

template <class T1, class T2>
inline typename tools::promote_args<T1, T2>::type
   gamma_p_inva(T1 x, T2 p)
{
   return boost::math::gamma_p_inva(x, p, policies::policy<>());
}

template <class T1, class T2>
inline typename tools::promote_args<T1, T2>::type
   gamma_q_inva(T1 x, T2 q)
{
   return boost::math::gamma_q_inva(x, q, policies::policy<>());
}

} // namespace math
} // namespace boost

#endif // BOOST_MATH_SP_DETAIL_GAMMA_INVA


///////////////////////////////////////////////////////////////////////////////
// importance_sampling::sampler.hpp                                          //
//                                                                           //
//  Copyright 2009 Erwann Rogard. 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_STATISTICS_DETAIL_IMPORTANCE_SAMPLING_SAMPLER_HPP_ER_2009
#define BOOST_STATISTICS_DETAIL_IMPORTANCE_SAMPLING_SAMPLER_HPP_ER_2009
#include <vector>
#include <boost/mpl/not.hpp>
#include <boost/mpl/assert.hpp>
#include <boost/call_traits.hpp>
#include <boost/type_traits.hpp>
#include <boost/utility.hpp>
#include <boost/range.hpp>
#include <boost/random/variate_generator.hpp>
#include <boost/random/uniform_real.hpp>
// TODO #include <boost/random/discrete_distributionhpp> when becomes avail
#include <boost/random/discrete_distribution_sw_2009.hpp>

namespace boost{
namespace statistics{
namespace detail{
namespace importance_sampling{

// Samples by SIR given a set of proposal values and their unnormalized weights
//
// Models RandomDistribution
//
// R1:  type of a range values
// W:   type of each weight
template<typename R1,typename W>
class sampler{
    typedef typename remove_reference<R1>::type         const_values_;
    typedef typename remove_cv<const_values_>::type     values_t;
    typedef typename range_size<values_t>::type         size_;
    typedef is_reference<R1>                            is_ref_;
    public:
    typedef typename range_value<const_values_>::type   result_type;
    private:
    typedef random::discrete_distribution<size_,W>      discr_dist_t;
    public:
    typedef typename discr_dist_t::input_type           input_type;

    sampler(){}
    template<typename R0>
    sampler(
        const R0& unnormalized_weights,
        typename call_traits<R1>::param_type values
    ):discr_dist_(
        boost::begin(unnormalized_weights),
        boost::end(unnormalized_weights)
    ),values_(values){
        BOOST_ASSERT(
            boost::size(unnormalized_weights) == boost::size(this->values())
        );
    }

    sampler(const sampler& that)
        :discr_dist_(that.discr_dist_),values_(that.values_){}

    sampler& operator=(const sampler& that)
    {
        if(&that!=this){
            discr_dist_ = that.discr_dist_;
            values_ = that.values_;
        }
        return (*this);
    }

    template<typename U>
    result_type operator()(U& urng)const
    {
        typedef typename discr_dist_t::result_type k_t;
        k_t k = discr_dist_(urng);
        BOOST_ASSERT( k < boost::size(this->values()) );
        return (*boost::next(boost::begin(this->values()),k));
    }
    const discr_dist_t& discrete_distribution()const
    {
        return this->discr_dist_;
    }

    // TODO os/is

    typename call_traits<R1>::const_reference values()const{
        return this->values_;
    }

    private:
    discr_dist_t discr_dist_;
    typename call_traits<R1>::value_type values_;
};

    template<typename R0,typename R1>
    sampler<
        R0,
        typename remove_cv<
            typename remove_reference<
                typename range_value<R0>::type
            >::type
        >::type
    >
    make_sampler(
        const R0& unnormalized_weights,
        typename call_traits<R1>::param_type values
    )
    {
        typedef sampler<
            R0,
            typename remove_cv<
                typename remove_reference<
                    typename range_value<R0>::type
                >::type
            >::type
        > result_;
        return result_(unnormalized_weights,values);
    }

}// importance_sampling
}// detail
}// statistics
}// boost

#endif


// smooth_feedback: Control theory on Lie groups
// https://github.com/pettni/smooth_feedback
//
// Licensed under the MIT License <http://opensource.org/licenses/MIT>.
//
// Copyright (c) 2021 Petter Nilsson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

#ifndef SMOOTH__FEEDBACK__ASIF_FUNC_HPP_
#define SMOOTH__FEEDBACK__ASIF_FUNC_HPP_

/**
 * @file
 * @brief Functions for active Set Invariance (ASI) filtering on Lie groups.
 */

#include <Eigen/Core>
#include <boost/numeric/odeint.hpp>
#include <smooth/compat/odeint.hpp>
#include <smooth/diff.hpp>
#include <smooth/lie_group.hpp>

#include "common.hpp"
#include "qp.hpp"

namespace smooth::feedback {

/**
 * @brief Active set invariance problem definition.
 *
 * The active set invariance problem is
 * \f[
 * \begin{cases}
 *  \min_u        & (u \ominus u_{des})' W_u (u \ominus u_{des})  \\
 *  \text{s.t.}   & x(0) = x_0 \\
 *                & u \in ulim \\
 *                & h(x(t)) \geq 0, \quad t \in [0, T]    \\
 * \end{cases}
 * \f]
 * for a system \f$ \mathrm{d}^r x_t = f(x(t), u(t)) \f$.
 */
template<LieGroup G, Manifold U>
  requires(Dof<G> > 0 && Dof<U> > 0)
struct ASIFProblem
{
  /// time horizon
  double T{1};
  /// initial state
  G x0{Default<G>()};
  /// desired input
  U u_des{Default<U>()};
  /// weights on desired input
  Eigen::Matrix<double, Dof<U>, 1> W_u{Eigen::Matrix<double, Dof<U>, 1>::Ones()};
  /// input bounds
  ManifoldBounds<U> ulim{};
};

/**
 * @brief Parameters for asif_to_qp
 */
struct ASIFtoQPParams
{
  /// number of constraint instances (equally spaced over the time horizon)
  std::size_t K{10};
  /// barrier function time constant \f$ \alpha \f$ s.t. \f$ \dot h - \alpha h \geq 0 \f$.
  double alpha{1};
  /// maximal integration time step
  double dt{0.1};
  /// relaxation cost
  double relax_cost{100};
};

/**
 * @brief Allocate QP matrices (part 1 of asif_to_qp())
 *
 * @param[in] K number of constraint instances
 * @param[in] nu_ineq number in inequalities in input constraint
 * @param[in] nh number of barrier constraints
 * @param[out] qp allocated QP with zero matrices
 */
template<LieGroup G, Manifold U>
  requires(Dof<G> > 0 && Dof<U> > 0)
void asif_to_qp_allocate(
  std::size_t K, std::size_t nu_ineq, std::size_t nh, QuadraticProgram<-1, -1, double> & qp)
{
  static constexpr int nx = Dof<G>;
  static constexpr int nu = Dof<U>;

  static_assert(nx > 0, "State space dimension must be static");
  static_assert(nu > 0, "Input space dimension must be static");

  const int M = K * nh + nu_ineq + 1;
  const int N = nu + 1;

  qp.A.setZero(M, N);
  qp.l.setZero(M);
  qp.u.setZero(M);

  qp.P.setZero(N, N);
  qp.q.setZero(N);
}

/**
 * @brief Fill QP matrices (part 2 of asif_to_qp())
 *
 * Note that the (dense) QP matrices must be pre-allocated and filled with zeros.
 */
template<
  LieGroup G,
  Manifold U,
  typename Dyn,
  typename SafeSet,
  typename BackupU,
  diff::Type DT = diff::Type::Default>
  requires(Dof<G> > 0 && Dof<U> > 0)
void asif_to_qp_fill(
  const ASIFProblem<G, U> & pbm,
  const ASIFtoQPParams & prm,
  Dyn && f,
  SafeSet && h,
  BackupU && bu,
  QuadraticProgram<-1, -1, double> & qp)
{
  using boost::numeric::odeint::euler, boost::numeric::odeint::vector_space_algebra;
  using std::placeholders::_1;

  static constexpr int nx = Dof<G>;
  static constexpr int nu = Dof<U>;
  static constexpr int nh = std::invoke_result_t<SafeSet, double, G>::SizeAtCompileTime;

  euler<G, double, Tangent<G>, double, vector_space_algebra> state_stepper{};
  euler<TangentMap<G>, double, TangentMap<G>, double, vector_space_algebra> sensi_stepper{};

  const int nu_ineq = pbm.ulim.A.rows();

  [[maybe_unused]] const int M = prm.K * nh + nu_ineq + 1;
  [[maybe_unused]] const int N = nu + 1;

  assert(qp.A.rows() == M);
  assert(qp.A.cols() == N);
  assert(qp.l.rows() == M);
  assert(qp.u.rows() == M);

  assert(qp.P.rows() == N);
  assert(qp.P.cols() == N);
  assert(qp.q.rows() == N);

  // iteration variables
  const double tau     = pbm.T / static_cast<double>(prm.K);
  const double dt      = std::min<double>(prm.dt, tau);
  double t             = 0;
  G x                  = pbm.x0;
  TangentMap<G> dx_dx0 = TangentMap<G>::Identity();

  // define ODEs for closed-loop dynamics and its sensitivity
  const auto x_ode = [&f, &bu](const G & xx, Tangent<G> & dd, double tt) {
    dd = f(tt, xx, bu(tt, xx));
  };

  const auto dx_dx0_ode = [&f, &bu, &x](const auto & S_v, auto & dS_dt_v, double tt) {
    auto f_cl = [&]<typename T>(const CastT<T, G> & vx) { return f(T(tt), vx, bu(T(tt), vx)); };
    const auto [fcl, dr_fcl_dx] = diff::dr<1, DT>(std::move(f_cl), wrt(x));
    dS_dt_v                     = (-ad<G>(fcl) + dr_fcl_dx) * S_v;
  };

  // value of dynamics at call time
  const auto [f0, d_f0_du] = diff::dr<1, DT>(
    [&]<typename T>(const CastT<T, U> & vu) { return f(T(t), cast<T>(x), vu); }, wrt(pbm.u_des));

  // loop over constraint number
  for (auto k = 0u; k != prm.K; ++k) {
    // differentiate barrier function w.r.t. x
    const auto [hval, dh_dtx] = diff::dr<1, DT>(
      [&h]<typename T>(const T & vt, const CastT<T, G> & vx) { return h(vt, vx); }, wrt(t, x));

    const Eigen::Matrix<double, nh, 1> dh_dt  = dh_dtx.template leftCols<1>();
    const Eigen::Matrix<double, nh, nx> dh_dx = dh_dtx.template rightCols<nx>();

    // insert barrier constraint
    const Eigen::Matrix<double, nh, nx> dh_dx0 = dh_dx * dx_dx0;
    qp.A.template block<nh, nu>(k * nh, 0)     = dh_dx0 * d_f0_du;
    qp.l.template segment<nh>(k * nh)          = -dh_dt - prm.alpha * hval - dh_dx0 * f0;
    qp.u.template segment<nh>(k * nh).setConstant(std::numeric_limits<double>::infinity());

    // integrate system and sensitivity forward until next constraint
    double dt_act = std::min(dt, tau * (k + 1) - t);
    while (t < tau * (k + 1)) {
      state_stepper.do_step(x_ode, x, t, dt_act);
      sensi_stepper.do_step(dx_dx0_ode, dx_dx0, t, dt_act);
      t += dt_act;
    }
  }

  // relaxation of barrier constraints
  qp.A.block(0, nu, prm.K * nh, 1).setConstant(1);

  // input bounds
  qp.A.block(prm.K * nh, 0, nu_ineq, nu) = pbm.ulim.A;
  qp.l.segment(prm.K * nh, nu_ineq)      = pbm.ulim.l - pbm.ulim.A * rminus(pbm.u_des, pbm.ulim.c);
  qp.u.segment(prm.K * nh, nu_ineq)      = pbm.ulim.u - pbm.ulim.A * rminus(pbm.u_des, pbm.ulim.c);

  // upper and lower bounds on delta
  qp.A(prm.K * nh + nu_ineq, nu) = 1;
  qp.l(prm.K * nh + nu_ineq)     = 0;
  qp.u(prm.K * nh + nu_ineq)     = std::numeric_limits<double>::infinity();

  qp.P.template block<nu, nu>(0, 0) = pbm.W_u.asDiagonal();

  qp.P(nu, nu) = prm.relax_cost;
  qp.q(nu)     = 0;
}

/**
 * @brief Convert an ASIFProblem to a QuadraticProgram.
 *
 * The objective is to impose constraints on the current input \f$ u \f$ of a system \f$
 * \mathrm{d}^r x_t = f(x, u) \f$ s.t.
 * \f[
 *    \frac{\mathrm{d}}{\mathrm{d}t} h(\phi(t; x_0, bu(\cdot)))
 *    \geq \alpha h(\phi(t; x_0, bu(\cdot)))
 * \f]
 * which enforces forward invariance of the set \f$ \{ x
 * : h(t, x) \geq 0 \} \f$ along the **backup trajectory** \f$ bu \f$. The constraint is enforced at
 * \f$K\f$ look-ahead time steps \f$ t_k = k \tau\f$ for \f$k = 0, \ldots, K \f$.
 *
 * This function encodes the problem as a QuadraticProgram that solves
 * \f[
 *  \begin{cases}
 *   \min_{\mu}  & \left\| u - u_{des} \right\|^2 \\
 *   \text{s.t.} & \text{constraint above holds for } u = u_{des} + \mu
 *   \end{cases}
 * \f]
 * A solution \f$ \mu^* \f$ to the QuadraticProgram corresponds to an input \f$ u_{des} \oplus \mu^*
 * \f$ applied to the system.
 *
 * @tparam G state LieGroup type \f$\mathbb{G}\f$
 * @tparam U input Manifold type \f$\mathbb{G}\f$
 *
 * @param pbm problem definition
 * @param prm algorithm parameters
 * @param f system model \f$f : \mathbb{R} \times \mathbb{G} \times \mathbb{U} \rightarrow
 * \mathbb{R}^{\dim \mathfrak g}\f$ s.t. \f$ \mathrm{d}^r x_t = f(t, x, u) \f$
 * @param h safe set \f$h : \mathbb{R} \times \mathbb{G} \rightarrow \mathbb{R}^{n_h}\f$ s.t. \f$
 * S(t) = \{ h(t, x) \geq 0 \} \f$ denotes the safe set at time \f$ t \f$
 * @param bu backup controller \f$ub : \mathbb{R} \times \mathbb{G} \rightarrow \mathbb{U} \f$
 *
 * \note The algorithm relies on automatic differentiation. The following supplied functions must be
 * differentiable (i.e. be templated on the scalar type if an automatic differentiation method is
 * selected):
 *   * \p f differentiable w.r.t. x and u
 *   * h differentiable w.r.t. t and x
 *   * bu differentiable w.r.t. x
 *
 * @return QuadraticProgram modeling the ASIF filtering problem
 */
template<
  LieGroup G,
  Manifold U,
  typename Dyn,
  typename SS,
  typename BackupU,
  diff::Type DT = diff::Type::Default>
QuadraticProgram<-1, -1, double> asif_to_qp(
  const ASIFProblem<G, U> & pbm, const ASIFtoQPParams & prm, Dyn && f, SS && h, BackupU && bu)
{
  static constexpr int nh = std::invoke_result_t<SS, double, G>::SizeAtCompileTime;

  static_assert(Dof<G> > 0, "State space dimension must be static");
  static_assert(Dof<U> > 0, "Input space dimension must be static");
  static_assert(nh > 0, "Safe set dimension must be static");

  const int nu_ineq = pbm.ulim.A.rows();
  QuadraticProgram<-1, -1, double> qp;
  asif_to_qp_allocate<G, U>(prm.K, nu_ineq, nh, qp);
  asif_to_qp_fill(pbm, prm, f, h, bu, qp);
  return qp;
}

}  // namespace smooth::feedback

#endif  // SMOOTH__FEEDBACK__INTERNAL__ASIF_FUNC_HPP_


//=======================================================================
// Copyright 2015 - 2020 Jeff Linahan
//
// 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)
//=======================================================================
#include "theorem4.h"
#include "strutil.h"
#include "graphutil.h"
#include "lemmas.h"
#include "BFSVisitor.h"
#include <boost/graph/copy.hpp>
#include <boost/graph/graph_concepts.hpp>
#include <boost/graph/is_straight_line_drawing.hpp>
#include <boost/graph/make_maximal_planar.hpp>
#include <boost/graph/copy.hpp>
#include <boost/bimap.hpp>
#include <boost/config.hpp>
#include <iostream>
using namespace std;
using namespace boost;

// may be able to inline this and hoist it out of the if blocks of theorem4_disconnected if we're able to get rid of bigger_than_two_thirds
uint lowest_i(uint n, uint num_components, vector<uint> const& num_verts_per_component)
{
        uint total_cost = 0;
        uint i = 0;
        for( ; i < num_components; ++i ){
                total_cost += num_verts_per_component[i];
                if ( total_cost*3 > n ) return i;
        }

        return -1;
}

// L[l] = # of vertices on level l
Partition theorem4_connected(GraphCR g, vector<uint> const& L, uint l[3], uint r, Graph* g_shrunk, BFSVisitorData const* bfs2)
{
        uint n = num_vertices(g);
        cout << "g:\n";
        //print_graph(g);
        vertex_t v;

        // l1 = the level such that the sum of costs in levels 0 thru l1-1 < 1/2, but the sum of costs in levels 0 thru l1 is >= 1/2
        uint total = 0;
        uint level = 0;
        while( total < n/2 ){
                total += L[level++];
        }

        uint l1 = level;

        BOOST_ASSERT(1 == L[0]);
        // Partition the vertices into levels according to their distance from some vertex v.
        //uint r;
        /*If r is the maximum distance of any vertex from v, define additional levels -1 and r+1 containing no vertices*/
        //uint l[3];
        //(If no such l1 exists, the total cost of all vertices < 1/2, and B = C = {} and return true) */
        uint k = 0;// = # of vertices on levels 0 thru l1.
        for( uint i = 0; i <= l1; ++i ) k += L[i];

        /*Find a level l0 such that l0 <= l1 and |L[l0]| + 2(l1-l0) <= 2sqrt(k)
        Find a level l2 such that l1+1 <= l2 and |L[l2] + 2(l2-l1-1) <= 2sqrt(n-k) */
        uint sqrtk = 2*sqrt(k);
        uint sqrtnk = 2*sqrt(n-k);

        uint l0 = l1;
        while( L[l0] + 2*(l1-l0) > sqrtk ) --l0;

        uint l2 = l1+1;
        while( L[l2] + 2*(l2-l1-1) > sqrtnk) ++l2;

        BOOST_ASSERT(l0 <= l1);
        BOOST_ASSERT(l1+1 <= l2);

        BFSVisitorData bfs(&g, *vertices(g).first);
        breadth_first_search(g, bfs.root, boost::visitor(BFSVisitor(bfs)));
        cout << "bfsroot: " << bfs.root << '\n';

        vector<vertex_t> cycle;
        l1 -= 2;
        --l2;
        Partition p = lemma3(g, L, l1, l2, r, bfs, *bfs2, cycle, g_shrunk);

        uint climit = sqrtk - sqrtnk;
        BOOST_ASSERT(p.verify_edges(g));
        //BOOST_ASSERT(p.verify_sizes_lemma3(L, l1, l2));

        /*If 2 such levels exist, then by Lemma 3 the vertices of G can be partitioned into three sets A, B, C such that no edge joins a vertex in A with a vertex in B,
        neither A or C has cost > 2/3, and C contains no more than 2(sqrt(k) + sqrt(n-k)) vertices.
        But 2(sqrt(k) + sqrt(n-k) <= 2(sqrt(n/2) + sqrt(n/2)) = 2sqrt(2)sqrt(n)
        Thus the theorem holds if suitable levels l0 and l2 exist
                Suppose a suitable level l0 does not exist.  Then, for i <= l1, L[i] >= 2sqrt(k) - 2(l1-i)
                Since L[0] = 1, this means 1 >= 2sqrt(k) - 2l1 and l1 + 1/2 >= sqrt(k).  Thus l1 = floor(l1 + 1/2) >
                Contradiction*/

        return p;
}

Partition theorem4_ccbigger23(GraphCR g_all, Partition const& biggest_comp_p)
{
        uint alln = num_vertices(g_all);
        cout << "alln: " << alln << '\n';

        Graph g_comp(g_all); // connected component

        auto g_all_prop_map = get(vertex_index, g_all);
        auto g_comp_prop_map = get(vertex_index, g_comp);

        auto vertid_to_vert_t = [&](uint id, decltype(get(vertex_index, g_all)) prop_map)
        {
                VertIter vit, vjt;
                tie(vit, vjt) = vertices(g_all);
                for( ; vit != vjt; ++vit ){
                        if( prop_map[*vit] == id ) return *vit;
                }
                return Graph::null_vertex();
        };

        // for all vertices v in g_comp,
        vector<vertex_t> toremove;
        VertIter vit, vjt;
        tie(vit, vjt) = vertices(g_comp);
        for( ; vit != vjt; ++vit ){
                vertex_t v = *vit;
                cout << "currently examining " << v << " comp_propmap" << g_comp_prop_map[v] << '\n';
                // cv = lookup the corresponding vertex in g_all
                vertex_t cv = vertid_to_vert_t(g_all_prop_map[v], g_all_prop_map);
                cout << "       with cv " << cv << " all_propmap" << g_all_prop_map[v] << '\n';
                // if cv is not in biggest_comp_p,
                if( !biggest_comp_p.a.contains(cv) &&
                    !biggest_comp_p.b.contains(cv) &&
                    !biggest_comp_p.c.contains(cv) ){
                        // delete v from g_comp
                            toremove.push_back(v);
                    }

        }

        cout << "gcomp:\n";
        //print_graph(g_comp);
        //print_graph_addresses(g_comp);

        uint nremove = toremove.size();
        cout << "nremove: " << nremove << '\n';
        for(auto& v : toremove ){
                cout << "   nr: " << v << '\n';
                kill_vertex(v, g_comp);
        }

        reset_vertex_indices(g_comp);

        cout << "gcomp:\n";
        //print_graph(g_comp);
        //print_graph_addresses(g_comp);
        BFSVisitorData vd(&g_comp, *vertices(g_comp).first);
        breadth_first_search(g_comp, vd.root, boost::visitor(BFSVisitor(vd)));

        // disabled because they don't support multiple connected components
        //BOOST_ASSERT(assert_verts(g_orig, vis_data_orig));
        //BOOST_ASSERT(assert_verts(g_copy, vis_data_copy));

        uint n = num_vertices(g_comp);

        vector<uint> L(vd.num_levels + 1, 0);
	cout << "L levels: " << L.size() << '\n';
        for( auto& d : vd.verts ){
                cout << "level: " << d.second.level << '\n';
	       	++L[d.second.level];
	}

        uint k = L[0];
        uint l[3];
        l[1] = 0;
        while( k <= n/2 ){
                uint indx = ++l[1];
                uint lsize = L.size();
                if( indx >= lsize ) break;
	       	k += L.at(indx);
	}

        float sq  = 2 * sqrt(k);
        float snk = 2 * sqrt(num_vertices(g_comp) - k);
        cout << "sq:     " << sq << '\n';
        cout << "snk:    " << snk << '\n';
        cout << "L size: " << L.size() << '\n';

        l[0] = l[1];
        cout << "l[0]:   " << l[0] << '\n';
        while( l[0] < L.size() ){
                float val = L.at(l[0]) + 2*(l[1] - l[0]);
                if( val <= sq ) break;
                --l[0];
        }
        cout << "l0: " << l[0] << "     highest level <= l1\n";

        l[2] = l[1] + 1;
        cout << "l[2]" << l[2] << '\n';
        while( l[2] < L.size() ){
                float val = L.at(l[2]) + 2*(l[2] - l[1] - 1);
                if( val <= snk ) break;
                ++l[2];
        }
        cout << "l2: " << l[2] << "     lowest  level >= l1 + 1\n";

        uint r = vd.num_levels - 1;

        Partition star_p = theorem4_connected(g_comp, L, l, r, nullptr, nullptr);
        // ????

        Partition p;
        p.c = star_p.c;
        p.b.insert(STLALL(toremove));
        if( star_p.a.size() > star_p.b.size() ){
                p.a = star_p.a;
                p.b.insert(STLALL(star_p.b));
        } else {
                p.a = star_p.b;
                p.b.insert(STLALL(star_p.a));
        }
        return p;
}

Partition theorem4_disconnected(GraphCR g, uint n, uint num_components, associative_property_map<vertex_map> const& vertid_to_component, vector<uint> const& num_verts_per_component, Partition const& biggest_comp_p)
{
        vector<vector<VertIter>> vertex_sets; // set of vertex ids, indexed by component number (second vector should be set but compiler did not like call to .insert())

        vertex_sets.resize(num_components);

        // populate vertex sets
        VertIter vit, vjt;
        tie(vit, vjt) = vertices(g);
        for( uint i = 0; vit != vjt; ++vit, ++i ){
                uint component = vertid_to_component[*vit];
                vector<VertIter>& vset = vertex_sets[component];
                vset.push_back(vit);
        }

        // Let G1, G2, ... , Gk be the connected components of G, with vertex sets V1, V2, ... , Vk respectively.
        bool bigger_than_one_third = false;
        bool bigger_than_two_thirds = false;
        for( uint i = 0; i < num_components; ++i ){
                if( num_verts_per_component[i] > n    /3.0 ) bigger_than_one_third  = true;
                if( num_verts_per_component[i] > n*2.0/3.0 ){ bigger_than_two_thirds = true; break;}
        }

        if( !bigger_than_one_third ){

                // If no connected component has total vertex cost > 1/3, let i be the minimum index such that the total cost of V1 U V2 U ... U Vi > 1/3
                int i = lowest_i(n, num_components, num_verts_per_component); // -1 indicates no such index found

                Partition p;

                // populate partition A = V1 U V2 U ... U Vi
                for( int j = 0; j <= i; ++j ){
                        vector<VertIter>& vec = vertex_sets[j];
                        for( VertIter& v : vec ) p.a.insert(*v);
                }

                // populate partition B = Vi+1 U Vi+2 U ... U Vk
                for( int j = i+1; j < (int)num_components; ++j ){
                        vector<VertIter>& vec = vertex_sets[j];
                        for( VertIter& v : vec ) p.b.insert(*v);
                }

                // Since i is minimum and the cost of Vi <= 1/3, the cost of A <= 2/3. return true;
                cout << "not bigger than one third\n";
                return p;
        } else if( !bigger_than_two_thirds ){
                // If some connected component (say Gi) has total vertex cost between 1/3 and 2/3,
                Partition p;

                int i = lowest_i(n, num_components, num_verts_per_component);
                BOOST_ASSERT(i >= 0 && num_verts_per_component[i] >= n/3 && num_verts_per_component[i] <= 2*n/3);

                // populate partition A
                cout << "!! populating partition A\n";
                for( uint j = 0; j <= i; ++j ){
                        for( uint v = 0; v < vertex_sets[j].size(); ++v ){
                                p.a.insert(*vertex_sets[j][v]);
                        }

                }
                cout << '\n';

                // populate partition B, should be everything except what's in partition A
                cout << "!! populating partition B\n";
                for( uint j = i+1; j < num_components; ++j ){

                        vector<VertIter>& vset = vertex_sets[j];
                        for( VertIter& v : vset ){
                                p.b.insert(*v);
                                cout << "v: " << *v << endl;
                        }
                }

                // partition C should be empty
                BOOST_ASSERT(p.c.empty());

                p.print(&g);

                cout << "bigger than one third but less than two thirds\n";
                return p;
        }

        return theorem4_ccbigger23(g, biggest_comp_p);
}

/* Theorem 4: Let G be any (possibly disconnected) n-vertex planar graph having nonnegative vertex costs summing to no more than one.
Then the vertices of G can be partitioned into three sets A, B, C such that no edge joins a vertex
in A with a vertex in B, neither A nor B has total cost exceeding 2/3, and C contains no more than
2sqrt(2)sqrt(n) vertices :*/
Partition theorem4(GraphCR g, associative_property_map<vertex_map> const& vertid_to_component, vector<uint> const& num_verts_per_component, Partition const& biggest_comp_p)
{
	uint n = num_vertices(g);
	uint num_components = num_verts_per_component.size();
	bool is_graph_connected = (num_components == 1);

	if( is_graph_connected ){
		cout << "graph is connected\n";
                vector<uint> L;
                uint l[3];
                uint r;
                BOOST_ASSERT(0);
                return theorem4_connected(g, L, l, r, nullptr, nullptr);
	} else {
                cout << "graph is disconnected with " << num_components << " components\n";
                return theorem4_disconnected(g, n, num_components, vertid_to_component, num_verts_per_component, biggest_comp_p);
	}
}


/**
 * @file EqRoutine.hpp
 * @author A. Domahidi [domahidi@embotech.com]
 * @license (new license) License BSD-3-Clause
 * @copyright Copyright (c) [2012-2015] Automatic Control Lab, ETH Zurich & embotech GmbH, Zurich, Switzerland.
 * @date 2012-2015
 * @brief ECOS - Embedded Conic Solver
 *
 * Modified to c++ code by New York University and Max Planck Gesellschaft, 2017
 */

#pragma once

#include <memory>
#include <Eigen/Sparse>
#include <solver/interface/Cone.hpp>
#include <solver/interface/SolverSetting.hpp>

namespace solver {

  /*! Equilibration routine to improve condition number of
   *  matrices involved in the optimization problem. The method
   *  provided by default is Ruiz equilibration.
   */
  class EqRoutine
  {
    public:
	  EqRoutine(){}
	  ~EqRoutine(){}

	  void setEquilibration(const Cone& cone, const SolverSetting& stgs, SolverStorage& stg);
	  void unsetEquilibration(SolverStorage& stg);
	  void scaleVariables(OptimizationVector& opt);

      Vector& equilVec() { return equil_vec_; }
      const Vector& equilVec() const { return equil_vec_; }

    private:
	  void ruizEquilibration(SolverStorage& stg);
	  void maxRowsCols(double *row_vec, double *col_vec, const Eigen::SparseMatrix<double>& mat);
	  void equilibrateRowsCols(const double *row_vec, const double *col_vec, Eigen::SparseMatrix<double>& mat);
	  void unequilibrateRowsCols(const double *row_vec, const double *col_vec, Eigen::SparseMatrix<double>& mat);

    private:
	  Vector equil_vec_;
	  std::shared_ptr<Cone> cone_;
	  std::shared_ptr<SolverSetting> stgs_;
  };

}


/*  $Id: bullcutter.cpp 667 2011-02-14 22:13:45Z anders.e.e.wallin $
 *
 *  Copyright 2010 Anders Wallin (anders.e.e.wallin "at" gmail.com)
 *
 *  This file is part of OpenCAMlib.
 *
 *  OpenCAMlib 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.
 *
 *  OpenCAMlib 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 OpenCAMlib.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <boost/foreach.hpp>

#include "bullcutter.h"
#include "numeric.h"
#include "ellipse.h"

namespace ocl
{

BullCutter::BullCutter() {
    std::cout << " usage: BullCutter( double diameter, double corner_radius, double length ) \n";
    assert(0);
}

BullCutter::BullCutter(double d, double r, double l) {
    diameter = d;               assert( d > 0.0 );
    radius = d/2.0;        // total cutter radius
    radius1 = d/2.0 - r;   // cylindrical middle part radius
    radius2 = r;                assert( radius1 > 0.0 ); // corner radius
    length = l;                 assert( l > 0.0 );
    xy_normal_length = radius1;
    normal_length = radius2;
    center_height = radius2;
}

MillingCutter* BullCutter::offsetCutter(double d) const {
    return new BullCutter(diameter+2*d, radius2+d, length+d) ;
}

// height of cutter at radius r
double BullCutter::height(double r) const {
    if ( r <= radius1 )
        return 0.0; // cylinder
    else if ( r <= radius )
        return radius2 - sqrt( square(radius2) - square(r-radius1) ); // toroid
    else {
        assert(0);
        return -1;
    }
}

// width of cutter at height h
double BullCutter::width(double h) const {
    return ( h >= radius2 ) ? radius : radius1 + sqrt(square(radius2)-square(radius2-h)) ;
}

// drop-cutter: vertex and facet are handled in base-class

// drop-cutter: Toroidal cutter edge-test
CC_CLZ_Pair BullCutter::singleEdgeDropCanonical( const Point& u1, const Point& u2 ) const {
    if ( isZero_tol( u1.z - u2.z ) ) {  // horizontal edge special case
        return CC_CLZ_Pair( 0 , u1.z - height(u1.y) );
    } else { // the general offset-ellipse case
        double b_axis = radius2;                            // short axis of ellipse = radius2
        double theta = atan( (u2.z - u1.z) / (u2.x-u1.x) ); // theta is the slope of the line
        double a_axis = fabs( radius2/sin(theta) );         // long axis of ellipse = radius2/sin(theta)
        Point ellcenter(0,u1.y,0);
        Ellipse e = Ellipse( ellcenter, a_axis, b_axis, radius1);
        int iters = e.solver_brent();
        assert( iters < 200 );
        e.setEllipsePositionHi(u1,u2); // this selects either EllipsePosition1 or EllipsePosition2 and sets it to EllipsePosition_hi
        // pseudo cc-point on the ellipse/cylinder, in the CL=origo system
        Point ell_ccp = e.ePointHi();         assert( fabs( ell_ccp.xyNorm() - radius1 ) < 1E-5); // ell_ccp should be on the cylinder-circle
        Point cc_tmp_u = ell_ccp.closestPoint(u1,u2); // find real cc-point
        return CC_CLZ_Pair( cc_tmp_u.x , e.getCenterZ()-radius2);
    }
}

// push-cutter: vertex and facet handled by base-class

bool BullCutter::generalEdgePush(const Fiber& f, Interval& i,  const Point& p1, const Point& p2) const {
    //std::cout << " BullCutter::generalEdgePush() \n";
    bool result = false;

    if ( isZero_tol( (p2-p1).xyNorm() ) ) { // this would be a vertical edge
        return result;
    }

    if ( isZero_tol( p2.z-p1.z ) ) // this would be a horizontal edge
        return result;
    assert( fabs(p2.z-p1.z) > 0.0 ); // no horiz edges allowed hereafter

    // p1+t*(p2-p1) = f.p1.z+radius2   =>
    double tplane = (f.p1.z + radius2 - p1.z ) / (p2.z-p1.z); // intersect edge with plane at z = ufp1.z
    Point ell_center = p1+tplane*(p2-p1);
    assert( isZero_tol( fabs(ell_center.z - (f.p1.z+radius2)) ) );
    Point major_dir = (p2-p1);
    assert( major_dir.xyNorm() > 0.0 );

    major_dir.z = 0;
    major_dir.xyNormalize();
    Point minor_dir = major_dir.xyPerp();
    double theta = atan( (p2.z - p1.z) / (p2-p1).xyNorm() );
    double major_length = fabs( radius2/sin(theta) ) ;
    double minor_length = radius2;
    AlignedEllipse e(ell_center, major_length, minor_length, radius1,  major_dir, minor_dir );
    if ( e.aligned_solver( f ) ) { // now we want the offset-ellipse point to lie on the fiber
        Point pseudo_cc  = e.ePoint1(); // pseudo cc-point on ellipse and cylinder
        Point pseudo_cc2 = e.ePoint2();
        CCPoint cc  = pseudo_cc.closestPoint(p1,p2);
        CCPoint cc2 = pseudo_cc2.closestPoint(p1,p2);
        cc.type  = EDGE_POS;
        cc2.type = EDGE_POS;
        Point cl  = e.oePoint1() - Point(0,0,center_height);
        assert( isZero_tol( fabs(cl.z - f.p1.z)) );
        Point cl2 = e.oePoint2() - Point(0,0,center_height);
        assert( isZero_tol( fabs(cl2.z - f.p1.z)) );
        double cl_t  = f.tval(cl);
        double cl_t2 = f.tval(cl2);
        if ( i.update_ifCCinEdgeAndTrue( cl_t, cc, p1, p2, true ) )
            result = true;
        if ( i.update_ifCCinEdgeAndTrue( cl_t2, cc2, p1, p2, true ) )
            result = true;
    }
    //std::cout << " BullCutter::generalEdgePush() DONE result= " << result << "\n";
    return result;
}

std::string BullCutter::str() const {
    std::ostringstream o;
    o << *this;
    return o.str();
}

std::ostream& operator<<(std::ostream &stream, BullCutter c) {
  stream << "BullCutter(d=" << c.diameter << ", r1=" << c.radius1 << " r2=" << c.radius2 << ", L=" << c.length <<  ")";
  return stream;
}

} // end namespace
// end file bullcutter.cpp


/*!
  * \file
  * \brief Include this into the main file to have access to the MHD library and code.
  */
#pragma once

// std includes
#include <iostream>
#include <fstream>
#include <cmath>
#include <memory>
// external libraries
#include <armadillo>
#include "viennacl/vector.hpp"
#include "viennacl/matrix.hpp"
#include "viennacl/compressed_matrix.hpp"
#include "viennacl/linalg/prod.hpp"
#include "viennacl/linalg/gmres.hpp"
#include "viennacl/linalg/bicgstab.hpp"
#include "viennacl/linalg/ilu.hpp"
#include "viennacl/linalg/lu.hpp"
#include "viennacl/linalg/sum.hpp"
#include "viennacl/linalg/maxmin.hpp"
#include "viennacl/tools/random.hpp"
#include "viennacl/linalg/inner_prod.hpp"
#include "viennacl/tools/timer.hpp"
#include "viennacl/forwards.h"

typedef viennacl::vector<double> vcl_vec; /*!< GPU arma::vector with ViennaCL */
typedef viennacl::compressed_matrix<double> vcl_sp_mat; /*!< GPU sparse matrix with ViennaCL */
typedef viennacl::matrix<double> vcl_mat; /*!< GPU dense matrix with ViennaCL */

/*!
 * Everything in CoMFi will be in this namespace for safety. This is the top level namespace.
 */
namespace comfi
{

/*!
* All the classes, enums and types can be found here.
*/
namespace types {

/*!
 * \struct Settings
 * \brief The settings for the simulation (used in main loop) are here.
 */
struct Settings {
  /// Maximum number of time steps before simulation end.
  int max_time_steps = -1;
  /// Maximum time before simulation end in normalized units. \f$(t_0)\f$
  double max_time = -1.0;
  /// gmres tolerance
  double tolerance = 1.e-6;
  /// Save solution every X time steps
  int save_dn = -1;
  /// Save solution every X time steps
  double save_dt = -1;
  /// If this run restarts from a previous run
  bool restart = false;
  /// Other flags for runtime.
  uint flags = 0;
};

/*!
 * \brief The Boundary Condition enum.
 */
enum BoundaryCondition {PERIODIC, /*!< Periodic boundary condition (i+1 at right boundary is i=0) */
                        DIRICHLET, /*!< Dirichlet boundary conditions. Inserting data up to user. */
                        NEUMANN, /*!< Neumann boundary conditions. (i+1 at right boundary = i) */
                        MIRROR, /*!< Mirror boundary conditions. (i+1 at right boundary = i or negative that value in case of x component) */
                        DIMENSIONLESS /*!< This boundary is not in a free dimension. */
                       };

/*!
 * \brief The simulation context. Everything needed about the simulation is in here.
 */
class Context {
  /// Current \f$\Delta t\f$ in normalized units of \ref t_0 .
  double m_dt = 0.0;
  /// Current \f$t\f$ in normalized units of \ref t_0 .
  double m_time_elapsed = 0.0;
  /// Current step \f$n\f$ .
  uint m_time_step = 0;

public:
  // Simulation domain
  /// Number of grid points in x-direction (horizontal).
  const arma::uword nx;
  /// Number of grid points in z-direction (vertical).
  const arma::uword nz;
  /// Upper boundary condition in the z-direction.
  const BoundaryCondition bc_up;
  /// Lower boundary condition in the z-direction.
  const BoundaryCondition bc_down;
  /// Right boundary condition in the x-direction.
  const BoundaryCondition bc_right;
  /// Left boundary condition in the x-direction.
  const BoundaryCondition bc_left;

  /// Settings passed on from main.
  const Settings settings;

  // Constants
  /// ratio of specific heats
  //const double gammamono = 5.0/3.0;
  const double gammamono = 2.0;
  /// \f$alpha_p\f$ divergence error propogation parameter
  const double alpha_p = 0.18;
  /// π
  const double pi = arma::datum::pi;
  /// mass of proton in kg
  const double m_i = arma::datum::m_p;
  /// permeability of free space
  const double mu_0 = arma::datum::mu_0;
  /// boltzmann constant in SI
  const double k_b = arma::datum::k;
  /// elementary charge in C
  const double e_ = arma::datum::ec;
  // Normalization constants
  /// Length normalization constant in meters.
  const double l_0 = 1.0;
  /// \f$m^{-3}\f$ density normalization constant.
  const double n_0 = 25.0/(36.0*arma::datum::pi);
  /// Magnetic field normalization constant in Teslas.
  const double B_0 = 1.0/std::sqrt(4.0*arma::datum::pi);
  /// Width (x-direction) in normalized units.
  const double width = l_0;
  /// Height (z-direction) in normalized units.
  const double height = l_0;
  // Derived constants
  /// \f$\Delta x\f$ in normalized units.
  const double dx = (width/nx)/l_0;
  /// \f$\Delta z\f$ in normalized units.
  const double dz = (height/nz)/l_0;
  /// \f$\Delta s\f$ (smallest grid length) in normalized units.
  const double ds = (dx>dz)*dz + (dz>=dx)*dx;
  /// Normalized mass of electron in kg.
  const double m_e = arma::datum::m_e/m_i;
  /// Derived normalization constant of Alfven velocity in m/s. \f$V_0 = V_A = B_0 / \sqrt{\mu_0 n_0 m_i}\f$
  const double V_0 = B_0/std::sqrt(mu_0*n_0*m_i);
  /// Derived normalization constant of time in seconds.
  const double t_0 = l_0/V_0;
  /// Derived normalization constant of pressure in Pascals. \f$p_0 = B_0^2 / \mu_0\f$
  const double p_0 = B_0*B_0/mu_0;
  /// Derived normalization constant of temperature in Kelvin. \f$T_0 = p_0 / (n_0 k_b)\f$
  const double T_0 = p_0/(n_0*k_b);
  /// In normalized units, default is sun value of 0.27395 km/s/s.
  const double g = 0.27395e3*t_0/V_0;
  /// Derived normalization constant for heat transfer.
  const double q_0               = p_0*V_0;
  /// Derived normalization constant for charge in Coulombs.
  const double e_0               = B_0/(mu_0*n_0*V_0*l_0);
  /// Normalized elementary charge constant.
  const double q                 = e_/e_0;
  /// Derived normalization constant for heat transfer coefficient.
  const double kappa_0           = q_0*l_0/T_0;

  // Solution index
  /// Result vector top level index
  const arma::uword Bx    = 0;
  /// Result vector top level index
  const arma::uword Bz    = 1;
  /// Result vector top level index
  const arma::uword Bp    = 2;
  /// Result vector top level index
  const arma::uword n_n   = 3;
  /// Result vector top level index
  const arma::uword Ux    = 4;
  /// Result vector top level index
  const arma::uword Uz    = 5;
  /// Result vector top level index
  const arma::uword Up    = 6;
  /// Result vector top level index
  const arma::uword n_p   = 7;
  /// Result vector top level index
  const arma::uword Vx    = 8;
  /// Result vector top level index
  const arma::uword Vz    = 9;
   /// Result vector top level index
  const arma::uword Vp    = 10;
  /// Result vector top level index
  const arma::uword E_n   = 11;
   /// Result vector top level index
  const arma::uword E_p   = 12;
  /// Result vector top level index
  const arma::uword GLM   = 13;

  // Ranges
  const viennacl::range r_grid = viennacl::range(0, nx*nz);
  const viennacl::range r_Np = viennacl::range(n_p, n_p+1);
  const viennacl::range r_Nn = viennacl::range(n_n, n_n+1);
  const viennacl::range r_NVx = viennacl::range(Vx, Vx+1);
  const viennacl::range r_NVz = viennacl::range(Vz, Vz+1);
  const viennacl::range r_NVp = viennacl::range(Vp, Vp+1);
  const viennacl::range r_NUx = viennacl::range(Ux, Ux+1);
  const viennacl::range r_NUz = viennacl::range(Uz, Uz+1);
  const viennacl::range r_NUp = viennacl::range(Up, Up+1);
  const viennacl::range r_Ep = viennacl::range(E_p, E_p+1);
  const viennacl::range r_En = viennacl::range(E_n, E_n+1);
  const viennacl::range r_Bx = viennacl::range(Bx, Bx+1);
  const viennacl::range r_Bz = viennacl::range(Bz, Bz+1);
  const viennacl::range r_Bp = viennacl::range(Bp, Bp+1);
  const viennacl::range r_GLM = viennacl::range(GLM, GLM+1);
  const arma::uword num_of_eq = 14;

  Context(arma::uword _nx,
          arma::uword _nz,
          BoundaryCondition _bc_up=NEUMANN,
          BoundaryCondition _bc_down=NEUMANN,
          BoundaryCondition _bc_right=NEUMANN,
          BoundaryCondition _bc_left=NEUMANN,
          const Settings _settings = Settings()
          ) : nx(_nx), nz(_nz),
              bc_up(_bc_up),
              bc_down(_bc_down),
              bc_right(_bc_right),
              bc_left(_bc_left),
              settings(_settings) {
  }

  /*!
   * \brief returns number of grid points per unknown.
   * \return Number of grid points per unknown.
   */
  arma::uword num_of_grid() const { return nz*nx; }

  /*!
   * \brief Set \f$\Delta t\f$ in normalized units of \ref t_0 .
   * \param dt Time step in normalized units of \ref t_0 .
   * Since \ref m_dt is a private member. This variable must be changed through this function.
   */
  void set_dt(const double &dt) {
    m_dt = dt;
  }

  /*!
   * \brief Advance the state of the simulation.
   * This will advance the state of the simulation so that the \ref time_elapsed()
   * and \ref time_step() are updated.
   */
  void advance() {
    m_time_elapsed += m_dt;
    m_time_step++;
  }

  /// Getter of \f$\c_h\f$
  double c_h() const { return ds/m_dt; }
  /// Getter of \f$\Delta t\f$
  double dt() const { return m_dt; }
  /// Getter of \f$t\f$
  double time_elapsed() const { return m_time_elapsed; }
  /// Getter of \f$n\f$
  uint time_step() const { return m_time_step; }

  // Matrix range returns
  inline viennacl::range range(const uint &var) const { return viennacl::range(var, var+1); }
  inline viennacl::matrix_range<vcl_mat> range(const uint &var, const vcl_mat &xn) {
    return project(xn, r_grid, range(var));
  }
  inline viennacl::matrix_range<vcl_mat> v_Np(const vcl_mat &xn) {
    return project(xn, r_grid, r_Np);
  }
  inline viennacl::matrix_range<vcl_mat> v_Nn(const vcl_mat &xn) {
    return project(xn, r_grid, r_Nn);
  }
  inline viennacl::matrix_range<vcl_mat> v_NVx(const vcl_mat &xn) {
    return project(xn, r_grid, r_NVx);
  }
  inline viennacl::matrix_range<vcl_mat> v_NVz(const vcl_mat &xn) {
    return project(xn, r_grid, r_NVz);
  }
  inline viennacl::matrix_range<vcl_mat> v_NVp(const vcl_mat &xn) {
    return project(xn, r_grid, r_NVp);
  }
  inline viennacl::matrix_range<vcl_mat> v_NUx(const vcl_mat &xn) {
    return project(xn, r_grid, r_NUx);
  }
  inline viennacl::matrix_range<vcl_mat> v_NUz(const vcl_mat &xn) {
    return project(xn, r_grid, r_NUz);
  }
  inline viennacl::matrix_range<vcl_mat> v_NUp(const vcl_mat &xn) {
    return project(xn, r_grid, r_NUp);
  }
  inline viennacl::matrix_range<vcl_mat> v_Bx(const vcl_mat &xn) {
    return project(xn, r_grid, r_Bx);
  }
  inline viennacl::matrix_range<vcl_mat> v_Bz(const vcl_mat &xn) {
    return project(xn, r_grid, r_Bz);
  }
  inline viennacl::matrix_range<vcl_mat> v_Bp(const vcl_mat &xn) {
    return project(xn, r_grid, r_Bp);
  }
  inline viennacl::matrix_range<vcl_mat> v_GLM(const vcl_mat &xn) {
    return project(xn, r_grid, r_GLM);
  }
  inline viennacl::matrix_range<vcl_mat> v_Ep(const vcl_mat &xn) {
    return project(xn, r_grid, r_Ep);
  }
  inline viennacl::matrix_range<vcl_mat> v_En(const vcl_mat &xn) {
    return project(xn, r_grid, r_En);
  }
};

} // namespace types

/*!
 * \brief Functions to build or apply the 'operators'.
 * Permutations of the matrix.
 */
namespace operators {

/*!
* \brief Permutate a solution down one cell.
* \param xn Solution matrix context.
* \param ctx Simulation context.
* \return Matrix of \f$A_{j-1}\f$
* Every index is replaced by the cell below it while maintaining boundary conditions.
*/
vcl_mat jm1(const vcl_mat &xn, comfi::types::Context ctx);

/*!
* \brief Permutate a solution up one cell.
* \param xn Solution matrix context.
* \param ctx Simulation context.
* \return Matrix of \f$A_{j+1}\f$
* Every index is replaced by the cell above it while maintaining boundary conditions.
*/
vcl_mat jp1(const vcl_mat &xn, comfi::types::Context ctx);

/*!
* \brief Permutate a solution left one cell.
* \param xn Solution matrix context.
* \param ctx Simulation context.
* \return Matrix of \f$A_{i-1}\f$
* Every index is replaced by the cell to the left of it while maintaining boundary conditions.
*/
vcl_mat im1(const vcl_mat &xn, comfi::types::Context ctx);

/*!
* \brief Permutate a solution right one cell.
* \param xn Solution matrix context.
* \param ctx Simulation context.
* \return Matrix of \f$A_{i+1}\f$
* Every index is replaced by the cell to the right of it while maintaining boundary conditions.
*/
vcl_mat ip1(const vcl_mat &xn, comfi::types::Context ctx);

/*!
 * \brief Permutate a result arma::vector right one cell.
 * \param ctx Simulation context.
 * \return Sparse matrix to multiply a result arma::vector by.
 * Every index is replaced by the cell to the right of it while maintaining boundary conditions.
 */
const arma::sp_mat buildPip1(comfi::types::Context &ctx);

/*!
 * \brief Permutate a result arma::vector left one cell.
 * \param ctx Simulation context.
 * \return Sparse matrix to multiply a result arma::vector by.
 * Every index is replaced by the cell to the left of it while maintaining boundary conditions.
 */
const arma::sp_mat buildPim1(comfi::types::Context &ctx);

} // namespace operators

/*!
 * Utility functions. This namespace contains non-mathematical functions to assist with IO, monitoring, logging and error checking. Also included are RHS and LHS functions.
 */
namespace util
{
/*!
 * \brief gettimestr Get current time.
 * \return const std::string in YYYY-MM-DD-HH-MM format.
 */
std::string gettimestr();

/*!
 * \brief save_solution Saves the solution to the output folder
 * \param x0 Solution matrix.
 * \param ctx Simulation context.
 * \param data_name data set name to save it under in hdf5 tree. Defaults to "unknowns"
 * \return Success or fail.
 */
bool save_solution(const vcl_mat &x0,
                   comfi::types::Context &ctx,
                   const std::string &data_name = "unknowns");

/*!
 * \brief sendtolog Send message to log file.
 * \param message Message std::string to show in log file.
 * \param filename The log filename make sure this is always the same.
 */
void sendtolog(const std::string message, const std::string filename);

/*!
 * \brief Calculate the initial condition based on Orszang-Tang Vortex.
 * \param ctx Simulation context
 * \return initial condition matrix
 */
vcl_mat ot_vortex_ic(comfi::types::Context &ctx);

/*!
 * \brief Calculate the initial condition based on Sod's Shock Tube.
 * \param ctx Simulation context
 * \return initial condition matrix
 */
vcl_mat shock_tube_ic(comfi::types::Context &ctx);

// Misc

/*!
 * \brief getmaxV get fast mode speed + local speed OR resistivity speed whichever is faster
 * \param x0 Solution matrix
 * \param ctx Simulation context
 * \return Max characteristic speed in normalized units
 */
double getmaxV(const vcl_mat &x0, comfi::types::Context &ctx);

/*!
 * \brief vec_to_mat Change a column vector to a viennacl::matrix type with one column
 * \param vec viennacl::vector
 * \return viennacl::matrix of one column
 */
vcl_mat vec_to_mat(const vcl_vec &vec);

/*!
 * \brief Write a binary file of a field arma::vector. Seperate binary files for each direction.
 * \param x0 result arma::vector
 * \param name Name of field, will be used as filename prefix.
 * \param timestep Time step to save as
 * \return True if no error.
 */
bool saveField(const arma::vec &x0, const std::string name, const int timestep);

/*!
 * \brief Read the arguments passed to the executable and change the Settings struct.
 */
void interpret_arguments(comfi::types::Settings &settings, int argc, char** argv);

/*!
 * \brief saveScalar Write a binary file of a scalar arma::vector.
 * \param x0 result arma::vector
 * \param name Name of scalar, will be used as filename prefix.
 * \param timestep Time step to save as
 * \return True if no error.
 */
bool saveScalar(const arma::vec &x0, const std::string name, const int timestep);

/*!
 * \brief saveField Write a binary file of a field arma::vector. Seperate binary files for each direction.
 * \param x0 result arma::vector
 * \param name Name of field, will be used as filename prefix.
 * \param timestep Time step to save as
 * \return True if no error.
 */
bool saveField(const vcl_vec &x0, const std::string name, const int timestep);

/*!
 * \brief saveScalar Write a binary file of a scalar arma::vector.
 * \param x0 result arma::vector
 * \param name Name of scalar, will be used as filename prefix.
 * \param timestep Time step to save as
 * \return True if no error.
 */
bool saveScalar(const vcl_vec &x0, const std::string name, const int timestep);

} // namespace util

/*!
 * Linear algebra routines and deriving variables and such are all found here.
 */
namespace routines
{

/*!
 * \brief Build an eigenvalue matrix (wave speeds) for the Lax-Friedrichs scheme in the \f$x\f$-direction.
 * \param p_eig Ion wave speeds.
 * \param n_eig Neutral wave speeds.
 */
vcl_mat build_eig_matrix_x(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief Build an eigenvalue matrix (wave speeds) for the Lax-Friedrichs scheme in the \f$z\f$-direction.
 * \param p_eig Ion wave speeds.
 * \param n_eig Neutral wave speeds.
 */
vcl_mat build_eig_matrix_z(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief pressure_n Returns the plasma pressure based on the type of energy being used
 * \param xn Solution matrix
 * \param ctx Simulation context
 * \return Simulation pressure matrix in (grid, 1) dimensions
 */
vcl_mat pressure_p(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief pressure_n Returns the neutral pressure based on the type of energy being used
 * \param xn Solution matrix
 * \param ctx Simulation context
 * \return Simulation pressure matrix in (grid, 1) dimensions
 */
vcl_mat pressure_n(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief Fx Get flux values in the x direction.
 * \param xn Solution matrix in the cell edge.
 * \param xn_ij Solution matrix in the cell center.
 * \param ctx Simulation context.
 * \return Flux function matrix results.
 */
vcl_mat Fx(const vcl_mat &xn, const vcl_mat &xn_ij, comfi::types::Context &ctx);

/*!
 * \brief Fz Get flux values in the z direction.
 * \param xn Solution matrix in the cell edge.
 * \param xn_ij Solution matrix in the cell center.
 * \param ctx Simulation context.
 * \return Flux function matrix results.
 */
vcl_mat Fz(const vcl_mat &xn, const vcl_mat &xn_ij, comfi::types::Context &ctx);

/*!
 * \brief Re_MUSCL Do flux reconstruction with the MUSCL scheme and any source terms.
 * \param xn Solution matrix
 * \param ctx Simulation context
 * \return Solution.
 */
vcl_mat Re_MUSCL(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief Flux limiter
 * \param r ratio of gradients
 * \return Limiter function values
 */
vcl_mat fluxl(const vcl_mat &r);

/*!
 * \brief Get plasma sound speed
 * \param xn matrix of results
 * \param ctx Simulation context
 * \return vector (vcl_mat) of plasma sound speed
 */
vcl_mat sound_speed_p(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief Get neutral sound speed
 * \param xn matrix of results
 * \param ctx Simulation context
 * \return vector (vcl_mat) of neutral sound speed
 */
vcl_mat sound_speed_n(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief fast_speed_z Get vertical fast mode speed
 * \param xn result matrix
 * \param ctx Simulation context.
 * \return Column vector vcl_mat of vertical fast mode speed of ions.
 */
vcl_mat fast_speed_z(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief fast_speed_x Get horizontal fast mode speed
 * \param xn Result matrix.
 * \param ctx	Simulation context.
 * \return Column vector (vcl_mat) of horizontal fast mode speeds.
 */
vcl_mat fast_speed_x(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief polyval compute's polynomial using Herner's scheme
 * \param p polynomial coefficient vector of size (polynomial degree + 1)
 * \param x vector to be computed
 * \return result of the polynomial
 */
vcl_vec polyval(const arma::vec &p, const vcl_vec &x);

/*!
 * \brief computeRHS_Euler Compute right hand side using Eulerian time stepping.
 * \param xn Current time step result matrix
 * \param ctx Simulation context
 * \return Right hand side result solution matrix
 */
vcl_mat computeRHS_Euler(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief computeRHS_RK4 Compute right hand side using Runge-Kutta 4 time stepping.
 * \param xn Current time step result arma::vector
 * \param dt Change in time
 * \param t Time elapsed
 * \param op operators
 * \param bg Background data
 * \return  Right hand side result arma::vector
 */
vcl_mat computeRHS_RK4(const vcl_mat &xn, comfi::types::Context &ctx);

/*!
 * \brief Shock Tube dirichlet boundary conditions
 * \param Lxn Left state
 * \param Rxn Right state
 * \param ctx Context
 */
void topbc_shock_tube(vcl_mat &Lxn, vcl_mat &Rxn, comfi::types::Context &ctx);

/*!
 * \brief Shock Tube dirichlet boundary conditions
 * \param Lxn Left state
 * \param Rxn Right state
 * \param ctx Context
 */
void bottombc_shock_tube(vcl_mat &Lxn, vcl_mat &Rxn, comfi::types::Context &ctx);

// inlines

/* inline const vcl_vec div_tvd(const vcl_vec &s_iph, const vcl_vec &s_imh, const vcl_vec &s_jph, const vcl_vec &s_jmh) */
/* { */
/*   const vcl_vec xpart = (s_iph-s_imh)/dx; */
/*   const vcl_vec zpart = (s_jph-s_jmh)/dz; */
/*   return xpart+zpart; */
/* } */

} // namespace routines

/*!
 * Solar chromosphere math functions.
 */
namespace sol
{

inline double nu_nn(const double &nn, const double &T, const comfi::types::Context &ctx)
{
  const double sigma_nn = 7.73e-19; //m-2
  const double m_nn = 1.007825*arma::datum::m_u;
  const double nn_coeff = sigma_nn * std::sqrt(16.0*arma::datum::k/(arma::datum::pi*m_nn));
  return (nn_coeff * nn * ctx.n_0 * std::sqrt(std::abs(T*ctx.T_0)) * ctx.t_0); // ion-neutral collision rate
}

inline arma::vec nu_nn(const arma::vec &nn, const arma::vec &T, const comfi::types::Context &ctx)
{
  const double sigma_nn = 7.73e-19; //m-2
  const double m_nn = 1.007825*arma::datum::m_u;
  const double nn_coeff = sigma_nn * std::sqrt(16.0*arma::datum::k/(arma::datum::pi*m_nn));
  return (nn_coeff * ctx.n_0 * ctx.t_0 * nn % arma::sqrt(arma::abs(T*ctx.T_0)) ); // ion-neutral collision rate
}

inline vcl_vec nu_nn(const vcl_vec &nn, const vcl_vec &T, const comfi::types::Context &ctx)
{
  const double sigma_nn = 7.73e-19; //m-2
  const double m_nn = 1.007825*arma::datum::m_u;
  const double nn_coeff = sigma_nn * std::sqrt(16.0*arma::datum::k/(arma::datum::pi*m_nn));
  return nn_coeff * ctx.n_0 * ctx.t_0  * viennacl::linalg::element_prod(nn, viennacl::linalg::element_sqrt(T*ctx.T_0)); // ion-neutral collision rate
}

inline double nu_in(const double &nn, const double &T, const comfi::types::Context &ctx)
{
  const double sigma_in = 1.16e-18; //m-2
  const double m_in = 1.007276466879*1.007825/(1.007276466879+1.007825);
  const double in_coeff = sigma_in * std::sqrt(8.0*arma::datum::k/(arma::datum::pi*m_in));
  return (in_coeff * nn * ctx.n_0 * std::sqrt(std::abs(T*ctx.T_0)) * ctx.t_0); // ion-neutral collision rate
}

inline arma::vec nu_in(const arma::vec &nn, const arma::vec &T, const comfi::types::Context &ctx)
{
  const double sigma_in = 1.16e-18; //m-2
  const double m_in = 1.007276466879*1.007825/(1.007276466879+1.007825);
  const double in_coeff = sigma_in * std::sqrt(8.0*arma::datum::k/(arma::datum::pi*m_in));
  return (in_coeff * ctx.n_0 * ctx.t_0 * nn % arma::sqrt(arma::abs(T*ctx.T_0)) ); // ion-neutral collision rate
}

inline vcl_vec nu_in(const vcl_vec &nn, const vcl_vec &T, const comfi::types::Context &ctx)
{
  const double sigma_in = 1.16e-18; //m-2
  const double m_in = 1.007276466879*1.007825/(1.007276466879+1.007825);
  const double in_coeff = sigma_in * std::sqrt(8.0*arma::datum::k/(arma::datum::pi*m_in));
  return in_coeff * ctx.n_0 * ctx.t_0  * viennacl::linalg::element_prod(nn, viennacl::linalg::element_sqrt(T*ctx.T_0)); // ion-neutral collision rate
}

inline double nu_en(const double &nn, const double &T, const comfi::types::Context &ctx)
{
  const double sigma_en = 1.0e-19; //m-2
  const double m_en = arma::datum::m_e*1.007825*arma::datum::m_u/(arma::datum::m_e+1.007825*arma::datum::m_u);
  const double en_coeff = sigma_en * std::sqrt(8.0*arma::datum::k/(arma::datum::pi*m_en));
  return (en_coeff * nn * ctx.n_0 * sqrt(std::abs(T*ctx.T_0)) * ctx.t_0); // electron-neutral collision rate
}

inline vcl_vec nu_en(const vcl_vec &nn, const vcl_vec &T, const comfi::types::Context &ctx)
{
  const double sigma_en = 1.0e-19; //m-2
  const double m_en = arma::datum::m_e*1.007825*arma::datum::m_u/(arma::datum::m_e+1.007825*arma::datum::m_u);
  const double en_coeff = sigma_en * std::sqrt(8.0*arma::datum::k/(arma::datum::pi*m_en));
  return en_coeff * ctx.n_0 * ctx.t_0 * viennacl::linalg::element_prod(nn, viennacl::linalg::element_sqrt(T*ctx.T_0)); // electron-neutral collision rate
}

inline double nu_ei(const double &ne, const double &T, const comfi::types::Context &ctx)
{
  const double coloumb_logarithm = 10.0;
  const double r = arma::datum::ec*arma::datum::ec/(4.0*arma::datum::pi*arma::datum::eps_0*arma::datum::k);
  const double sigma_ei = coloumb_logarithm*arma::datum::pi*r*r;
  const double coeff_ei = (4.0/3.0) * sigma_ei * std::sqrt(2.0*arma::datum::k/(arma::datum::pi*arma::datum::m_e));
  return (coeff_ei * ne * ctx.n_0 * std::pow(std::abs(T*ctx.T_0),-1.5) * ctx.t_0); // electron-ion collision rate
}

inline vcl_vec nu_ei(const vcl_vec &ne, const vcl_vec &Te, const comfi::types::Context &ctx)
{
  const double coloumb_logarithm = 10.0;
  const double r = arma::datum::ec*arma::datum::ec/(4.0*arma::datum::pi*arma::datum::eps_0*arma::datum::k);
  const double sigma_ei = coloumb_logarithm*arma::datum::pi*r*r;
  const double coeff_ei = (4.0/3.0) * sigma_ei * std::sqrt(2.0*arma::datum::k/(arma::datum::pi*arma::datum::m_e));
  return coeff_ei * ctx.n_0 * ctx.t_0* viennacl::linalg::element_prod(ne,viennacl::linalg::element_exp(-1.5*viennacl::linalg::element_log(ctx.T_0*Te))); // electron-ion collision rate
}

inline vcl_vec nu_ii(const vcl_vec &Np, const vcl_vec &Tp, const comfi::types::Context &ctx)
{
  const double coloumb_logarithm = 10.0;
  const double r = arma::datum::ec*arma::datum::ec/(4.0*arma::datum::pi*arma::datum::eps_0*arma::datum::k); //without T
  const double sigma_ii = coloumb_logarithm*arma::datum::pi*r*r;
  const double coeff_ii = (4.0/3.0) * sigma_ii * std::sqrt(2.0*arma::datum::k/(arma::datum::pi*arma::datum::m_p));
  return ctx.t_0 * ctx.n_0 * coeff_ii * viennacl::linalg::element_prod(Np, viennacl::linalg::element_exp(-1.5*viennacl::linalg::element_log(ctx.T_0*Tp))); // electron-ion collision rate
}

inline double nu_ii(const double &Np, const double &Tp, const comfi::types::Context &ctx)
{
  const double coloumb_logarithm = 10.0;
  const double r = arma::datum::ec*arma::datum::ec/(4.0*arma::datum::pi*arma::datum::eps_0*arma::datum::k); //without T
  const double sigma_ii = coloumb_logarithm*arma::datum::pi*r*r;
  const double coeff_ii = (4.0/3.0) * sigma_ii * std::sqrt(2.0*arma::datum::k/(arma::datum::pi*arma::datum::m_p));
  return coeff_ii * ctx.t_0 * Np * ctx.n_0 * std::exp(-1.5*std::log(ctx.T_0*Tp)); // electron-ion collision rate
}

inline vcl_vec resistivity(const vcl_vec &Np, const vcl_vec &Nn, const vcl_vec &Tp, const vcl_vec &Tn, const comfi::types::Context &ctx)
{
  return (ctx.m_e/(ctx.q*ctx.q))*viennacl::linalg::element_div((nu_ei(Np, Tp, ctx)+nu_en(Nn, 0.5*(Tp+Tn), ctx)),Np);
}
inline double resistivity(const double &Np, const double &Nn, const double &Tp, const double &Tn, const comfi::types::Context &ctx)
{
  return (ctx.m_e/(ctx.q*ctx.q)) * (nu_ei(Np, Tp, ctx)+nu_en(Nn, 0.5*(Tp+Tn), ctx))/Np;
}

inline double thermalvel(const double &T, const comfi::types::Context &ctx)
{
  return std::sqrt(2.0*arma::datum::k/arma::datum::m_p) * std::sqrt(T*ctx.T_0) / ctx.V_0;
}

inline vcl_vec thermalvel(const vcl_vec &T, const comfi::types::Context &ctx)
{
  using namespace viennacl::linalg;
  return std::sqrt(2.0*arma::datum::k/arma::datum::m_p) * element_sqrt(T*ctx.T_0) / ctx.V_0;
}

inline double kappa_n(const double &Tp, const double &Tn, const double &Np, const double &Nn, const comfi::types::Context &ctx)
{
  return Nn * thermalvel(Tn, ctx) * thermalvel(Tn, ctx) / nu_nn(Nn, Tn, ctx);
}

inline vcl_vec kappa_n(const vcl_vec &Tp, const vcl_vec &Tn, const vcl_vec &Np, const vcl_vec &Nn, const comfi::types::Context &ctx)
{
  using namespace viennacl::linalg;
  const vcl_vec v2 = element_prod(thermalvel(Tn, ctx), thermalvel(Tn, ctx));
  const vcl_vec nv2 = element_prod(Nn, v2);
  const vcl_vec nunn = nu_nn(Nn, Tn, ctx);
  return element_div(nv2, nunn);
}

inline double kappa_p(const double &Tp, const double &Tn, const double &Np, const double &Nn, const comfi::types::Context &ctx)
{
  return Np * thermalvel(Tp, ctx) * thermalvel(Tp, ctx) / (nu_in(Nn, 0.5*(Tp+Tn), ctx) + nu_ii(Np, Tp, ctx));
}

inline vcl_vec kappa_p(const vcl_vec &Tp, const vcl_vec &Tn, const vcl_vec &Np, const vcl_vec &Nn, const comfi::types::Context &ctx)
{
  using namespace viennacl::linalg;
  const vcl_vec v2 = element_prod(thermalvel(Tp, ctx), thermalvel(Tp, ctx));
  const vcl_vec nv2 = element_prod(Np, v2);
  const vcl_vec nuin = nu_in(Nn, 0.5*(Tp+Tn), ctx);
  const vcl_vec nuii = nu_ii(Np, Tp, ctx);
  return element_div(nv2, nuin+nuii);
}

/*!
 * \brief Calculate recombination coefficient for ions to recombine to neutrals.
 * \param Tp Ion temperature
 * \return Vector of normalized volumetric recombination rate.
 */
inline vcl_vec recomb_coeff(const vcl_vec &T, const comfi::types::Context &ctx)
{
  using namespace viennacl::linalg;
  static const vcl_vec one = viennacl::scalar_vector<double>(T.size(), 1.0);
  const vcl_vec beta = 0.6 * 13.6 * 1.6021766208e-19 * element_div(one, arma::datum::k*ctx.T_0*T);
  vcl_vec coeff = 0.4288*one + 0.5*element_log(beta) + 0.4698*element_pow(beta, -(1.0/3.0)*one);
  coeff = element_prod(5.20e-20*element_sqrt(beta), coeff);
  return ctx.t_0 * ctx.n_0 * coeff;
}

inline double recomb_coeff(const double &T, const comfi::types::Context &ctx)
{
  const double beta = 0.6 * 13.6 * 1.6021766208e-19 / (arma::datum::k*ctx.T_0*T);
  return ctx.t_0 * ctx.n_0 * 5.20e-20 * std::sqrt(beta) * (0.4288+std::log(beta)+0.4698*std::pow(beta, -(1.0/3.0)));
}

/*!
 * \brief Calculate normalized ionization coefficient for neutrals to ionize due to collisions (Draine p 134)
 * \param Tn Neutral temperature
 * \return Vector of normalized volumetric ionization rate
 */
inline vcl_vec ionization_coeff(const vcl_vec &T, const comfi::types::Context &ctx)
{
  using namespace viennacl::linalg;
  static const vcl_vec one = viennacl::scalar_vector<double>(T.size(), 1.0);
  const vcl_vec beta = 0.6 * 13.6 * 1.6021766208e-19 * element_div(one, arma::datum::k*ctx.T_0*T);
  const vcl_vec coeff = element_div(2.34e-14*element_exp(-beta), element_sqrt(beta));
  return ctx.t_0 * ctx.n_0 * coeff;

}

inline double ionization_coeff(const double &T, const comfi::types::Context &ctx)
{
  const double beta = 0.6 * 13.6 * 1.6021766208e-19 / (arma::datum::k*ctx.T_0*T);
  return ctx.t_0 * ctx.n_0 * 2.34e-14 * std::exp(-beta) / std::sqrt(beta);
}
} // namespace sol

} // namespace mhdsim

/*!
 * \brief Find scalar vector index.
 * \param i Horizontal index.
 * \param j Vertical index
 * \return Scalar vector index.
 */
inline arma::uword inds(const arma::uword &i, const arma::uword &j, comfi::types::Context &ctx)
{
  return i+j*ctx.nx;
}

/*
vim: tabstop=2
vim: shiftwidth=2
vim: smarttab
vim: expandtab
*/