#include "mtf/SSM/SL3.h"
#include "mtf/SSM/HomographyEstimator.h"
#include "mtf/Utilities/warpUtils.h"
#include "mtf/Utilities/miscUtils.h"
#include "mtf/Utilities/excpUtils.h"

#include <unsupported/Eigen/MatrixFunctions>
#include "opencv2/calib3d/calib3d.hpp"
#include <boost/random/random_device.hpp>
#include <boost/random/seed_seq.hpp>

#define VALIDATE_SL3_WARP(warp) \
	assert(warp.determinant() == 1.0);

#define SL3_NORMALIZED_BASIS 0
#define SL3_ITERATIVE_SAMPLE_MEAN 1
#define SL3_SAMPLE_MEAN_MAX_ITERS 10
#define SL3_SAMPLE_MEAN_EPS 1e-4
#define SL3_DEBUG_MODE 0

#ifndef SL3_MAX_VALID_VAL
#define SL3_MAX_VALID_VAL 1e10
#endif
#define is_unbounded(eig_mat) (eig_mat.array().cwiseAbs().eval() > SL3_MAX_VALID_VAL).any()

_MTF_BEGIN_NAMESPACE

SL3Params::SL3Params(const SSMParams *ssm_params,
bool _normalized_init, bool _iterative_sample_mean,
int _sample_mean_max_iters, double _sample_mean_eps,
bool _debug_mode) :
SSMParams(ssm_params),
normalized_init(_normalized_init),
iterative_sample_mean(_iterative_sample_mean),
sample_mean_max_iters(_sample_mean_max_iters),
sample_mean_eps(_sample_mean_eps),
debug_mode(_debug_mode){}

//! copy/default constructor
SL3Params::SL3Params(const SL3Params *params) :
SSMParams(params),
normalized_init(SL3_NORMALIZED_BASIS),
iterative_sample_mean(SL3_ITERATIVE_SAMPLE_MEAN),
sample_mean_max_iters(SL3_SAMPLE_MEAN_MAX_ITERS),
sample_mean_eps(SL3_SAMPLE_MEAN_EPS),
debug_mode(SL3_DEBUG_MODE){
	if(params){
		normalized_init = params->normalized_init;
		iterative_sample_mean = params->iterative_sample_mean;
		sample_mean_max_iters = params->sample_mean_max_iters;
		sample_mean_eps = params->sample_mean_eps;
		debug_mode = params->debug_mode;
	}
}

SL3::SL3(
	const ParamType *_params) :
	ProjectiveBase(_params), params(_params){

	printf("\n");
	printf("Using SL3 SSM with:\n");
	printf("resx: %d\n", resx);
	printf("resy: %d\n", resy);
	printf("normalized_init: %d\n", params.normalized_init);
	printf("iterative_sample_mean: %d\n", params.iterative_sample_mean);
	printf("sample_mean_max_iters: %d\n", params.sample_mean_max_iters);
	printf("sample_mean_eps: %f\n", params.sample_mean_eps);
	printf("debug_mode: %d\n", params.debug_mode);

	name = "sl3";
	state_size = 8;
	curr_state.resize(state_size);

	log_fname = log_fname;

	lie_alg_mat = Matrix3d::Zero();
	warp_mat = Matrix3d::Identity();

	lieAlgBasis[0] <<
		1, 0, 0,
		0, -1, 0,
		0, 0, 0;
	lieAlgBasis[1] <<
		0, 0, 0,
		0, -1, 0,
		0, 0, 1;
	lieAlgBasis[2] <<
		0, -1, 0,
		1, 0, 0,
		0, 0, 0;
	lieAlgBasis[3] <<
		0, 1, 0,
		1, 0, 0,
		0, 0, 0;
	lieAlgBasis[4] <<
		0, 0, 1,
		0, 0, 0,
		0, 0, 0;
	lieAlgBasis[5] <<
		0, 0, 0,
		0, 0, 1,
		0, 0, 0;
	lieAlgBasis[6] <<
		0, 0, 0,
		0, 0, 0,
		1, 0, 0;
	lieAlgBasis[7] <<
		0, 0, 0,
		0, 0, 0,
		0, 1, 0;

	utils::getNormUnitSquarePts(norm_pts, norm_corners, resx, resy,
		-static_cast<double>(resx) / 2.0, -static_cast<double>(resy) / 2.0,
		static_cast<double>(resx) / 2.0, static_cast<double>(resy) / 2.0);
	utils::homogenize(norm_corners, norm_corners_hm);
	utils::homogenize(norm_pts, norm_pts_hm);

	init_corners = norm_corners;
	init_corners_hm = norm_corners_hm;
	init_pts = norm_pts;
	init_pts_hm = norm_pts_hm;

	if(params.debug_mode){
#ifdef _WIN32
		FILE *fid;
		errno_t err;
		if((err = fopen_s(&fid, log_fname, "r")) != 0) {
			throw utils::InvalidArgument(cv_format("SL3 :: Log file %s could not be opened successfully : %s\n",
				log_fname, strerror(err)));
		}
#else
		FILE *fid = fopen(log_fname, "w");
		if(!fid){
			throw utils::InvalidArgument(cv_format("SL3 :: Log file %s could not be opened successfully\n", log_fname));
		}
#endif
		fclose(fid);
	}
}

void SL3::setState(const VectorXd &ssm_state){
	validate_ssm_state(ssm_state);
	//utils::printMatrix(ssm_state.transpose(), " SL3::setState :: ssm_state");
	if(!ssm_state.allFinite()){
		utils::printMatrix(ssm_state.transpose(), "ssm_state");
		throw mtf::utils::InvalidTrackerState("SL3::setState::Invalid state provided");
	}
	curr_state = ssm_state;
	getWarpFromState(curr_warp, curr_state);
	curr_pts_hm.noalias() = curr_warp * init_pts_hm;
	curr_corners_hm.noalias() = curr_warp * init_corners_hm;
	utils::dehomogenize(curr_pts_hm, curr_pts);
	utils::dehomogenize(curr_corners_hm, curr_corners);
	if(params.debug_mode){
		utils::printMatrixToFile(curr_warp, "setState::curr_warp", log_fname);
		utils::printMatrixToFile(curr_corners, "setState::curr_corners", log_fname);
	}
}

void SL3::setCorners(const CornersT& corners){
	if(!corners.allFinite() || is_unbounded(corners)){
		utils::printMatrix(corners, "corners");
		throw mtf::utils::InvalidTrackerState("SL3::setCorners::Invalid corners provided");
	}
	curr_corners = corners;
	curr_warp = utils::computeHomographyDLT(norm_corners, curr_corners);
	double warp_det = curr_warp.determinant();
	if(!warp_mat.allFinite() || is_unbounded(warp_mat) || warp_det == 0 || std::isnan(warp_det) || std::isinf(warp_det)){
		printf("SL3::Cannot set SSM to the provided corners as the corresponding warp matrix is invalid: \n");
		utils::printMatrix(corners, "corners");
		utils::printMatrix(curr_warp, "warp");
		utils::printScalar(warp_det, "warp_det");
		throw mtf::utils::InvalidTrackerState("SL3::setCorners :: Cannot set SSM to the provided points as the corresponding warp matrix is invalid");
	}
	utils::homogenize(curr_corners, curr_corners_hm);
	curr_pts_hm = curr_warp * norm_pts_hm;
	utils::dehomogenize(curr_pts_hm, curr_pts);
	curr_warp = curr_warp / cbrt(warp_det);
	if(params.normalized_init){
		getStateFromWarp(curr_state, curr_warp);
	} else{
		init_corners = curr_corners;
		init_corners_hm = curr_corners_hm;
		init_pts = curr_pts;
		init_pts_hm = curr_pts_hm;
		curr_warp = Matrix3d::Identity();
		curr_state.fill(0);
	}
	if(params.debug_mode){
		utils::printMatrixToFile(init_pts, "setCorners::init_pts", log_fname);
		utils::printMatrixToFile(init_corners, "setCorners::init_corners", log_fname);
		utils::printMatrixToFile(curr_warp, "setCorners::curr_warp", log_fname);
	}
}

void SL3::compositionalUpdate(const VectorXd& state_update){
	validate_ssm_state(state_update);

	getWarpFromState(warp_update_mat, state_update);
	curr_warp = curr_warp * warp_update_mat;
	if(!curr_warp.allFinite() || is_unbounded(curr_warp)){
		utils::printMatrix(state_update.transpose(), "state_update");
		utils::printMatrix(curr_warp, "curr_warp");
		utils::printMatrix(warp_update_mat, "warp_update_mat");
		throw mtf::utils::InvalidTrackerState("SL3::compositionalUpdate::Invalid state update provided");

	}
	curr_pts_hm.noalias() = curr_warp * init_pts_hm;
	curr_corners_hm.noalias() = curr_warp * init_corners_hm;

	utils::dehomogenize(curr_pts_hm, curr_pts);
	utils::dehomogenize(curr_corners_hm, curr_corners);

	getStateFromWarp(curr_state, curr_warp);

}

void SL3::getWarpFromState(Matrix3d &warp_mat,
	const VectorXd& ssm_state){
	validate_ssm_state(ssm_state);
	assert(ssm_state.size() == 8);
	if(!ssm_state.allFinite() || is_unbounded(ssm_state)){
		utils::printMatrix(ssm_state.transpose(), "ssm_state");
		throw mtf::utils::InvalidTrackerState("SL3::getWarpFromState::Invalid state provided");
	}
	getLieAlgMatFromState(lie_alg_mat, ssm_state);
	if(!lie_alg_mat.allFinite() || is_unbounded(lie_alg_mat)){
		utils::printMatrix(ssm_state.transpose(), "ssm_state");
		utils::printMatrix(lie_alg_mat, "lie_alg_mat");
		throw mtf::utils::InvalidTrackerState("SL3::getWarpFromState::Invalid sl3 matrix corresponds to the given state");
	}
	warp_mat = lie_alg_mat.exp();
}

void SL3::getStateFromWarp(VectorXd &state_vec,
	const Matrix3d& warp_mat){
	validate_ssm_state(state_vec);
	double warp_det = warp_mat.determinant();
	if(!warp_mat.allFinite() || is_unbounded(warp_mat) || warp_det == 0 || std::isnan(warp_det) || std::isinf(warp_det)){
		utils::printMatrix(warp_mat, "warp_mat");
		utils::printScalar(warp_det, "warp_det");
		throw mtf::utils::InvalidTrackerState("SL3::getStateFromWarp :: Invalid warp matrix provided");
	}
	Matrix3d norm_warp_mat = warp_mat / cbrt(warp_det);
	if(!norm_warp_mat.allFinite() || is_unbounded(norm_warp_mat)){
		utils::printMatrix(warp_mat, "warp_mat");
		utils::printScalar(warp_det, "warp_det");
		utils::printMatrix(norm_warp_mat, "norm_warp_mat");
		throw mtf::utils::InvalidTrackerState("SL3::getStateFromWarp :: Invalid normalized warp matrix found");
	}
	lie_alg_mat = norm_warp_mat.log();
	getStateFromLieAlgMat(state_vec, lie_alg_mat);
}

void SL3::getLieAlgMatFromState(Matrix3d& lie_alg_mat,
	const VectorXd& ssm_state){
	validate_ssm_state(ssm_state);
	assert(ssm_state.size() == state_size);
	lie_alg_mat(0, 0) = ssm_state(0);
	lie_alg_mat(0, 1) = ssm_state(3) - ssm_state(2);
	lie_alg_mat(0, 2) = ssm_state(4);
	lie_alg_mat(1, 0) = ssm_state(3) + ssm_state(2);
	lie_alg_mat(1, 1) = -ssm_state(1) - ssm_state(0);
	lie_alg_mat(1, 2) = ssm_state(5);
	lie_alg_mat(2, 0) = ssm_state(6);
	lie_alg_mat(2, 1) = ssm_state(7);
	lie_alg_mat(2, 2) = ssm_state(1);
}

void SL3::getStateFromLieAlgMat(VectorXd &ssm_state,
	const Matrix3d& lie_alg_mat){
	validate_ssm_state(ssm_state);
	ssm_state(0) = lie_alg_mat(0, 0);
	ssm_state(1) = -lie_alg_mat(1, 1) - ssm_state(0);
	ssm_state(2) = (lie_alg_mat(1, 0) - lie_alg_mat(0, 1)) / 2.0;
	ssm_state(3) = (lie_alg_mat(1, 0) + lie_alg_mat(0, 1)) / 2.0;
	ssm_state(4) = lie_alg_mat(0, 2);
	ssm_state(5) = lie_alg_mat(1, 2);
	ssm_state(6) = lie_alg_mat(2, 0);
	ssm_state(7) = lie_alg_mat(2, 1);
	if(params.debug_mode){
		utils::printMatrixToFile(ssm_state.transpose(), "getStateFromLieAlgMat :: ssm_state", log_fname);
		utils::printMatrixToFile(lie_alg_mat, "getStateFromLieAlgMat :: lie_alg_mat", log_fname);
	}
}

void SL3::initializeSampler(const VectorXd &state_sigma,
	const VectorXd &state_mean){
	if(state_sigma.size() != 8){
		throw utils::InvalidArgument(
			cv::format("SL3::initializeSampler :: SSM sigma has invalid size: %d\n",
			state_sigma.size()));
	}
	covariance_mat = state_sigma.asDiagonal();
	printf("Using SL3 sampler with sigma:\n");
	utils::printMatrix(state_sigma.transpose(), nullptr, "%e");

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

	boost::random_device r;
	boost::random::seed_seq seed{ r(), r(), r(), r(), r(), r(), r(), r() };
	rand_gen[0] = SampleGenT(seed);
	rand_dist[0] = SampleDistT(0, 1);

	if(params.debug_mode){
		utils::printMatrixToFile(covariance_mat, "covariance_mat", log_fname);
	}
	is_initialized.sampler = true;
}

void SL3::setSampler(const VectorXd &state_sigma,
	const VectorXd &state_mean){
	assert(state_sigma.size() == state_size);
	covariance_mat = state_sigma.asDiagonal();
}

VectorXd SL3::getSamplerSigma(){
	VectorXd sampler_sigma = covariance_mat.diagonal();
	return sampler_sigma;
}

void SL3::generatePerturbation(VectorXd &perturbation){
	assert(perturbation.size() == state_size);
	Vector8d rand_vec;
	for(unsigned int state_id = 0; state_id < 8; state_id++){
		rand_vec(state_id) = rand_dist[0](rand_gen[0]);
	}
	perturbation = covariance_mat*rand_vec;
}

// use Random Walk model to generate perturbed sample
void SL3::compositionalRandomWalk(VectorXd &perturbed_state,
	const VectorXd &base_state){
	generatePerturbation(state_perturbation);
	perturbed_state = base_state + state_perturbation;
}
// use first order Auto Regressive model to generate perturbed sample
//void SL3::compositionalAutoRegression1(VectorXd &perturbed_state, VectorXd &perturbed_ar,
//	const VectorXd &base_state, const VectorXd &base_ar, double a){
//	ProjWarpT warp_perturbation, sl3_perturbation, sl3_base_ar, base_warp;
//	generatePerturbation(state_perturbation);
//	getLieAlgMatFromState(warp_perturbation, state_perturbation);
//	getLieAlgMatFromState(sl3_base_ar, base_ar);
//	getWarpFromState(base_warp, base_state);
//	sl3_perturbation = warp_perturbation + a*sl3_base_ar;
//	ProjWarpT SL3_perturbation = sl3_perturbation.exp();
//	ProjWarpT perturbed_warp = base_warp*SL3_perturbation;
//	ProjWarpT base_warp_inv = base_warp.inverse();
//	if(params.debug_mode){
//		utils::printMatrixToFile(state_perturbation.transpose(), "rand_perturbation", log_fname);
//		utils::printMatrixToFile(warp_perturbation.transpose(), "warp_perturbation", log_fname);
//		utils::printMatrixToFile(sl3_perturbation, "sl3_perturbation", log_fname);
//		utils::printMatrixToFile(base_ar.transpose(), "base_ar", log_fname);
//		utils::printMatrixToFile(sl3_base_ar, "sl3_base_ar", log_fname);
//		utils::printMatrixToFile(SL3_perturbation, "SL3_perturbation", log_fname);
//		utils::printMatrixToFile(base_warp, "base_warp", log_fname);
//		utils::printMatrixToFile(base_warp_inv, "base_warp_inv", log_fname);
//		utils::printMatrixToFile(perturbed_warp, "perturbed_warp", log_fname);
//		utils::printMatrixToFile(perturbed_state.transpose(), "perturbed_state", log_fname);
//		utils::printMatrixToFile(perturbed_ar.transpose(), "perturbed_ar", log_fname);
//	}
//	ProjWarpT SL3_perturbed_ar = a*(base_warp_inv*perturbed_warp);
//	ProjWarpT sl3_perturbed_ar = SL3_perturbed_ar.log();
//	if(params.debug_mode){
//		utils::printMatrixToFile(sl3_perturbed_ar, "sl3_perturbed_ar", log_fname);
//	}
//	getStateFromWarp(perturbed_state, perturbed_warp);
//	getStateFromLieAlgMat(perturbed_ar, sl3_perturbed_ar);
//
//}
void SL3::compositionalAutoRegression1(VectorXd &perturbed_state, VectorXd &perturbed_ar,
	const VectorXd &base_state, const VectorXd &base_ar, double a){
	ProjWarpT sl3_perturbation, lie_alg_base_ar, base_warp;
	generatePerturbation(state_perturbation);
	getLieAlgMatFromState(sl3_perturbation, state_perturbation);
	getLieAlgMatFromState(lie_alg_base_ar, base_ar);
	getWarpFromState(base_warp, base_state);
	ProjWarpT SL3_perturbation = (a*lie_alg_base_ar + sl3_perturbation).exp();
	ProjWarpT perturbed_warp = base_warp*SL3_perturbation;
	ProjWarpT base_warp_inv = base_warp.inverse();
	ProjWarpT lie_alg_perturbed_ar = a*(base_warp_inv*perturbed_warp).log();
	getStateFromWarp(perturbed_state, perturbed_warp);
	getStateFromLieAlgMat(perturbed_ar, lie_alg_perturbed_ar);
	if(params.debug_mode){
		utils::printMatrixToFile(state_perturbation.transpose(), "state_perturbation", log_fname);
		utils::printMatrixToFile(sl3_perturbation, "sl3_perturbation", log_fname);
		utils::printMatrixToFile(base_ar.transpose(), "base_ar", log_fname);
		utils::printMatrixToFile(lie_alg_base_ar, "lie_alg_base_ar", log_fname);
		utils::printMatrixToFile(SL3_perturbation, "SL3_perturbation", log_fname);
		utils::printMatrixToFile(base_warp, "base_warp", log_fname);
		utils::printMatrixToFile(perturbed_warp, "perturbed_warp", log_fname);
		utils::printMatrixToFile(lie_alg_perturbed_ar, "lie_alg_perturbed_ar", log_fname);
	}
}

void SL3::getInitPixGrad(Matrix2Xd &ssm_grad, int pt_id) {
	double x = init_pts(0, pt_id);
	double y = init_pts(1, pt_id);

	ssm_grad <<
		x, -x, -y, y, 1, 0, -x*x, -x*y,
		-x, -2*y, x, x, 0, 1, -y*x, -y*y;
}

void SL3::getCurrPixGrad(Matrix2Xd &ssm_grad, int pt_id) {
	double x = init_pts(0, pt_id);
	double y = init_pts(1, pt_id);

	double curr_x = curr_pts(0, pt_id);
	double curr_y = curr_pts(1, pt_id);

	double xx = x*curr_x, xy = x*curr_y, yy = y*curr_y, yx = y*curr_x;

	Matrix3d curr_warp;
	getWarpFromState(curr_warp, curr_state);
	double a1 = curr_warp(0, 0), a2 = curr_warp(0, 1), a3 = curr_warp(0, 2);
	double a4 = curr_warp(1, 0), a5 = curr_warp(1, 1), a6 = curr_warp(1, 2);
	double a7 = curr_warp(2, 0), a8 = curr_warp(2, 1), a9 = curr_warp(2, 2);

	double p11 = a1*x - a2*y - a7*xx + a8*yx;
	double p12 = a1*y - a7*yx;
	double p13 = a1 - a7*curr_x;
	double p14 = a2*x - a8*xx;
	double p15 = a2*y - a3 - a8*yx + a9*curr_x;
	double p16 = a2 - a8*curr_x;
	double p17 = a3*x - a9*xx;
	double p18 = a3*y - a9*yx;

	double p21 = a4*x - a5*y - a7*xy + a8*yy;
	double p22 = a4*y - a7*yy;
	double p23 = a4 - a7*curr_y;
	double p24 = a5*x - a8*xy;
	double p25 = a5*y - a6 - a8*yy + a9*curr_y;
	double p26 = a5 - a8*curr_y;
	double p27 = a6*x - a9*xy;
	double p28 = a6*y - a9*yy;

	ssm_grad <<
		p11, p12, p13, p14, p15, p16, p17, p18,
		p21, p22, p23, p24, p25, p26, p27, p28;

	double inv_d = 1.0 / curr_pts_hm(2, pt_id);
	ssm_grad *= inv_d;
}

void SL3::cmptInitPixJacobian(MatrixXd &dI_dp,
	const PixGradT &dI_dw){
	validate_ssm_jacobian(dI_dp, dI_dw);

	unsigned int ch_pt_id = 0;
	for(unsigned int pt_id = 0; pt_id < n_pts; ++pt_id){
		double x = init_pts(0, pt_id);
		double y = init_pts(1, pt_id);

		for(unsigned int ch_id = 0; ch_id < n_channels; ++ch_id){
			double Ix = dI_dw(ch_pt_id, 0);
			double Iy = dI_dw(ch_pt_id, 1);

			double Ixx = Ix * x;
			double Iyy = Iy * y;
			double Ixy = Ix * y;
			double Iyx = Iy * x;

			dI_dp(ch_pt_id, 0) = Ixx - Iyx;
			dI_dp(ch_pt_id, 1) = -(2 * Iyy + Ixx);
			dI_dp(ch_pt_id, 2) = Iyx - Ixy;
			dI_dp(ch_pt_id, 3) = Iyx + Ixy;
			dI_dp(ch_pt_id, 4) = Ix;
			dI_dp(ch_pt_id, 5) = Iy;
			dI_dp(ch_pt_id, 6) = -Ixx*x - Iyy*x;
			dI_dp(ch_pt_id, 7) = -Ixx*y - Iyy*y;

			++ch_pt_id;
		}
	}
}

void SL3::cmptWarpedPixJacobian(MatrixXd &dI_dp,
	const PixGradT &dI_dw) {
	validate_ssm_jacobian(dI_dp, dI_dw);

	double a00 = curr_warp(0, 0);
	double a01 = curr_warp(0, 1);
	double a10 = curr_warp(1, 0);
	double a11 = curr_warp(1, 1);
	double a20 = curr_warp(2, 0);
	double a21 = curr_warp(2, 1);

	unsigned int ch_pt_id = 0;
	for(unsigned int pt_id = 0; pt_id < n_pts; ++pt_id) {
		double w_x = curr_pts(0, pt_id);
		double w_y = curr_pts(1, pt_id);

		double D = curr_pts_hm(2, pt_id);
		double inv_det = 1.0 / D;

		double dwx_dx = (a00 - a20*w_x);
		double dwx_dy = (a01 - a21*w_x);
		double dwy_dx = (a10 - a20*w_y);
		double dwy_dy = (a11 - a21*w_y);

		double x = init_pts(0, pt_id);
		double y = init_pts(1, pt_id);

		//double Ix = pix_grad(pt_id, 0);
		//double Iy = pix_grad(pt_id, 1);
		for(unsigned int ch_id = 0; ch_id < n_channels; ++ch_id){
			double Ix = (dwx_dx*dI_dw(ch_pt_id, 0) + dwy_dx*dI_dw(ch_pt_id, 1))*inv_det;
			double Iy = (dwx_dy*dI_dw(ch_pt_id, 0) + dwy_dy*dI_dw(ch_pt_id, 1))*inv_det;

			double Ixx = Ix * x;
			double Ixy = Ix * y;
			double Iyy = Iy * y;
			double Iyx = Iy * x;

			dI_dp(ch_pt_id, 0) = Ixx - Iyx;
			dI_dp(ch_pt_id, 1) = -(2 * Iyy + Ixx);
			dI_dp(ch_pt_id, 2) = Iyx - Ixy;
			dI_dp(ch_pt_id, 3) = Iyx + Ixy;
			dI_dp(ch_pt_id, 4) = Ix;
			dI_dp(ch_pt_id, 5) = Iy;
			dI_dp(ch_pt_id, 6) = -Ixx*x - Iyy*x;
			dI_dp(ch_pt_id, 7) = -Ixx*y - Iyy*y;

			++ch_pt_id;
		}
	}
}

void SL3::cmptPixJacobian(MatrixXd &dI_dp,
	const PixGradT &pix_jacobian){
	validate_ssm_jacobian(dI_dp, pix_jacobian);

	for(unsigned int pt_id = 0; pt_id < n_pts; ++pt_id){
		double x = init_pts(0, pt_id);
		double y = init_pts(1, pt_id);
		double Ix = pix_jacobian(pt_id, 0);
		double Iy = pix_jacobian(pt_id, 1);

		double curr_x = curr_pts(0, pt_id);
		double curr_y = curr_pts(1, pt_id);
		double inv_d = 1.0 / curr_pts_hm(2, pt_id);

		double Ixx = Ix * x;
		double Iyy = Iy * y;
		double Ixy = Ix * y;
		double Iyx = Iy * x;

		dI_dp(pt_id, 0) = (Ixx - Iyy) * inv_d;
		dI_dp(pt_id, 1) = Ixy * inv_d;
		dI_dp(pt_id, 2) = Ix * inv_d;
		dI_dp(pt_id, 3) = Iyx * inv_d;
		dI_dp(pt_id, 4) = (Ix*curr_x + Iy*(y + curr_y)) * inv_d;
		dI_dp(pt_id, 5) = Iy * inv_d;
		dI_dp(pt_id, 6) = (-Ixx*curr_x - Iyx*curr_y) * inv_d;
		dI_dp(pt_id, 7) = (-Ixy*curr_x - Iyy*curr_y) * inv_d;
	}
	//dI_dp.array().colwise() /= curr_pts_hm.array().row(2).transpose();
}

void SL3::cmptApproxPixJacobian(MatrixXd &dI_dp,
	const PixGradT &pix_jacobian) {
	validate_ssm_jacobian(dI_dp, pix_jacobian);

	curr_warp /= curr_warp(2, 2);

	double h00_plus_1 = curr_warp(0, 0);
	double h01 = curr_warp(0, 1);
	double h10 = curr_warp(1, 0);
	double h11_plus_1 = curr_warp(1, 1);
	double h20 = curr_warp(2, 0);
	double h21 = curr_warp(2, 1);

	for(unsigned int pt_id = 0; pt_id < n_pts; ++pt_id){

		double Nx = curr_pts_hm(0, pt_id);
		double Ny = curr_pts_hm(1, pt_id);
		double D = curr_pts_hm(2, pt_id);
		double D_sqr_inv = 1.0 / (D*D);

		double a = (h00_plus_1*D - h21*Nx) * D_sqr_inv;
		double b = (h01*D - h21*Nx) * D_sqr_inv;
		double c = (h10*D - h20*Ny) * D_sqr_inv;
		double d = (h11_plus_1*D - h21*Ny) * D_sqr_inv;
		double inv_det = 1.0 / ((a*d - b*c)*D);

		double x = init_pts(0, pt_id);
		double y = init_pts(1, pt_id);

		double curr_x = curr_pts(0, pt_id);
		double curr_y = curr_pts(1, pt_id);

		double Ix = pix_jacobian(pt_id, 0);
		double Iy = pix_jacobian(pt_id, 1);

		double Ixx = Ix * x;
		double Ixy = Ix * y;
		double Iyy = Iy * y;
		double Iyx = Iy * x;

		double factor1 = b*curr_y - d*curr_x;
		double factor2 = c*curr_x - a*curr_y;

		dI_dp(pt_id, 0) = (Ixx*d + Ixy*b - Iyx*c - Iyy*a) * inv_det;
		dI_dp(pt_id, 1) = (Ixy*d - Iyy*c) * inv_det;
		dI_dp(pt_id, 2) = (Ix*d - Iy*c) * inv_det;
		dI_dp(pt_id, 3) = (Iyx*a - Ixx*b) * inv_det;
		dI_dp(pt_id, 4) = (Iyy*a - Ix*factor1 - Ixy*b - Iy*factor2) * inv_det;
		dI_dp(pt_id, 5) = (Iy*a - Ix*b) * inv_det;
		dI_dp(pt_id, 6) = (Ixx*factor1 + Iyx*factor2) * inv_det;
		dI_dp(pt_id, 7) = (Ixy*factor1 + Iyy*factor2) * inv_det;

		//dI_dp(i, 0) = (Ix*(d*x + b*y) - Iy*(c*x + a*y)) * inv_det;
		//dI_dp(i, 1) = (Ix*d*y - Iy*c*y) * inv_det;
		//dI_dp(i, 2) = (Ix*d - Iy*c) * inv_det;
		//dI_dp(i, 3) = (Iy*a*x - Ix*b*x) * inv_det;
		//dI_dp(i, 4) = (Ix*(d*curr_x - b*(y + curr_y)) + Iy*(-c*curr_x + a*(y + curr_y))) * inv_det;
		//dI_dp(i, 5) = (Iy*a - Ix*b) * inv_det;
		//dI_dp(i, 6) = (Ix*(-d*x*curr_x + b*x*curr_y) + Iy*(c*x*curr_x - a*x*curr_y)) * inv_det;
		//dI_dp(i, 7) = (Ix*(-d*y*curr_x + b*y*curr_y) + Iy*(c*y*curr_x - a*y*curr_y)) * inv_det;
	}
}

void SL3::estimateWarpFromCorners(VectorXd &state_update, const Matrix24d &in_corners,
	const Matrix24d &out_corners){
	validate_ssm_state(state_update);
	Matrix3d warp_update_mat = utils::computeHomographyDLT(in_corners, out_corners);
	getStateFromWarp(state_update, warp_update_mat);
}

void SL3::estimateWarpFromPts(VectorXd &state_update, vector<uchar> &mask,
	const vector<cv::Point2f> &in_pts, const vector<cv::Point2f> &out_pts,
	const EstimatorParams &est_params){
	cv::Mat warp_mat_cv = estimateHomography(in_pts, out_pts, mask, est_params);
	utils::copyCVToEigen<double, Matrix3d>(warp_mat, warp_mat_cv);
	getStateFromWarp(state_update, warp_mat);
}

void SL3::estimateMeanOfSamples(VectorXd &sample_mean,
	const std::vector<VectorXd> &samples, int n_samples){
	if(params.iterative_sample_mean){
		vector<ProjWarpT> lie_group_samples;
		lie_group_samples.resize(n_samples);
		// convert state vectors to SL3 matrices
		for(int sample_id = 0; sample_id < n_samples; ++sample_id){
			getWarpFromState(lie_group_samples[sample_id], samples[sample_id]);
		}
		ProjWarpT lie_group_mean = lie_group_samples[0];
		ProjWarpT lie_group_mean_inv = lie_group_mean.inverse();
		for(int iter_id = 0; iter_id < params.sample_mean_max_iters; ++iter_id){
			ProjWarpT lie_algebra_mean = ProjWarpT::Zero();
			for(int sample_id = 0; sample_id < n_samples; ++sample_id){
				lie_algebra_mean += (lie_group_mean_inv*lie_group_samples[sample_id]).log();
			}
			lie_algebra_mean /= n_samples;
			ProjWarpT lie_group_mean_upd = lie_algebra_mean.exp();
			lie_group_mean = lie_group_mean*lie_group_mean_upd;
			lie_group_mean_inv = lie_group_mean.inverse();
			double upd_norm = lie_group_mean_upd.squaredNorm();
			if(upd_norm < params.sample_mean_eps){
				break;
			}
		}
		getStateFromWarp(sample_mean, lie_group_mean);
	} else{
		ProjectiveBase::estimateMeanOfSamples(sample_mean, samples, n_samples);
	}
}

_MTF_END_NAMESPACE


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

#include <votca/tools/linalg.h>
#include <boost/numeric/ublas/matrix_proxy.hpp>

#include <gsl/gsl_linalg.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_eigen.h>

namespace votca { namespace tools {

using namespace std;

void linalg_cholesky_decompose( ub::matrix<double> &A){
        // Cholesky decomposition using GSL
        const size_t N = A.size1();

        gsl_matrix_view A_view = gsl_matrix_view_array(&A(0,0), N, N);

        // get the Cholesky matrices
        int status = gsl_linalg_cholesky_decomp ( &A_view.matrix );
}

void linalg_cholesky_solve(ub::vector<double> &x, ub::matrix<double> &A, ub::vector<double> &b){
    /* calling program should catch the error error code GSL_EDOM
     * thrown by gsl_linalg_cholesky_decomp and take
     * necessary steps
     */

    gsl_matrix_view m
        = gsl_matrix_view_array (&A(0,0), A.size1(), A.size2());

    gsl_vector_view gb
        = gsl_vector_view_array (&b(0), b.size());

    gsl_vector *gsl_x = gsl_vector_alloc (x.size());

    gsl_set_error_handler_off();
    int status = gsl_linalg_cholesky_decomp(&m.matrix);

    if( status == GSL_EDOM)
        throw std::runtime_error("Matrix not symmetric positive definite");

    gsl_linalg_cholesky_solve(&m.matrix, &gb.vector, gsl_x);

    for (size_t i =0 ; i < x.size(); i++)
        x(i) = gsl_vector_get(gsl_x, i);

    gsl_vector_free (gsl_x);
}

}}


//
// Created by Bryn Elesedy on 12/11/2018.
// Linear regression with a single output dimension

#include <getopt.h>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <Eigen/Dense>

using Eigen::MatrixXd;
using Eigen::VectorXd;

const double DEFAULT_LAMBDA_VALUE = 0;
const std::streamsize DEFAULT_PREC = 2;
const char DELIM = ',';
const std::string INPUT_FILE_EXTENSION = ".csv";
const std::string OUTPUT_PREDICTION_EXTENSION = ".fittedvalues";
const std::string OUTPUT_WEIGHTS_EXTENSION = ".weights";

// ----------------
// Parse command line arguments

struct command_line_args {
    command_line_args() : lambda(DEFAULT_LAMBDA_VALUE), prec(DEFAULT_PREC) {}

    double lambda;
    std::streamsize prec;
    std::string input_filename;

};

void print_help_message() {
    std::cout << "==========================================================\n";
    std::cout << ">>>>>>>>>>>>>>>>>L2 Regularised Regression<<<<<<<<<<<<<<<<\n";
    std::cout << "==========================================================\n";
    std::cout << "Usage: linregress <path/to/input_file> [--lambda] [--prec]\n";
    std::cout << "----------------------------------------------------------\n";
    std::cout << "Parameters:\n";
    std::cout << "input_file: File containing training data, must be csv.\n"
                 "            Row stacked training examples.\n"
                 "            Final column the y values, others form X.\n";
    std::cout << "--lambda (numeric): Ridge regularisation parameter.\n"
                 "                    Must be >=0, 0 gives OLS regression.\n"
                 "                    Default is " << DEFAULT_LAMBDA_VALUE << "\n";
    std::cout << "--prec (int): Significant figures for console output.\n"
                 "              Default is " << DEFAULT_PREC << "\n";
    std::cout << "==========================================================\n";
    return;
}

command_line_args parse_args(int argc, char *argv[]) {
    const char *const short_opts = "l:h";
    const option long_opts[] = {
            {"lambda", required_argument, nullptr, 'l'},
            {"prec",   required_argument,       nullptr, 'p'},
            {"help",   no_argument,       nullptr, 'h'},
    };

    command_line_args output;
    int opt;
    int option_index = 0;
    while ((opt = getopt_long(argc, argv, short_opts, long_opts, &option_index)) != -1) {
        switch (opt) {
            case 'l':
                output.lambda = std::stod(optarg);
                break;
            case 'p':
                output.prec = (std::streamsize) std::stoi(optarg);
                break;
            case 'h':
            case '?':
            default:
                print_help_message();
                exit(EXIT_SUCCESS);
        }
    }

    if (optind < argc) {
        output.input_filename = argv[optind++];
    } else {
        std::cerr << "Failure: must give an input file." << std::endl;
        exit(EXIT_FAILURE);
    }

    if (optind < argc) {
        std::cerr << "Warning: " << argc - optind << " additional arguments given, all ignored!\n" << std::endl;
    }

    return output;
}

void check_args(command_line_args &args) {

    if (args.lambda < 0) {
        std::cerr << "Error: Ridge regularisation parameter lambda must be >= 0" << std::endl;
        exit(EXIT_FAILURE);
    }

    if (args.prec < 0) {
        std::cerr << "Warning: given output precision < 0." << std::endl;
    }

    return;
}

// ----------------

// ----------------
// Parse inputs into Eigen matrices

struct data_size {
    std::size_t n_features;
    std::size_t n_examples;
};

struct input_data {
    MatrixXd X;
    VectorXd y;
};

data_size get_data_size(std::string input_filename, char delim) {
    std::ifstream infile(input_filename);

    if (!infile.is_open()) {
        std::cerr << "Error: Failed to open " << input_filename << std::endl;
        exit(EXIT_FAILURE);
    }

    data_size input_sizes;
    std::string line;
    std::size_t lc = 0;
    while (std::getline(infile, line)) {
        if (line.empty()) {
            continue;

        }

        if (lc == 0) {
            input_sizes.n_features = std::count(line.begin(), line.end(), delim);
        }
        lc++;
    }
    input_sizes.n_examples = lc;

    infile.close();
    return input_sizes;
}

input_data read_input_file(std::string input_filename, data_size input_sizes, char delim) {

    std::ifstream infile(input_filename);

    if (!infile.is_open()) {
        std::cerr << "Error: Failed to open " << input_filename << std::endl;
        exit(EXIT_FAILURE);
    }

    MatrixXd X(input_sizes.n_examples, input_sizes.n_features);
    VectorXd y(input_sizes.n_examples);

    std::size_t col = 0;
    std::size_t row = 0;
    std::string line;
    while (std::getline(infile, line)) {
        std::stringstream this_line(line);
        std::string item;

        while (std::getline(this_line, item, delim)) {
            if (item.empty()) {
                continue;
            }

            if (col == input_sizes.n_features) {
                y(row) = std::stod(item);
                col = 0;
                row++;
            } else {
                X(row, col) = std::stod(item);
                col++;
            }
        }

    }
    infile.close();

    input_data output;
    output.X = X;
    output.y = y;
    return output;
}

// ----------------

// ----------------
// Do the Regression

MatrixXd append_ones_column(MatrixXd X) {
    MatrixXd new_x(X.rows(), X.cols() + 1);
    VectorXd c = VectorXd::Constant(X.rows(), 1.);
    new_x << X, c;
    return new_x;
}

class RidgeRegression {
private:
    double lambda;
    bool has_constant;
    bool is_fitted;

    Eigen::HouseholderQR <MatrixXd> qr_decomp;
    VectorXd weights;

public:
    RidgeRegression(double lambda) : lambda(lambda), is_fitted(false) {}

    void fit(MatrixXd X, VectorXd y, bool add_constant);

    VectorXd predict(MatrixXd x_vals);

    inline VectorXd get_weights() { return this->weights; }

    inline bool has_const() { return this->has_constant; }
};

void RidgeRegression::fit(MatrixXd X, VectorXd y, bool add_constant) {
    if (X.rows() != y.rows()) {
        throw std::invalid_argument("X and y must have same number of rows.");
    }

    MatrixXd X_train = (add_constant) ? append_ones_column(X) : X;

    Eigen::MatrixXd eye = Eigen::MatrixXd::Identity(X_train.cols(), X_train.cols());

    this->qr_decomp = (this->lambda * eye + X_train.transpose() * X_train).householderQr();
    this->weights = this->qr_decomp.solve(X_train.transpose() * y);
    this->has_constant = add_constant;
    this->is_fitted = true;
    return;
}

VectorXd RidgeRegression::predict(MatrixXd xvals) {
    if (!this->is_fitted) {
        throw std::domain_error("Need to fit regression before prediction.");
    }

    if (has_constant) {
        double offset = this->weights(this->weights.rows() - 1);
        VectorXd slope = this->weights.head(this->weights.rows() - 1);
        return ((xvals * slope).array() + offset).matrix();
    } else {
        return xvals * this->weights;
    }
}
// ----------------

// ----------------
// Output

double rmse(VectorXd y1, VectorXd y2) {
    return std::sqrt((y1 - y2).array().pow(2).mean());
}

void print_results(RidgeRegression regressor, double rmse_, std::streamsize prec) {

    VectorXd weights = regressor.get_weights();

    std::streamsize curr_prec = std::cout.precision();
    std::cout << "Fitted Model: y = ";
    std::cout << std::setprecision(prec);
    int i = 0;
    while (i < weights.size() - 1) {
        std::cout << weights(i) << " x" << i + 1 << " + ";
        i++;
    }

    std::cout << weights(i) << std::endl;
    std::cout << "RMSE: " << rmse_ << std::endl;
    std::cout << std::setprecision(curr_prec);
    return;
}

std::string get_output_filename(std::string input_filename, std::string old_ext, std::string new_ext) {
    std::string prefix;
    int d = input_filename.size() - old_ext.size();
    if (d > 0 && input_filename.compare(d, old_ext.size(), old_ext) == 0) {
        prefix = input_filename.substr(0, d);
    } else {
        prefix = input_filename;
    }
    return prefix + new_ext;
}

void save_vector(VectorXd fitted_values, std::string filename) {
    std::ofstream outfile(filename);

    if (!outfile.is_open()) {
        std::cerr << "Failed out open output file " << filename << " values not written." << std::endl;
        return;
    }

    for (std::size_t i = 0; i < fitted_values.rows(); i++) {
        outfile << fitted_values(i) << std::endl;
    }

    outfile.close();
    return;
}

// ----------------

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

    command_line_args args = parse_args(argc, argv);

    check_args(args);

    data_size dims = get_data_size(args.input_filename, DELIM);
    input_data data = read_input_file(args.input_filename, dims, DELIM);

    std::cout << "=================================================" << std::endl;
    std::cout << "Least Squares Regression. L2 regularisation = " << args.lambda << std::endl;
    std::cout << "-------------------------------------------------" << std::endl;
    std::cout << "Read data from " << args.input_filename << ", found " << dims.n_features << " features and ";
    std::cout << dims.n_examples << " examples" << std::endl;
    std::cout << "-------------------------------------------------" << std::endl;

    RidgeRegression regressor(args.lambda);
    // Previously considered giving user option to not include constant
    // in regression but I changed my mind.
    regressor.fit(data.X, data.y, true);

    VectorXd fitted_values = regressor.predict(data.X);

    print_results(regressor, rmse(data.y, fitted_values), args.prec);
    save_vector(
            fitted_values,
            get_output_filename(args.input_filename, INPUT_FILE_EXTENSION, OUTPUT_PREDICTION_EXTENSION)
            );

    save_vector(
            regressor.get_weights(),
            get_output_filename(args.input_filename, INPUT_FILE_EXTENSION, OUTPUT_WEIGHTS_EXTENSION)
    );
    std::cout << "=================================================" << std::endl;

    return 0;

}


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

/*
 * dilate-impl.hpp
 *
 * Bitwise integer dilation/undilation. Algorithms courtesy of
 * "Converting to and from Dilated Integers" by Raman and Wise
 * (http://www.cs.indiana.edu/~dswise/Arcee/castingDilated-comb.pdf).
 *
 *  Created on: Sep 16, 2015
 *      Author: David A. Boyuka II
 */
#ifndef SRC_UTIL_IMPL_DILATE_IMPL_HPP_
#define SRC_UTIL_IMPL_DILATE_IMPL_HPP_

//#include <boost/math/special_functions/pow.hpp>
#include <boost/integer/static_log2.hpp>

/*

// WIP - general multiplication-based implementation
template<typename word_t, int d>
struct DilaterHelper {
	static constexpr int w = std::numeric_limits<word_t>::digits;
	static constexpr int s = w / d;

	template<int i>
	class DilaterRound {
		static constexpr int di = std::pow(d, i);
		static constexpr int bitcount = (di < s ? di : s);

		template<int j>
		static constexpr word_t z = z<j-1> | (z<0> >> (j * di * d));

		template<>
		static constexpr word_t z<0> = ((1 << bitcount) - 1) << (d * (s - 1) + 1 - bitcount);

	};
};
*/

template<typename word_t, int d>
word_t Dilater<word_t, d>::dilate(word_t word) { abort(); }

template<typename word_t, int d>
word_t Dilater<word_t, d>::undilate(word_t word) { abort(); }

// Partial specialization for d=2
template<typename word_t>
class Dilater<word_t, 2> {
public:
	static word_t dilate(word_t word);
	static word_t undilate(word_t word);
};

namespace detail {
	template<typename word_t> constexpr int S() { return std::numeric_limits<word_t>::digits/2; } // num dilateable bits
	template<typename word_t> constexpr int ROUNDS() { return boost::static_log2< S<word_t>() >::value; }

	template<typename word_t, int power> constexpr word_t EXP2P1() { return ((word_t)1<<((word_t)1<<power)) + (word_t)1; }

	template<typename word_t, int round>
	struct rounds {
		// DILATION

		// Useful <rounds> go from 1 to ROUNDS inclusive (0 is just a base case, though it is used by UNDILATE_MASK below)
		// MASK<0> is a set mask of the lower half bits of word_t
		// Each successive MASK chops all groups of set bits in half (keeping the lower halves), then
		// duplicates those groups up by twice the width of these now smaller groups
		// Examples: for 8-bit: MASK<0> = 00001111, MASK<1> = 00110011, MASK<2> = 01010101
		static constexpr word_t DILATE_MASK() { return rounds<word_t, round-1>::DILATE_MASK() / EXP2P1<word_t, ROUNDS<word_t>()-round>() * EXP2P1<word_t, ROUNDS<word_t>()-round+1>(); }
		static constexpr word_t dilate_round(word_t word) {
			return
				word = rounds<word_t, round-1>::dilate_round(word),
				(word | (word << (S<word_t>() >> round))) & rounds<word_t, round>::DILATE_MASK();
		}

		// UNDILATION

		// Useful <rounds> go from 1 to ROUNDS inclusive (0 is just a base case)
		// UNDILATE_MASK is just DILATE_MASK in reverse order, shifted by one down
		// Examples: UNDILATE_MASK<1> is DILATE_MASK<ROUNDS-1>, UNDILATE_MASK<ROUNDS> is DILATE_MASK<0>
		static constexpr word_t UNDILATE_MASK() { return rounds<word_t, ROUNDS<word_t>()-round>::DILATE_MASK(); }
		static constexpr word_t undilate_round(word_t word) {
			return
				word = rounds<word_t, round-1>::undilate_round(word),
				(word | (word >> ((word_t)1 << (round-1)))) & rounds<word_t, round>::UNDILATE_MASK();
		}
	};

	template<typename word_t>
	struct rounds<word_t, 0> {
		static constexpr word_t DILATE_MASK() { return ((word_t)1 << S<word_t>()) - 1; }
		static constexpr word_t dilate_round(word_t word) { return word; }
		static constexpr word_t undilate_round(word_t word) { return word; }
	};

	template<typename word_t> constexpr word_t dilate(word_t word) { return rounds<word_t, ROUNDS<word_t>()>::dilate_round(word); }
	template<typename word_t> constexpr word_t undilate(word_t word) { return rounds<word_t, ROUNDS<word_t>()>::undilate_round(word); }
} // namespace detail

template<typename word_t>
word_t Dilater<word_t, 2>::dilate(word_t word) {
	return detail::dilate<word_t>(word);
}

template<typename word_t>
word_t Dilater<word_t, 2>::undilate(word_t word){
	return detail::undilate<word_t>(word);
}

#endif /* SRC_UTIL_IMPL_DILATE_IMPL_HPP_ */


/**
 *  testFrameConversion.cpp
 *
 *  Test the frame conversion between inertial and body frames.
 *
 *  Created by Yinan Li on Nov. 27, 2020.
 *  Hybrid Systems Group, University of Waterloo.
 */

#include <iostream>
#include <cmath>
#include <boost/numeric/odeint.hpp>
#include <cstdlib>
#include <ctime>

#include "src/definitions.h"

/* define dynamics */
struct car_dynamics {
    rocs::Rn u;
    car_dynamics (const rocs::Rn param): u (param) {}

    void operator() (rocs::Rn &x, rocs::Rn &dxdt, double t) const
    {
	dxdt[0] = u[0]*std::cos(x[2]);
	dxdt[1] = u[0]*std::sin(x[2]);
	dxdt[2] = u[1];
    }
};

struct twoagent_ode {
    rocs::Rn u;
    rocs::Rn d;
    twoagent_ode (const rocs::Rn p1, const rocs::Rn p2): u(p1), d(p2){}
    /**
     * ODE model
     * @param x system state: [x,y,theta], n=3
     * @param dxdt vector field
     * @param t time
     */
    void operator() (rocs::Rn &x, rocs::Rn &dxdt, double t) const
    {
	dxdt[0] = -u[0] + d[0]*cos(x[2]) + u[1]*x[1];
	dxdt[1] = d[0]*sin(x[2]) - u[1]*x[0];
	dxdt[2] = d[1] - u[1];
    }
};

void inertia_to_body(rocs::Rn &z, rocs::Rn z0) {
    double x = std::cos(z0[2])*z[0] + std::sin(z0[2])*z[1];
    double y = -std::sin(z0[2])*z[0] + std::cos(z0[2])*z[1];
    z[0] = x;
    z[1] = y;
}

void body_to_inertia(rocs::Rn &z, rocs::Rn z0) {
    double x = std::cos(z0[2])*z[0] - std::sin(z0[2])*z[1];
    double y = std::sin(z0[2])*z[0] + std::cos(z0[2])*z[1];
    z[0] = x;
    z[1] = y;
}

int main() {
    boost::numeric::odeint::runge_kutta_cash_karp54<rocs::Rn> rk45;
    rocs::Rn z0{5.08561,10.1089,1.8}, z1{3.7, 10.8, 0.0};
    rocs::Rn u{0.9, -0.9}, d{0.6, 0.6/4.61};
    rocs::Rn zr(3);
    for(int i = 0; i < 3; ++i) {
	zr[i] = z1[i]-z0[i];
    }
    std::cout << "z1 position to z0 in inertial frame initially: ";
    for(int i = 0; i < 3; ++i) {
	std::cout << zr[i] << ' ';
    }
    std::cout << '\n';

    inertia_to_body(zr, z0);
    std::cout << "z1 position to z0 in z0 body frame initially: ";
    for(int i = 0; i < 3; ++i) {
	std::cout << zr[i] << ' ';
    }
    std::cout << '\n';

    const double h = 0.3;  // sampling time
    const double dt = 0.001; //integration step size for odeint
    /* Integrate z0 and z1 separately */
    boost::numeric::odeint::integrate_const(rk45, car_dynamics(u), z0, 0.0, h, dt);
    boost::numeric::odeint::integrate_const(rk45, car_dynamics(d), z1, 0.0, h, dt);
    std::cout << "z1-z0 = ";
    for(int i = 0; i < 3; ++i) {
	std::cout << z1[i]-z0[i] << ' ';
    }
    std::cout << '\n';
    /* Integrate by the local model */
    boost::numeric::odeint::integrate_const(rk45, twoagent_ode(u,d), zr, 0.0, h, dt);
    std::cout << "z1 position to z0 in z0 body frame after 0.3s: ";
    for(int i = 0; i < 3; ++i) {
	std::cout << zr[i] << ' ';
    }
    std::cout << '\n';

    body_to_inertia(zr, z0);
    std::cout << "z1 position to z0 in inertial frame after 0.3s: ";
    for(int i = 0; i < 3; ++i) {
	std::cout << zr[i] << ' ';
    }
    std::cout << '\n';
}


#pragma once
#include "numeric/dense.hpp"
#include <Eigen/Jacobi>

namespace cpb { namespace compute {

namespace detail {
    template<class real_t>
    static void tridiagonal_qr_step(real_t* diag, real_t* subdiag, int start, int end) {
        auto td = (diag[end-1] - diag[end]) * real_t(0.5);
        auto e = subdiag[end-1];
        auto mu = diag[end];
        if (td == 0) {
            mu -= std::abs(e);
        }
        else {
            auto e2 = Eigen::numext::abs2(subdiag[end-1]);
            auto h = Eigen::numext::hypot(td, e);
            if (e2 == 0)
                mu -= (e / (td + (td>0 ? 1 : -1))) * (e / h);
            else
                mu -= e2 / (td + (td>0 ? h : -h));
        }

        auto x = diag[start] - mu;
        auto z = subdiag[start];
        for (auto k = start; k < end; ++k) {
            Eigen::JacobiRotation<real_t> rot;
            rot.makeGivens(x, z);

            // do T = G' T G
            auto sdk = rot.s() * diag[k] + rot.c() * subdiag[k];
            auto dkp1 = rot.s() * subdiag[k] + rot.c() * diag[k+1];

            diag[k] = rot.c() * (rot.c() * diag[k] - rot.s() * subdiag[k])
                - rot.s() * (rot.c() * subdiag[k] - rot.s() * diag[k+1]);
            diag[k+1] = rot.s() * sdk + rot.c() * dkp1;
            subdiag[k] = rot.c() * sdk - rot.s() * dkp1;

            if (k > start)
                subdiag[k - 1] = rot.c() * subdiag[k-1] - rot.s() * z;

            x = subdiag[k];
            if (k < end - 1) {
                z = -rot.s() * subdiag[k+1];
                subdiag[k + 1] = rot.c() * subdiag[k+1];
            }
        }
    }
}

template<class Derived, class scalar_t = typename Derived::Scalar>
inline ArrayX<scalar_t> tridiagonal_eigenvalues(const DenseBase<Derived>& alpha,
                                                const DenseBase<Derived>& beta)
{
    ArrayX<scalar_t> eigenvalues = alpha;
    ArrayX<scalar_t> temp = beta;

    auto start = 0;
    auto end = static_cast<int>(eigenvalues.size()) - 1;
    auto iter = 0;
    constexpr auto max_iterations = 30;

    while (end > 0) {
        for (auto i = start; i < end; ++i) {
            auto a = std::abs(temp[i]);
            auto b = std::abs(eigenvalues[i]) + std::abs(eigenvalues[i + 1]);
            // if a is much smaller than b
            if (a < b * std::numeric_limits<scalar_t>::epsilon())
                temp[i] = 0;
        }

        while (end > 0 && temp[end-1] == 0)
            end--;

        if (end <= 0)
            break;

        if (++iter > max_iterations * eigenvalues.size())
            throw std::runtime_error{"Tridiagonal QR error"};

        start = end - 1;
        while (start > 0 && temp[start-1] != 0)
            start--;

        detail::tridiagonal_qr_step(eigenvalues.data(), temp.data(), start, end);
    }

    return eigenvalues;
}

}} // namespace cpb::compute


/* =========================================================================
   Copyright (c) 2010-2016, Institute for Microelectronics,
                            Institute for Analysis and Scientific Computing,
                            TU Wien.
   Portions of this software are copyright by UChicago Argonne, LLC.

                            -----------------
                  ViennaCL - The Vienna Computing Library
                            -----------------

   Project Head:    Karl Rupp                   rupp@iue.tuwien.ac.at

   (A list of authors and contributors can be found in the PDF manual)

   License:         MIT (X11), see file LICENSE in the base directory
============================================================================= */

/** \example armadillo-with-viennacl.cpp
*
*   This tutorial shows how data can be directly transferred from the Armadillo Library to ViennaCL objects using the built-in convenience wrappers.
*
*   The first step is to include the necessary headers and activate the Armadillo convenience functions in ViennaCL:
**/

// System headers
#include <iostream>

// Armadillo headers (disable BLAS and LAPACK to avoid linking issues)
#define ARMA_DONT_USE_BLAS
#define ARMA_DONT_USE_LAPACK
#include <armadillo>

// IMPORTANT: Must be set prior to any ViennaCL includes if you want to use ViennaCL algorithms on Armadillo objects
#define VIENNACL_WITH_ARMADILLO 1

// ViennaCL includes
#include "viennacl/vector.hpp"
#include "viennacl/matrix.hpp"
#include "viennacl/compressed_matrix.hpp"
#include "viennacl/linalg/prod.hpp"

// Helper functions for this tutorial:
#include "vector-io.hpp"

/**
*    The following function contains the main code for this tutorial.
*    It consists of the following steps:
*      - Creates Armadillo matrices and vectors
*      - Initializes them with data
*      - Create ViennaCL objects
*      - Copy them over to the respective ViennaCL objects
*      - Compute matrix-vector products in both Armadillo and ViennaCL and compare results.
*
**/
template<typename NumericT>
void run_tutorial()
{
  typedef arma::SpMat<NumericT>  ArmaSparseMatrix;
  typedef arma::Mat<NumericT>    ArmaMatrix;
  typedef arma::Col<NumericT>    ArmaVector;

  /**
  * Create and fill dense matrices from the Armadillo library:
  **/
  ArmaMatrix arma_densemat(6, 5);
  ArmaMatrix arma_densemat2(6, 5);
  arma_densemat(0,0) = 2.0;   arma_densemat(0,1) = -1.0;
  arma_densemat(1,0) = -1.0;  arma_densemat(1,1) =  2.0;  arma_densemat(1,2) = -1.0;
  arma_densemat(2,1) = -1.0;  arma_densemat(2,2) = -1.0;  arma_densemat(2,3) = -1.0;
  arma_densemat(3,2) = -1.0;  arma_densemat(3,3) =  2.0;  arma_densemat(3,4) = -1.0;
                              arma_densemat(5,4) = -1.0;  arma_densemat(4,4) = -1.0;

  /**
  * Create and fill sparse matrices from the Armadillo library:
  **/
  ArmaSparseMatrix arma_sparsemat(6, 5);
  ArmaSparseMatrix arma_sparsemat2(6, 5);
  arma_sparsemat(0,0) = 2.0;   arma_sparsemat(0,1) = -1.0;
  arma_sparsemat(1,1) = 2.0;   arma_sparsemat(1,2) = -1.0;
  arma_sparsemat(2,2) = -1.0;  arma_sparsemat(2,3) = -1.0;
  arma_sparsemat(3,3) = 2.0;   arma_sparsemat(3,4) = -1.0;
  arma_sparsemat(5,4) = -1.0;

  /**
  * Create and fill a few vectors from the Armadillo library:
  **/
  ArmaVector arma_rhs(5);
  ArmaVector arma_result(6);
  ArmaVector arma_temp(6);

  arma_rhs(0) = 10.0;
  arma_rhs(1) = 11.0;
  arma_rhs(2) = 12.0;
  arma_rhs(3) = 13.0;
  arma_rhs(4) = 14.0;

  /**
  * Create the corresponding ViennaCL objects:
  **/
  viennacl::vector<NumericT> vcl_rhs(5);
  viennacl::vector<NumericT> vcl_result(6);
  viennacl::matrix<NumericT> vcl_densemat(6, 5);
  viennacl::compressed_matrix<NumericT> vcl_sparsemat(6, 5);

  /**
  * Directly copy the Armadillo objects to ViennaCL objects
  **/
  viennacl::copy(arma_rhs.memptr(), arma_rhs.memptr() + arma_rhs.n_elem, vcl_rhs.begin());  //method 1: via iterator interface (cf. std::copy())
  viennacl::copy(arma_rhs, vcl_rhs);  //method 2: via built-in wrappers (convenience layer)

  viennacl::copy(arma_densemat, vcl_densemat);
  viennacl::copy(arma_sparsemat, vcl_sparsemat);
  std::cout << "VCL sparsematrix dimensions: " << vcl_sparsemat.size1() << ", " << vcl_sparsemat.size2() << std::endl;

  // For completeness: Copy matrices from ViennaCL back to Eigen:
  viennacl::copy(vcl_densemat, arma_densemat2);
  viennacl::copy(vcl_sparsemat, arma_sparsemat2);

  /**
  * Run dense matrix-vector products and compare results:
  **/
  arma_result = arma_densemat * arma_rhs;
  vcl_result = viennacl::linalg::prod(vcl_densemat, vcl_rhs);
  viennacl::copy(vcl_result, arma_temp);
  std::cout << "Difference for dense matrix-vector product: " << norm(arma_result - arma_temp) << std::endl;
  std::cout << "Difference for dense matrix-vector product (Armadillo -> ViennaCL -> Armadillo): "
            << norm(arma_densemat2 * arma_rhs - arma_temp) << std::endl;

  /**
  * Run sparse matrix-vector products and compare results:
  **/
  arma_result = arma_sparsemat * arma_rhs;
  vcl_result = viennacl::linalg::prod(vcl_sparsemat, vcl_rhs);
  viennacl::copy(vcl_result, arma_temp);
  std::cout << "Difference for sparse matrix-vector product: " << norm(arma_result - arma_temp) << std::endl;
  std::cout << "Difference for sparse matrix-vector product (Armadillo -> ViennaCL -> Armadillo): "
            << norm(arma_sparsemat2 * arma_rhs - arma_temp) << std::endl;
}

/**
*   In the main() routine we only call the worker function defined above with both single and double precision arithmetic.
**/
int main(int, char *[])
{
  std::cout << "----------------------------------------------" << std::endl;
  std::cout << "## Single precision" << std::endl;
  std::cout << "----------------------------------------------" << std::endl;
  run_tutorial<float>();

#ifdef VIENNACL_HAVE_OPENCL
  if ( viennacl::ocl::current_device().double_support() )
#endif
  {
    std::cout << "----------------------------------------------" << std::endl;
    std::cout << "## Double precision" << std::endl;
    std::cout << "----------------------------------------------" << std::endl;
    run_tutorial<double>();
  }

  /**
  *   That's it. Print a success message and exit.
  **/
  std::cout << std::endl;
  std::cout << "!!!! TUTORIAL COMPLETED SUCCESSFULLY !!!!" << std::endl;
  std::cout << std::endl;

}


#include <iostream>     // std::cout
#include <algorithm>    // std::count
#include <unordered_map>//unordered_map
#include <vector>       // std::vector
#include <string>       // std::string, std::to_string
#include <iostream>     // std::cout
#include <random>       // std::random_device std::default_random_engine std::uniform_int_distribution std::normal_distribution
#include <numeric>      // std::inner_product
#include <chrono>       // std::chrono:
#include <boost/math/distributions/normal.hpp> //boost::math::normal
#include <math.h>  //std::sqrt

constexpr size_t n = 4;
int p = 101;
int s[n] = { 1, 2, 3, 4}; //{ 0, 0, 0 }; //{ 1, 2, 3 };

// Declare EQUATION struct type
struct EQUATION {
    int a[n];
    int b;
    int e;

};

// mappings
std::unordered_map<std::string, std::vector<EQUATION>> equationMemo;
std::unordered_map<int, std::unordered_map<int, long>> guessCount;  // {position, {guess, count}

// error distributions
/// rounded gaussian
float mu = 0;
float sigma = 2;

std::random_device rd;
std::default_random_engine generator(rd());
std::uniform_int_distribution<int> uniform_int(0, p-1);  //min, max
std::normal_distribution<float> gaussian(mu, sigma);  // mean, std

/// centered binomial
constexpr size_t upper = 2;  // has to be odd
std::binomial_distribution<int> binomial(upper, 0.5);  //upper bound, probability of success

// oracle
void oracleLWE(EQUATION& equation){

    for (int i = 0; i < n; i++) {
        equation.a[i] = uniform_int(generator);
    }

    // discrete gaussian
    //equation.e = round(gaussian(generator));

    // centered binomial (used in NewHope reference implementation)
    equation.e = binomial(generator) - upper/2;

    equation.b = std::inner_product(s, s+n, equation.a, 0.0) + equation.e;
    equation.b %= p;
}

void encodePosOfZeros(std::string& key, int(&coeffs)[n]) {
    for (const auto& coeff : coeffs)
        if (coeff == 0)
            key.push_back('0');
        else
            key.push_back('1');
}
void initializeEquationMemo(int n, int arr[], int i) {
    if (i == n) {
        std::string key;
        for (int j=0; j < n; j++) {
            key += std::to_string(arr[j]);
        }
        std::vector<EQUATION> equations;
        equationMemo.insert({ key, equations });
        //std::cout << key << std::endl;
        return;
    }

    // First assign "0" at ith position
    // and try for all other permutations
    // for remaining positions
    arr[i] = 0;
    initializeEquationMemo(n, arr, i + 1);

    // And then assign "1" at ith position
    // and try for all other permutations
    // for remaining positions
    arr[i] = 1;
    initializeEquationMemo(n, arr, i + 1);
}

// fast pow for integer using square and multiply algo
int integerPow(int a, int ex, int mod) {
    if (a == 1) return 1;

    long r = 1;
    while (ex) {
        if (ex & 1)r = (r * a) % mod;
        ex = ex >> 1;
        a = (a * a) % mod;
    }
    return r;
}

void compareEquationsSub(std::string& key, EQUATION& prevEquation, EQUATION& equation) {
    for (int i = 0; i < n; ++i) {
        if ((prevEquation.a[i] - equation.a[i]) % p == 0)
            key.push_back('0');
        else
            key.push_back('1');
    }
}

void compareEquationsAdd(std::string& key, EQUATION& prevEquation, EQUATION& equation) {
    for (int i = 0; i < n; ++i) {
        if ((prevEquation.a[i] + equation.a[i]) % p == 0)
            key.push_back('0');
        else
            key.push_back('1');
    }
}

void furtherElimination(std::string& prevKey, EQUATION& prevEquation) {
    auto equations = equationMemo.at(prevKey);

    std::string key;
    for (auto& equation : equations) {
        EQUATION newEquation;

        // compare for subtraction
        key.clear();
        compareEquationsSub(key, prevEquation, equation);

        if (std::count(key.begin(), key.end(), '1') == 1) {
            for (int i = 0; i < n; ++i)
                newEquation.a[i] = (prevEquation.a[i] - equation.a[i]) % p;
            newEquation.b = (prevEquation.b - equation.b) % p;
            int pos = key.find('1');
            int coeff = newEquation.a[pos];
            int s = newEquation.b*integerPow(coeff, p-2, p)%p;  // solving linear congruence fast by using little fermat
            if (s < 0)
                s += p;

            guessCount.at(pos).at(s) += 1;
        }
        else {
            if (std::count(key.begin(), key.end(), '1') < std::count(prevKey.begin(), prevKey.end(), '1')) {
                for (int i = 0; i < n; ++i)
                    newEquation.a[i] = (prevEquation.a[i] - equation.a[i]) % p;
                newEquation.b = (prevEquation.b - equation.b) % p;

                furtherElimination(key, newEquation);

                equationMemo.at(key).push_back(equation);
            }
        }

        // compare for addition
        key.clear();
        compareEquationsAdd(key, prevEquation, equation);

        if (std::count(key.begin(), key.end(), '1') == 1) {
            for (int i = 0; i < n; ++i)
                newEquation.a[i] = (prevEquation.a[i] + equation.a[i]) % p;
            newEquation.b = (prevEquation.b + equation.b) % p;
            int pos = key.find('1');
            int coeff = newEquation.a[pos];
            int s = newEquation.b * integerPow(coeff, p - 2, p) % p;  // solving linear congruence fast by using little fermat
            if (s < 0)
                s += p;

            guessCount.at(pos).at(s) += 1;
        }
        else {
            if (std::count(key.begin(), key.end(), '1') < std::count(prevKey.begin(), prevKey.end(), '1')) {
                for (int i = 0; i < n; ++i)
                    newEquation.a[i] = (prevEquation.a[i] + equation.a[i]) % p;
                newEquation.b = (prevEquation.b + equation.b) % p;

                furtherElimination(key, newEquation);

                equationMemo.at(key).push_back(equation);
            }
        }
    }
}

int main(int argc, char** argv) {
    // initialize equation memory
    int arr[n];
    initializeEquationMemo(n, arr, 0);

    // initialize map to track the count of guesses for a possible solution
    for (int pos = 0; pos < n; pos++) {
        std::unordered_map<int, long> tmp;
        for (int guess = 0; guess < p; guess++)
            tmp.insert(std::make_pair(guess, 0));

        guessCount.insert(std::make_pair(pos, tmp));
    }

    std::chrono::steady_clock::time_point begin = std::chrono::steady_clock::now();

    long i = 1;
    while (1) {
        EQUATION equation;
        oracleLWE(equation);

        std::string key;
        encodePosOfZeros(key, equation.a);

        furtherElimination(key, equation);

        equationMemo.at(key).push_back(equation);

        // early stopping
        if (i % 10 == 0) {
            bool ready = true;
            for (auto const& pos : guessCount) {
                float mean = 0;
                for (auto const& guess : pos.second)
                    mean += guess.second;
                if (mean < 20) {
                    std::cout << "mean: " << mean << std::endl;
                    goto endOfLoop;
                }
                mean /= p;

                float var = 0;
                for (auto const& guess : pos.second)
                    var += (guess.second - mean) * (guess.second - mean); // fast pow
                var /= p - 1;
                float sigma = std::sqrt(var);

                int guessCount = 0;
                for (auto const& guess : pos.second) {
                    if (mean + sigma < guess.second) {
                        guessCount += 1;
                        std::cout << "Pos: " << pos.first << "\t\tGuess: " << guess.first << "\t\tCount: " <<
                            guess.second << "\t\tRight solution: " << s[pos.first] << std::endl;
                    }
                }
                std::cout << std::endl;
                if (guessCount != 1) {
                    ready = false;
                }
            }
            if (ready) {
                std::cout << "Count of equations: " << i << std::endl << std::endl;
                break;
            }
        }
    endOfLoop: {}
        std::cout << "Actual equations: " << i << "\t\tProcessing Time = " <<
            std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - begin).count() / 1000 <<
            "[s]" << std::endl;
        i += 1;
    }

    std::chrono::steady_clock::time_point end = std::chrono::steady_clock::now();

    std::cout << "Time difference = " << std::chrono::duration_cast<std::chrono::milliseconds>(end - begin).count()/1000 << "[s]" << std::endl;
    //return 0;

    // p-value empirical distribution
    for (auto const& pos : guessCount) {
        float mean = 0;
        for (auto const& guess : pos.second)
            mean += guess.second;
        mean /= p;

        float var = 0;
        for (auto const& guess : pos.second)
            var += (guess.second - mean)*(guess.second - mean); // fast pow
        var /= p-1;  // empirical variance
        float sigma = std::sqrt(var);
        std::cout << "mean + sigma " << mean+sigma << std::endl;
    }
    std::cout << std::endl;

    // p-value gaussian approximation of binomial
    /*for (auto const& pos : guessCount) {
        long sum = 0;
        for (auto const& guess : pos.second)
            sum += guess.second;
        float mean = sum/p;
        float sigma = std::sqrt(mean * (1 - 1 / p));
        boost::math::normal dist(mean, sigma);
        double q = quantile(dist, 0.99);
        std::cout << "quantile = " << q << std::endl;
    }*/

    // determine guess

    // debug
    for (auto const& pos : guessCount) {
        float mean = 0;
        for (auto const& guess : pos.second)
            mean += guess.second;
        mean /= p;

        float var = 0;
        for (auto const& guess : pos.second)
            var += (guess.second - mean) * (guess.second - mean); // fast pow
        var /= p-1;
        float sigma = std::sqrt(var);

        for (auto const& guess : pos.second) {
            if (mean + sigma < guess.second) {
                std::cout << "Pos: " << pos.first << "\t\tguess: " << guess.first << "\t\tcount: " << guess.second << "\t\tright: " << s[pos.first] << '\n';
            }
            //if (guess.second != 0)
            //    std::cout << "guess: " << guess.first << "    count: " << guess.second <<'\n';
        }
    }
    return 0;

    // debug
    for (auto const& pos : guessCount) {
        int trackGuess=0, trackMax=0;
        for (auto const& guess : pos.second) {
            if (trackMax < guess.second) {
                trackGuess = guess.first;
                trackMax = guess.second;
            }
            //if (guess.second != 0)
            //    std::cout << "guess: " << guess.first << "    count: " << guess.second <<'\n';
        }
        std::cout << "Pos: " << pos.first << "\t\t\tguess: " << trackGuess << "\t\t\count: " << trackMax << "\t\t\tright: " << s[pos.first] << '\n';
    }
    return 0;

    for (auto const& tmp : equationMemo) {
        std::cout << "Key: " << tmp.first << '\n';
        std::cout << "tmp contains " << tmp.second.size() << " elements.\n";
        continue;
        for (auto const& equation : tmp.second) {
            for (int i = 0; i < n; i++)
                std::cout << "a[" << i << "]: " << equation.a[i] << '\n';

            std::cout << "b: " << equation.b << '\n';
            std::cout << "e: " << equation.e << '\n' << '\n';
        }

    }

    return 0;
}


#ifndef STAN_MATH_PRIM_SCAL_FUN_BESSEL_FIRST_KIND_HPP
#define STAN_MATH_PRIM_SCAL_FUN_BESSEL_FIRST_KIND_HPP

#include <boost/math/special_functions/bessel.hpp>
#include <stan/math/prim/scal/err/check_not_nan.hpp>

namespace stan {
namespace math {

/**
 *
   \f[
   \mbox{bessel\_first\_kind}(v, x) =
   \begin{cases}
     J_v(x) & \mbox{if } -\infty\leq x \leq \infty \\[6pt]
     \textrm{error} & \mbox{if } x = \textrm{NaN}
   \end{cases}
   \f]

   \f[
   \frac{\partial\, \mbox{bessel\_first\_kind}(v, x)}{\partial x} =
   \begin{cases}
     \frac{\partial\, J_v(x)}{\partial x} & \mbox{if } -\infty\leq x\leq \infty
 \\[6pt] \textrm{error} & \mbox{if } x = \textrm{NaN} \end{cases} \f]

   \f[
   J_v(x)=\left(\frac{1}{2}x\right)^v
   \sum_{k=0}^\infty \frac{\left(-\frac{1}{4}x^2\right)^k}{k!\, \Gamma(v+k+1)}
   \f]

   \f[
   \frac{\partial \, J_v(x)}{\partial x} = \frac{v}{x}J_v(x)-J_{v+1}(x)
   \f]
 *
 */
template <typename T2>
inline T2 bessel_first_kind(int v, const T2 z) {
  check_not_nan("bessel_first_kind", "z", z);
  return boost::math::cyl_bessel_j(v, z);
}

}  // namespace math
}  // namespace stan
#endif


#include "core/Initializer.h"

#include <Eigen/Core>
#include <Eigen/SVD>
#include <Eigen/Dense>

using namespace Eigen;
using namespace std;

namespace EllipsoidSLAM
{
    Initializer::Initializer(int rows, int cols) {
        miImageRows = rows;
        miImageCols = cols;
    }

    // ascending sort
    bool cmp(pair<int, double>a, pair<int, double>b)
    {
        return a.second > b.second;
    }

    // Initialize an ellipsoid from several bounding boxes and camera poses.
    g2o::ellipsoid Initializer::initializeQuadric(MatrixXd &pose_mat, MatrixXd &detection_mat, Matrix3d &calib) {
        mbResult = false;

        int input_size = pose_mat.rows();

        // 1) Generate tangent planes from bounding boxes and camera cneter.
        //    Those invalid observations are ignored.
        MatrixXd planesHomo = getPlanesHomo(pose_mat, detection_mat, calib);

        int plane_size = planesHomo.cols();
        int invalid_plane = input_size*4 - plane_size;
        if( invalid_plane > 0)
            std::cout << " * invalid_plane: " << invalid_plane << std::endl;

        if(plane_size < 9)  // at least 9 planes are needed
        {
            g2o::ellipsoid e_bad;
            mbResult = false;
            return e_bad;
        }

        // Using SVD to generate a quadric
        MatrixXd planesVector = getVectorFromPlanesHomo(planesHomo);
        Matrix4d QStar = getQStarFromVectors(planesVector);

        // generate ellipsoid from the quadric
        g2o::ellipsoid e = getEllipsoidFromQStar(QStar);

        // Set color : blue.
        Eigen::Vector3d blueScalar(0,0,255);
        e.setColor(blueScalar);
        return e;
    }

    MatrixXd Initializer::getPlanesHomo(MatrixXd &pose_mat, MatrixXd &detection_mat, Matrix3d &calib) {
        assert( pose_mat.rows() == detection_mat.rows() && " Two matrices should match. " );
        assert( pose_mat.rows() > 2 && " At least 3 measurements are required. " );

        MatrixXd planes_all(4,0);

        int rows = pose_mat.rows();
        for(int i=0;i<rows;i++)
        {
            VectorXd pose = pose_mat.row(i);
            VectorXd detection = detection_mat.row(i);

            // filter invalid detections
            if( detection(0) < 1 && detection(1) < 1 && detection(2) < 1 && detection(3) < 1  )
                continue;

            Vector7d pose_vec = pose.tail(7);
            g2o::SE3Quat campose_wc(pose_vec);
            // get projection matrix
            MatrixXd P = generateProjectionMatrix(campose_wc.inverse(), calib);

            MatrixXd lines = fromDetectionsToLines(detection);
            MatrixXd planes = P.transpose() * lines;

            // add to matrix
            for( int m=0;m<planes.cols();m++)
            {
                planes_all.conservativeResize(planes_all.rows(), planes_all.cols()+1);
                planes_all.col(planes_all.cols()-1) = planes.col(m);
            }
        }

        return planes_all;
    }

    Matrix3Xd Initializer::generateProjectionMatrix(const SE3Quat& campose_cw, const Matrix3d& Kalib) const
    {
        Matrix3Xd identity_lefttop;
        identity_lefttop.resize(3, 4);
        identity_lefttop.col(3)=Vector3d(0,0,0);
        identity_lefttop.topLeftCorner<3,3>() = Matrix3d::Identity(3,3);

        Matrix3Xd proj_mat = Kalib * identity_lefttop;

        proj_mat = proj_mat * campose_cw.to_homogeneous_matrix();

        return proj_mat;
    }

    MatrixXd Initializer::fromDetectionsToLines(VectorXd &detections) {
        bool flag_openFilter = true;        // filter those lines lying on the image boundary

        double x1 = detections(0);
        double y1 = detections(1);
        double x2 = detections(2);
        double y2 = detections(3);

        Vector3d line1 (1, 0, -x1);
        Vector3d line2 (0, 1, -y1);
        Vector3d line3 (1, 0, -x2);
        Vector3d line4 (0, 1, -y2);

        // those lying on the image boundary have been marked -1
        MatrixXd line_selected(3, 0);
        MatrixXd line_selected_none(3, 0);
        if( !flag_openFilter || ( x1>0 && x1<miImageCols-1 ))
        {
            line_selected.conservativeResize(3, line_selected.cols()+1);
            line_selected.col(line_selected.cols()-1) = line1;
        }
        if( !flag_openFilter || (y1>0 && y1<miImageRows-1 ))
        {
            line_selected.conservativeResize(3, line_selected.cols()+1);
            line_selected.col(line_selected.cols()-1) = line2;
        }
        if( !flag_openFilter || (x2>0 && x2<miImageCols-1 ))
        {
            line_selected.conservativeResize(3, line_selected.cols()+1);
            line_selected.col(line_selected.cols()-1) = line3;
        }
        if( !flag_openFilter || (y2>0 && y2<miImageRows-1 ))
        {
            line_selected.conservativeResize(3, line_selected.cols()+1);
            line_selected.col(line_selected.cols()-1) = line4;
        }

        return line_selected;
    }

    MatrixXd Initializer::getVectorFromPlanesHomo(MatrixXd &planes) {
        int cols = planes.cols();

        MatrixXd planes_vector(10,0);

        for(int i=0;i<cols;i++)
        {
            VectorXd p = planes.col(i);
            Vector10d v;

            v << p(0)*p(0),2*p(0)*p(1),2*p(0)*p(2),2*p(0)*p(3),p(1)*p(1),2*p(1)*p(2),2*p(1)*p(3),p(2)*p(2),2*p(2)*p(3),p(3)*p(3);

            planes_vector.conservativeResize(planes_vector.rows(), planes_vector.cols()+1);
            planes_vector.col(planes_vector.cols()-1) = v;
        }

        return planes_vector;
    }

    Matrix4d Initializer::getQStarFromVectors(MatrixXd &planeVecs) {
        MatrixXd A = planeVecs.transpose();

        // svd decompose
        JacobiSVD<Eigen::MatrixXd> svd(A, ComputeThinU | ComputeThinV );
        MatrixXd V = svd.matrixV();

        VectorXd qj_hat = V.col(V.cols()-1);

        // Get QStar
        Matrix4d QStar;
        QStar <<
            qj_hat(0),qj_hat(1),qj_hat(2),qj_hat(3),
            qj_hat(1),qj_hat(4),qj_hat(5),qj_hat(6),
            qj_hat(2),qj_hat(5),qj_hat(7),qj_hat(8),
            qj_hat(3),qj_hat(6),qj_hat(8),qj_hat(9);

        return QStar;
    }

    g2o::ellipsoid Initializer::getEllipsoidFromQStar(Matrix4d &QStar) {
        g2o::ellipsoid e;

        Matrix4d Q = QStar.inverse() * cbrt(QStar.determinant());

        SelfAdjointEigenSolver<Matrix4d> es(Q);    // ascending order by default
        MatrixXd D = es.eigenvalues().asDiagonal();
        MatrixXd V = es.eigenvectors();

        VectorXd eigens = es.eigenvalues();

        // For an ellipsoid, the signs of the eigenvalues must be ---+ or +++-
        int num_pos = int(eigens(0)>0) +int(eigens(1)>0) +int(eigens(2)>0) +int(eigens(3)>0);
        int num_neg = int(eigens(0)<0) +int(eigens(1)<0) +int(eigens(2)<0) +int(eigens(3)<0);
        if( !(num_pos ==3 && num_neg == 1) && !(num_pos ==1 && num_neg == 3) ){
            cout << " Not Ellipsoid : pos/neg  " << num_pos << " / " << num_neg << endl;
            cout << "eigens :" << eigens.transpose() << endl;
            mbResult = false;
            return e;
        }
        else
            mbResult = true;

        if( eigens(3) > 0  )  // normalize to - - - +
        {
            Q=-Q;
            SelfAdjointEigenSolver<Matrix4d> es_2(Q);
            D = es_2.eigenvalues().asDiagonal();
            V = es_2.eigenvectors();

            eigens = es_2.eigenvalues();
        }

        // Solve ellipsoid parameters from matrix Q
        Vector3d lambda_mat = eigens.head(3).array().inverse();

        Matrix3d Q33 = Q.block(0,0,3,3);

        double k = Q.determinant()/Q33.determinant();

        Vector3d value = -k*(lambda_mat);
        Vector3d s = value.array().abs().sqrt();

        Vector4d t = QStar.col(3);
        t = t/t(3);
        Vector3d translation = t.head(3);

        SelfAdjointEigenSolver<Matrix3d> es2(Q33);
        MatrixXd D_Q33 = es2.eigenvalues().asDiagonal();
        MatrixXd rot = es2.eigenvectors();

        double r,p,y;
        rot_to_euler_zyx<double>(rot, r, p, y);
        Vector3d rpy(r,p,y);

        // generate ellipsoid
        Vector9d objectVec;
        objectVec << t(0),t(1),t(2),rpy(0),rpy(1),rpy(2),s(0),s(1),s(2);
        e.fromMinimalVector(objectVec);

        return e;

    }

    void Initializer::sortEigenValues(VectorXcd &eigens, MatrixXcd &V) {
        vector<pair<int, double>> pairs;
        for(int i=0;i<4;i++)
            pairs.push_back( make_pair( i, eigens(i).real() ) );
        sort(pairs.begin(), pairs.end(), cmp);

        // Construct a new matrix in order
        MatrixXcd V_new(4,4);
        Vector4cd eigens_new;

        for(int i=0;i<4;i++)
        {
            int oldID = pairs[i].first;
            V_new.col(i) = V.col(oldID);
            eigens_new(i) = complex<double>(pairs[i].second, 0);
        }

        eigens = eigens_new;
        V = V_new;
    }

    double Initializer::quadricErrorWithPlanes(MatrixXd &pose_mat, MatrixXd &detection_mat, Matrix3d &calib,
                                               g2o::ellipsoid &e) {
        Matrix4d QStar = e.generateQuadric();           // get Q^*
        Vector10d qj_hat;
        qj_hat << QStar(0,0),QStar(0,1),QStar(0,2),QStar(0,3),QStar(1,1),QStar(1,2),QStar(1,3),QStar(2,2),QStar(2,3),QStar(3,3);

        // Get planes vector
        MatrixXd planesHomo = getPlanesHomo(pose_mat, detection_mat, calib);
        MatrixXd planesVector = getVectorFromPlanesHomo(planesHomo);

        // Get error
        VectorXd result = planesVector.transpose() * qj_hat;
        return result.transpose() * result;
    }

    bool Initializer::getInitializeResult() {
        return mbResult;
    }

    g2o::ellipsoid Initializer::initializeQuadric(EllipsoidSLAM::Observations &obs, Matrix3d &calib) {
        MatrixXd pose_mat;
        MatrixXd detection_mat;

        // get pose matrix and detection matrix from observations
        getDetectionAndPoseMatFromObservations(obs, pose_mat, detection_mat);

        g2o::ellipsoid e = initializeQuadric(pose_mat, detection_mat, calib);

        if( getInitializeResult() )
        {
            e.miLabel = obs[0]->label;
        }

        return e;
    }

    void Initializer::getDetectionAndPoseMatFromObservations(EllipsoidSLAM::Observations &obs, MatrixXd &pose_mat,
                                                             MatrixXd &detection_mat) {
        int frameSize = obs.size();
        pose_mat.resize(frameSize, 7);  // x y z qx qy qz qw
        detection_mat.resize(frameSize, 5); // x1 y1 x2 y2 accuracy

        int id = 0;
        for (auto iter = obs.begin(); iter!=obs.end(); iter++)
        {
            Vector5d det_vec;
            det_vec << (*iter)->bbox, (*iter)->rate;
            Vector7d pos_vec = (*iter)->pFrame->cam_pose_Twc.toVector();

            pose_mat.row(id) = pos_vec;
            detection_mat.row(id) = det_vec;

            id ++;
        }
    }
}


#include <iostream>
#include <Eigen/Dense>
#include <unsupported/Eigen/LevenbergMarquardt>
#include <unsupported/Eigen/MatrixFunctions>

using namespace Eigen;
using namespace std;

void flatten(Matrix<double, 4, 4> &g, Matrix<double,1,12> &y) {
    y(0) = g(0,0);
    y(1) = g(0,1);
    y(2) = g(0,2);
    y(3) = g(1,0);
    y(4) = g(1,1);
    y(5) = g(1,2);
    y(6) = g(2,0);
    y(7) = g(2,1);
    y(8) = g(2,2);
    y(9) = g(0,3);
    y(10) = g(1,3);
    y(11) = g(2,3);
}

void unflatten(Matrix<double, 1, 12> &y, Matrix<double, 4, 4> &g) {
    g(0,0) = y(0);
    g(0,1) = y(1);
    g(0,2) = y(2);
    g(1,0) = y(3);
    g(1,1) = y(4);
    g(1,2) = y(5);
    g(2,0) = y(6);
    g(2,1) = y(7);
    g(2,2) = y(8);
    g(0,3) = y(9);
    g(1,3) = y(10);
    g(2,3) = y(11);
    g(3,0) = 0;
    g(3,1) = 0;
    g(3,2) = 0;
    g(3,3) = 1;
}

void skew(Matrix<double, 3, 1> &x, Matrix<double, 3, 3> &y) {
    y(0,1) = -x(2);
    y(0,2) = x(1);
    y(1,0) = x(2);
    y(1,2) = -x(0);
    y(2,0) = -x(1);
    y(2,1) = x(0);
    y(0,0) = 0;
    y(1,1) = 0;
    y(2,2) = 0;
}

void adjoint(Matrix<double, 6, 1> &x, Matrix<double, 6, 6> &y) {
    Vector3d w = x.head<3>();
    Vector3d v = x.tail<3>();
    Matrix3d w_skew;
    Matrix3d v_skew;

    skew(w,w_skew);
    skew(v,v_skew);

    y.block<3,3>(0,0) = w_skew;
    y.block<3,3>(3,0) = v_skew;
    y.block<3,3>(0,3) = Matrix3d::Zero();
    y.block<3,3>(3,3) = w_skew;
}

void se(Matrix<double, 6, 1> &x, Matrix<double, 4, 4> &y) {
    Vector3d w = x.head<3>();
    Vector3d v = x.tail<3>();
    Matrix3d w_skew;
    skew(w,w_skew);
    y.block<3,3>(0,0) = w_skew;
    y.block<3,1>(0,3) = v;
    y(3,0) = 0;
    y(3,1) = 0;
    y(3,2) = 0;
    y(3,3) = 0;

}

// initial implementation of rod
class Rod {
    private:
        double D;
        double L;
        double E;
        double G;
        double rho;
        double mu;
        double ds;
        int N;
        Matrix<double, 6, 6> K; //not necessarily diagonal, but for now its alright
        DiagonalMatrix<double, 6> M; //always diagonal
        DiagonalMatrix<double, 6> V; //I suppose not always diagonal
        Matrix<double, 6, 1> xi_ref;
    public:
        Matrix<double, Dynamic, 6> xi;
        Matrix<double, Dynamic, 6> eta;
        Matrix<double, Dynamic, 12> g;
        Rod(double D, double L, double E, double rho, double mu, int N);
        double energy(void);
        void step(double dt);
        VectorXd condition(Rod *prev, double dt, VectorXd xi0);
        void integrate(Rod *prev, double dt, VectorXd xi0);

};

// function object, for some reason called a functor in c++
class LM_Functor : public DenseFunctor<double> {
    double dt;
    Rod *prev;
    Rod &next;
    public:
        LM_Functor(Rod *r, Rod &c, double t) : DenseFunctor<double>(6, 6), prev(r), next(c), dt(t) {};
        int operator()(const VectorXd &x, VectorXd &fvec) const {
            // x is the input for xi0
            // fvec is the tip condition values
            fvec = next.condition(prev, dt, x);
            return 0;
        }
};

Rod::Rod(double d, double l, double e, double r, double m, int n) {
    D = d;
    L = l;
    E = e;
    G = E/3;
    rho = r;
    mu = m;
    N = n;
    ds = L/(N-1);

    double A = M_PI/4*pow(D,2);
    double I = M_PI/64*pow(D,4);
    double J = 2*I;

    //simplest initialization
    // seems like there would be a better way to initialize these
    int i;
    MatrixXd xi_temp(N,6);
    MatrixXd g_temp(N,12);
    for (i=0; i<N; i++) {
        xi_temp(i,0) = 0;
        xi_temp(i,1) = M_PI/(4*10e-2);
        xi_temp(i,2) = 0;
        xi_temp(i,3) = 0;
        xi_temp(i,4) = 0;
        xi_temp(i,5) = 1;

        g_temp(i,0) = 1;
        g_temp(i,1) = 0;
        g_temp(i,2) = 0;
        g_temp(i,3) = 0;
        g_temp(i,4) = 1;
        g_temp(i,5) = 0;
        g_temp(i,6) = 0;
        g_temp(i,7) = 0;
        g_temp(i,8) = 1;
        g_temp(i,9) = 0;
        g_temp(i,10) = 0;
        g_temp(i,11) = ds*i;
    }
    xi = xi_temp;
    g = g_temp;
    eta = MatrixXd::Constant(N,6,0);

    K = Matrix<double, 6, 6>::Zero();

    K.diagonal() << E*I, E*I, G*J, G*A, G*A, E*A;
    M.diagonal() << I, I, J, A, A, A;
    M = rho*M;
    V.diagonal() << 3*I, 3*I, J, A, A, 3*A;
    V = mu*V;

    xi_ref << 0, 0, 0, 0, 0, 1;

}

double Rod::energy(void) {
    double H = 0;
    int i;
    for (i=0; i<N; i++) {
        H += 0.5 * eta.row(i) * M * eta.row(i).transpose();
        H += 0.5 * (xi.row(i).transpose() - xi_ref).transpose() * K * (xi.row(i).transpose() - xi_ref);
    }
    return ds*H;
}

void Rod::step(double dt) {
    // solve the condition
    Rod prev = *this;
    VectorXd guess = xi.row(0);
    LM_Functor functor(&prev, *this, dt);
    DenseIndex nfev;
    int info;
    info = LevenbergMarquardt<LM_Functor>::lmdif1(functor, guess, &nfev);

//    this->integrate(&prev, dt, guess);
}

VectorXd Rod::condition(Rod *prev, double dt, VectorXd xi0) {

    this->integrate(prev, dt, xi0);
    return xi.row(N-1).transpose() - xi_ref;

}

void Rod::integrate(Rod *prev, double dt, VectorXd xi0) {

    Matrix<double, 6, 1> xi_half;
    Matrix<double, 6, 1> eta_half;
    Matrix<double, 6, 1> xi_half_next;
    Matrix<double, 6, 1> eta_half_next;
    Matrix<double, 6, 1> xi_dot;
    Matrix<double, 6, 1> eta_dot;
    Matrix<double, 6, 1> xi_der;
    Matrix<double, 6, 1> eta_der;
    Matrix<double, 6, 1> B_bar;

    Matrix<double, 6, 6> A_bar;
    Matrix<double, 6, 6> xi_ad;
    Matrix<double, 6, 6> eta_ad;

    Matrix4d eta_se;
    Matrix<double, 1, 12> g_row;

    xi.row(0) = xi0.transpose();

    Matrix4d G = Matrix4d::Identity();
    int i;
    for (i=0; i<N-1; i++) {
        xi_half = (xi.row(i) + prev->xi.row(i)).transpose()/2;
        eta_half = (eta.row(i) + prev->eta.row(i)).transpose()/2;

        xi_dot = (xi.row(i) - prev->xi.row(i)).transpose()/dt;
        eta_dot = (eta.row(i) - prev->eta.row(i)).transpose()/dt;

        A_bar = Matrix<double, 6, 6>::Zero();
        B_bar = Matrix<double, 6, 1>::Zero();

        B_bar = V * xi_dot;

        adjoint(xi_half, xi_ad);
        adjoint(eta_half, eta_ad);

        xi_der = K.ldlt().solve(((M * eta_dot) - (eta_ad.transpose() * M * eta_half) + (xi_ad.transpose() * K * (xi_half - xi_ref)) + B_bar));
        eta_der = xi_dot - xi_ad * eta_half;

        xi_half_next = xi_half + ds * xi_der;
        eta_half_next = eta_half + ds * eta_der;

        xi.row(i+1) = (2 * xi_half_next.transpose() - prev->xi.row(i+1));
        eta.row(i+1) = (2 * eta_half_next.transpose() - prev->eta.row(i+1));

    }

//    for (i=0; i<N; i++) {
//        g_row = prev->g.row(i);
//        unflatten(g_row, G);
//        eta_half = (eta.row(i)+prev->eta.row(i)).transpose()/2*ds;
//        se(eta_half, eta_se);
//        G = G * eta_se.exp();
//        flatten(G, g_row);
//        g.row(i) = g_row;
//    }

}

int main() {
    double dt = 0.05;
    Rod r1(1e-2,10e-2,1e6,1e3,0,40);
    cout << r1.energy() << endl;
    int i;
    for (i=0; i<1000; i++) {
        r1.step(dt);
        cout << r1.energy() << endl;
    }

//    Matrix<double, 6,1> xi;
//    xi << 1, 2, 3, 4, 5, 6;
//    Matrix4d xi_se;
//    Matrix<double, 6, 6> xi_ad;
//
//    cout << xi << endl;
//    se(xi,xi_se);
//    cout << xi_se << endl;
//    adjoint(xi,xi_ad);
//    cout << xi_ad << endl;

    return 0;
}


/*
 * Mesh.cpp
 *
 * This class contains the mesh and mesh properties to be used
 * in the nonlinear finite element method (FiniteElemNL.cpp).
 *
 * It constructs the mesh, then constructs the stiffness matrix and
 * calculates the boundary. It can be refined by calling
 * uniformrefine().
 *
 *  Created on: Aug 13, 2016
 *      Author: Ted Kwan
 */

#include "MeshMG.h"
#include "ArmaFuns.h"
#include <cmath>
#include <iostream>
#include <armadillo>

using namespace std;
using namespace arma;

MeshMG::MeshMG() {

}

/**
 * Constructor for MeshMG. It takes one input, a vector containing the
 * properties of the mesh to be created.
 *
 * @param meshprops - double vector with all of the necessary properties.
 */
MeshMG::MeshMG(vector<double> meshprops) {
	// Extract the mesh size value.
	h = meshprops[4];
	n = round(1 / h);
	// Create vectors to be used for mesh creation.
	vec xr = linspace<vec>(meshprops[0], meshprops[1], n + 1);
	vec yr = linspace<vec>(meshprops[2], meshprops[3], n + 1);
	// Create the mesh.
	makeMesh(xr, yr);
	// Assemble stiffness matrix
	vector<sp_mat> stiffmass = assembleMatrix();
	stiffness=stiffmass[0];
	mass=stiffmass[1];
	// Calculate boundary nodes.
	findBoundary();
}

/**
 * makeMeshMG creates the mesh for the 2D Cartesian grid.
 *
 * @param xr - input vector containing the x values.
 * @param yr - input vector containing the y values.
 */
void MeshMG::makeMesh(vec xr, vec yr) {
	// Findthe lengths to use.
	uword xrl = xr.n_rows;
	uword yrl = xr.n_rows;
	mat x = zeros<mat>(xrl, yrl);
	mat y = zeros<mat>(xrl, yrl);
	// Construct the meshes. This is just the same
	// as calling meshgrid() in MATLAB.
	for (uword i = 0; i < yrl; i++) {
		x.row(i) = xr.t();
	}
	for (uword i = 0; i < xrl; i++) {
		y.col(i) = yr;
	}
	// Vectorise the values to use for the nodes matrix.
	vec xv = vectorise(x);
	N = xv.n_rows;
	Ns.push_back(N);
	uword ni = x.n_rows;
	// Initialize node matrix.
	node = zeros<mat>(N, 2);
	// put all the values in the node matrix
	node.col(0) = xv;
	node.col(1) = vectorise(y);
	// Map the indices for the nodes to create elem matrix.
	uvec t2nidxMapnz = regspace<uvec>(0, N - ni - 1);
	uvec topNode = regspace<uvec>(ni - 1, ni, N - ni - 1);
	// Set all of the doubled nodes to not be included.
	t2nidxMapnz(topNode) = zeros<uvec>(topNode.n_rows);
	// Find nonzeros in the original indices map.
	uvec nnz = nonzeros(t2nidxMapnz);
	uvec k = zeros<uvec>(nnz.n_rows + 1);
	// Calculate the different values used to create nodes.
	k(span(1, nnz.n_rows)) = nnz;
	uword NE = k.n_rows;
	HBs.push_back(zeros<umat>(3,3));
	coarse2Fine.push_back(zeros<uvec>(3));
	Pro.push_back(zeros<sp_mat>(3,3));
	// Create the elements which will differ by odd and
	// even elements.
	umat elemup = zeros<umat>(NE, 3);
	umat elemdown = zeros<umat>(NE, 3);
	uvec niv = ones<uvec>(k.n_rows);
	niv = niv * ni;
	uvec onek = ones<uvec>(k.n_rows);
	// Map elements to nodes in order.
	elemup.col(0) = k + niv;
	elemup.col(1) = k + niv + onek;
	elemup.col(2) = k;
	elemdown.col(0) = k + onek;
	elemdown.col(1) = k;
	elemdown.col(2) = k + niv + onek;
	// join all columns together and stack them properly.
	elem = join_cols(elemup, elemdown);
	NT = elem.n_rows;
	refines=0;
}

/**
 * assembleMatrix assembles the stiffness matrix using the quick
 * construction method from armadillo for sparse matrices.
 *
 * @return - Stiffness matrix as a sparse matrix.
 */
vector<sp_mat> MeshMG::assembleMatrix() {
	// Initialize index maps and value vector.
	uvec ii = zeros<uvec>(9 * NT);
	uvec jj = zeros<uvec>(9 * NT);
	vec sA = zeros<vec>(9 * NT);
	// Calculate the area of each node.
	cube ve(NT, 2, 3);
	ve.slice(0) = node.rows(elem.col(2)) - node.rows(elem.col(1));
	ve.slice(1) = node.rows(elem.col(0)) - node.rows(elem.col(2));
	ve.slice(2) = node.rows(elem.col(1)) - node.rows(elem.col(0));
	// Find the area using the dot product on the second dimension.
	area = 0.5*abs((ve.slice(2).col(0) % ve.slice(1).col(1))
					-(ve.slice(2).col(1) % ve.slice(1).col(0)));
	uword index = 0;
	// Loop to map values and indices.
	for (uword i = 0; i < 3; i++) {
		for (uword j = 0; j < 3; j++) {
			// Setup indices to map in this iteration.
			uvec inds = regspace<uvec>(index, index + NT - 1);
			// Setup element maps for indices.
			ii(inds) = elem.col(i);
			jj(inds) = elem.col(j);
			// Calculate values of stiffness matrix at
			// these points.
			mat prod = ve.slice(i) % ve.slice(j);
			// Store calculated value.
			sA(inds) = sum(prod, 1) / (4 * area);
			index = index + NT;
		}
	}
	// Setup index map to be a 2x9NT matrix.
	umat inds = join_horiz(ii, jj);
	// Create the sparse matrix using the same
	// method as sparse(row indices, col indices, values, size)
	sp_mat A(true, inds.t(), sA, N, N, true, true);
	vector<sp_mat> matvec(2);
	matvec[0]=A;
	// Create mass matrix.
	vec Mv=accumArrayM(join_vert(elem.col(0), join_vert(elem.col(1),elem.col(2)))
			,join_vert(area, join_vert(area,area))/(3.0), N);
	uvec inds2=regspace<uvec>(0,Mv.n_rows-1);
	umat subs2=join_horiz(inds2,inds2);
	sp_mat M(true,subs2.t(),Mv,N,N,true,true);
	matvec[1]=M;
	masses.push_back(Mv);
	stiffs.push_back(A);
	return matvec;
}

/**
 * Find the boundary nodes and elements.
 *
 * This method calculates the boundary for the given mesh to be used
 * to set the boundary condition for the finite element method.
 *
 * Values are stored so that they can be accesssed later.
 *
 */
void MeshMG::findBoundary() {

	// Setup as two column vectors to find the edges.
	umat e1 = join_vert(join_horiz(elem.col(2), elem.col(1)),
			join_horiz(elem.col(0), elem.col(2)));
	// Calculate all of the values of the edges.
	umat totalEdge = join_vert(e1, join_horiz(elem.col(1), elem.col(0)));
	totalEdge = sort(totalEdge, "ascend", 1);
	vec onev = ones<vec>(totalEdge.n_rows);
	// Create sparse matrix containing the edges which are not being double counted
	// and are thus exterior edges.
	sp_mat fndmat(true, totalEdge.t(), onev, totalEdge.n_rows, totalEdge.n_rows,
			true, true);
	// Initialize edge matrix.
	umat bdEdge = zeros<umat>(totalEdge.n_rows, 2);
	vec s = nonzeros(fndmat);
	uvec ii = zeros<uvec>(s.n_elem);
	uvec jj = zeros<uvec>(s.n_elem);
	uword k = 0;
	// Get the indices for the edges which are exterior edges.
	// The indices map back to nodes.
	for (uword i = 0; i < fndmat.n_rows; i++) {
		for (uword j = 0; j < fndmat.n_cols; j++) {
			double spot = fndmat(i, j);
			// Only find edges where there are two boundary nodes.
			if (spot == 1.0 && k < s.n_elem) {
				ii(k) = i;
				jj(k) = j;
				k++;
			}
		}
	}
	uvec i1 = nonzeros(ii);
	uvec j1 = nonzeros(jj);
	// Setup boolean vector which has whether or not a node is a boundary
	// node.
	isbdNode = zeros<uvec>(N);
	isbdNode.rows(i1) = ones<uvec>(i1.n_rows);
	isbdNode.rows(j1) = ones<uvec>(j1.n_rows);
	isbdNode(0) = 1.0;
	// Find indices of boundary nodes.
	bdNode = find(isbdNode);
	uvec onebd = ones<uvec>(isbdNode.n_rows);
	// Find indices of interior nodes.
	freeNode = find(onebd - isbdNode);
	freeNodes.push_back(freeNode);

}

/**
 * accumArrayM implements the accumarray method from MATLAB.
 *
 * This is a direct implementation and does the exact same thing
 * as the MATLAB function.
 *
 * @param subs - vector of indices.
 * @param ar - vector of values.
 * @param N - size of the array to be made.
 * @return accumulated array (same as MATLAB would).
 */
vec MeshMG::accumArrayM(uvec subs,vec ar,uword N){

	vec S=zeros<vec>(N);

	for(uword i=0;i<N;i++){
		// Get the subscripts.
		uvec q1=find(subs ==(i));
		vec spot;
		double thesum=0.0;
		if(!q1.is_empty()){
			// Find elements at indices q1
			spot=ar.elem(q1);
			// Sum elements.
			thesum=sum(spot);
			// Set at position i.
			S(i)=thesum;
		}
		// If it is empty, add 0.
		else{
			S(i)=0.0;
		}
	}

	return S;
}

/**
 * uniformrefine refines the mesh and re-calculates the boundary
 * as well as all of the needed matrices and vectors.
 *
 * This function ensures that the mesh is refined by adding nodes at
 * the midpoints of all edges in each element.
 *
 */
void MeshMG::uniformrefine(){
	// Find all of the edges.
	umat e23=join_horiz(elem.col(1),elem.col(2));
	umat e31=join_horiz(elem.col(2),elem.col(0));
	umat e12=join_horiz(elem.col(0),elem.col(1));
	// Matrix with all edge indices (n1 -> n2),
	umat totalEdge=sort(join_vert(e23,join_vert(e31,e12)),"ascend",1);
	ArmaFuns funs;
	// Implementation of unique in MATLAB.
	umat edge=funs.uniqueu(totalEdge);
	uvec js=funs.uniqueidx;
	uword NE=edge.n_rows;
	// Map elems to edges.
	umat elem2edge=zeros<umat>(NT*3,1);
	elem2edge.col(0)=js;
	elem2edge.reshape(NT,3);
	// New nodes to create.
	mat nnode=(node.rows(edge.col(0))+node.rows(edge.col(1)))/(2.0);
	// Concatinate the new nodes with the old.
	node=join_vert(node,nnode);
	// Re-map edges including new nodes.
	uvec edge2newNode=regspace<uvec>(N,N+NE-1);
	// Make hierarchical basis.
	HB=join_horiz(edge2newNode,edge);
	uvec t=regspace<uvec>(0,NT-1);
	// Indices for the new elements.
	umat p= join_horiz(elem,
			join_horiz(edge2newNode.rows(elem2edge.col(0)),
			join_horiz(edge2newNode.rows(elem2edge.col(1)),
					edge2newNode.rows(elem2edge.col(2)))));
	// Add elements to the elem matrix.
	elem.rows(t)=join_horiz(p.col(0),join_horiz(p.col(4),p.col(5)));
	// Create new elements.
	umat nelem=zeros<umat>(4*NT,3);
	nelem.rows(span(NT,2*NT-1))=join_horiz(p.col(5),join_horiz(p.col(1),p.col(3)));
	nelem.rows(span(2*NT,3*NT-1))=join_horiz(p.col(4),join_horiz(p.col(3),p.col(2)));
	nelem.rows(span(3*NT,4*NT-1))=join_horiz(p.col(3),join_horiz(p.col(4),p.col(5)));
	nelem.rows(span(0,NT-1))=elem;
	elem=nelem;
	// Save new information.
	NT=elem.n_rows;
	N=node.n_rows;
	refines=refines+1;
	Ns.push_back(N);
	HBs.push_back(HB);
	ProHB();
	vector<sp_mat> stiffmass = assembleMatrix();
	stiffness=stiffmass[0];
	mass=stiffmass[1];
	findBoundary();
}

/**
 * ProHB creates the prolongation and restriction matrices,
 * as well as maps the fine nodes to the coarse nodes for the FAS.
 *
 * This function ensures that the multigrid method works properly and
 * interpolates the coarse grid to the fine grid, as well as restricts
 * the coarse grid and the fine grid.
 *
 */
void MeshMG::ProHB(){
	// Get the number of coarse nodes.
	uword nCoarse = Ns[Ns.size()-2];
	uword nTotal=Ns[Ns.size()-1];
	// Get the number of fine grid only nodes.
	uword nFineNode=nTotal-nCoarse;
	uvec coarseNode=regspace<uvec>(0,nCoarse-1);
	// Save data for use by FEM.
	coarse2Fine.push_back(coarseNode);
	// Setup indices for the prolongation matrix.
	uvec ii=join_vert(coarseNode,join_vert(HB.col(0),HB.col(0)));
	uvec jj=join_vert(coarseNode,join_vert(HB.col(1),HB.col(2)));
	vec ss=join_vert(ones<vec>(nCoarse),
			join_vert(0.5*ones<vec>(nFineNode),0.5*ones<vec>(nFineNode)));
	// Quick sparse construction of the prolongation matrix.
	sp_mat Procurr(true,join_horiz(ii,jj).t(),ss,nTotal,nCoarse);
	Pro.push_back(Procurr);
}

MeshMG::~MeshMG() {

}