#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <vector>
#include <Eigen/Dense>
#include <iomanip>
#include <cmath>
#include <math.h>
#include <float.h>
#include <algorithm>
#include "ANN.cpp"

using namespace Eigen;

MatrixXd getTril(MatrixXd& mat) {
	MatrixXd tril = mat.triangularView<Lower>();
	for (int i = 0; i < tril.rows(); i++)
		tril(i,i) = 0;
	return tril;
}

void appendVec2Mat(MatrixXd& mat, VectorXd vec) {
	mat.conservativeResize(mat.rows(), mat.cols()+1);
	mat.col(mat.cols()-1) = vec;
}

VectorXd StdVec2Eigen(std::vector <double> vec){
  VectorXd outvec(vec.size());
  for(int i=0;i<vec.size();i++)
    outvec(i)=vec[i];
  return outvec;
}

std::vector <double> EigenVec2Std(VectorXd vec){
  std::vector <double> outvec(vec.rows(),0.0);
  for(int i=0;i<vec.rows();i++)
    outvec[i]=vec(i);
  return outvec;
}

class ObjFunc{
public:
  	void init(ANN& input) {
  		nnet = input;
  		target = VectorXd::Constant(20,0);
  	}
  	void setTarget(VectorXd input) {
  		for (int i = 0; i < 10; i++) {
  			target(2*i) = input(i);
  			target(2*i + 1) = input(i);
  		}
  	}
 //  	VectorXd evalANN(VectorXd input){
 //  		std::vector<double> temperature = nnet.ComputeANN(EigenVec2Std(input));
	// 	VectorXd output(10);
	// 	for(int i=0;i<10;i++){
	// 		output(i)=	0.5*(0.15*temperature[3*i] +
	// 					0.85*temperature[3*i+1] +
	// 					0.85*temperature[3*i+2] +
	// 					0.15*temperature[3*i+3]);
	// 	}
	// 	return output;
	// }
  	VectorXd evalANN(VectorXd input){
  		std::vector<double> temperature = nnet.ComputeANN(EigenVec2Std(input));
		VectorXd output(20);
		for(int i=0;i<10;i++){
			output(2 * i) = (temperature[3*i] + temperature[3*i + 1]) * 0.015 / 2 +
			(1.5 * temperature[3*i + 1] + 0.5 * temperature[3*i + 2]) * 0.035 / 2;
			output(2 * i) /= 0.05;
			output(2 * i + 1) = (temperature[3*i + 2] + temperature[3*i + 3]) * 0.015 / 2 +
			(0.5 * temperature[3*i + 1] + 1.5 * temperature[3*i + 2]) * 0.035 / 2;
			output(2*i + 1) /= 0.05;
		}
		return output;
	}
	void eval(double& f, VectorXd& g, VectorXd input){
		VectorXd output = evalANN(input);
		// std::cout << output << std::endl;
		int n = input.rows();
		f = (output - target).dot(output - target);
		VectorXd gradient = VectorXd::Constant(n,0.0);
		double delta = 0.01;
		// std::cout << "BP\n";
		for (int i = 0; i < n; i++) {
			VectorXd d_input = VectorXd::Constant(n,0.0);
			d_input(i) = delta;
			VectorXd temp_output1 = evalANN(input + d_input);
			VectorXd temp_output2 = evalANN(input - d_input);
			gradient(i) = (temp_output1 - target).dot(temp_output1 - target) -
							(temp_output2 - target).dot(temp_output2 - target);
			gradient(i) /= 2*delta;
		}
		g = gradient;
	}
private:
	ANN nnet;
	VectorXd target;
};

void quickSort(std::vector<double>& arr, std::vector<int>& indices, int left, int right) {
	int i = left, j = right;
	double pivot = arr[(left + right) / 2];
	/*partition*/
	while (i <= j) {
		while (arr[i] < pivot)
			i++;
		while (arr[j] > pivot)
			j--;
		if (i <= j) {
			double tmp_v = arr[i];
			int tmp_i = indices[i];
			arr[i] = arr[j];
			indices[i] = indices[j];
			arr[j] = tmp_v;
			indices[j] = tmp_i;
			i++;
			j--;
		}
	}
	/*recursion*/
	if (left < j)
		quickSort(arr, indices, left, j);
	if (i < right)
		quickSort(arr, indices, i, right);
}

struct LBFGSBParam{
	int m;
	int max_iters;
	double tol;
	bool display;
	bool xhistory;
};

struct BPOutput{
	VectorXd t;
	VectorXd d;
	std::vector<int> F;
};

struct CauchyOutput{
	VectorXd xc;
	VectorXd c;
};

struct SubspaceMinOutput{
	VectorXd xbar;
	bool line_search_flag;
};

struct LBFGSB_Output {
	VectorXd x;
	double obj;
};

class LBFGSB{
public:
	LBFGSBParam param;
	LBFGSB_Output solve(ObjFunc& obj, VectorXd x0, VectorXd l, VectorXd u, LBFGSBParam params);
private:
	double f;
	VectorXd g;
	double getOptimality(VectorXd x, VectorXd g, VectorXd l, VectorXd u);
	BPOutput getBreakpoints(VectorXd x, VectorXd g, VectorXd l, VectorXd u);
	CauchyOutput getCauchyPoint(VectorXd x, VectorXd g, VectorXd l, VectorXd u,
								double theta, MatrixXd W, MatrixXd M);
	double findAlpha(VectorXd l, VectorXd u, VectorXd xc, VectorXd du,
					std::vector<int> free_vars_idx);
	SubspaceMinOutput subspaceMin(VectorXd x, VectorXd g, VectorXd l, VectorXd u,
						VectorXd xc, VectorXd c, double theta, MatrixXd W, MatrixXd M);
	double strongWolfe(ObjFunc& obj, VectorXd x0, double f0, VectorXd g0, VectorXd p);
	double alphaZoom(ObjFunc& obj, VectorXd x0, double f0, VectorXd g0, VectorXd p,
					double alpha_lo, double alpha_hi);
};

double LBFGSB::getOptimality(VectorXd x, VectorXd g, VectorXd l, VectorXd u){
	VectorXd projected_g = x - g;
	for(int i = 0; i < x.rows(); i++) {
		if (projected_g(i) < l(i))
			projected_g(i) = l(i);
		else if (projected_g(i) > u(i))
			projected_g(i) = u(i);
	}
	projected_g = projected_g - x;
	return projected_g.cwiseAbs().maxCoeff();
}

BPOutput LBFGSB::getBreakpoints(VectorXd x, VectorXd g, VectorXd l, VectorXd u){
	int nn = x.rows();
	VectorXd t = VectorXd::Constant(nn,0.0);
	VectorXd d = -g;
	for (int i = 0; i < nn; i++) {
		if (g(i) < 0)
			t(i) = (x(i) - u(i)) / g(i);
		else if (g(i) > 0)
			t(i) = (x(i) - l(i)) / g(i);
		else
			t(i) = DBL_MAX;
		if (t(i) < DBL_EPSILON)
			d(i) = 0.0;
	}
	// Sort elements of t and store indices in F
	std::vector<int> Fvec(nn, 0);
	std::vector<double> tvec(nn,0.0);
	for (int i = 0; i < nn; i++){
		tvec[i] = t(i);
		Fvec[i] = i;
		// std::cout << tvec[i] << " " << Fvec[i] <<std::endl;
	}
	//std::cout<<nn<<std::endl;
	// std::cout<<"BP_quicksort\n";
	quickSort(tvec, Fvec, 0, nn-1);
	// std::cout<<"BP_after_quicksort\n";

	BPOutput bpout;
	bpout.t = t;
	bpout.d = d;
	bpout.F = Fvec;
	return bpout;
}

CauchyOutput LBFGSB::getCauchyPoint(VectorXd x, VectorXd g, VectorXd l, VectorXd u,
								double theta, MatrixXd W, MatrixXd M){

	//std::cout << "BP0" << std::endl;
	BPOutput bpout = getBreakpoints(x, g, l, u);
	// std::cout << "t = " << std::endl << bpout.t << std::endl;
	// std::cout << "d = " << std::endl << bpout.d << std::endl;
	// std::cout << "F = " << std::endl;
	// for (int i = 0; i < bpout.F.size(); i++) {
	// 	std::cout << bpout.F[i] << std::endl;
	// }
	VectorXd xc = x;
	//std::cout << "BP01" << std::endl;
	VectorXd p = W.transpose() * bpout.d;
	VectorXd c = VectorXd::Constant(W.cols(), 0.0);
	double fp = - bpout.d.transpose() * bpout.d;
	double fpp = -theta * fp - p.dot(M * p);
	double fpp0 = -theta * fp;
	double dt_min = -fp / fpp;
	double t_old = 0;
	int i;
	//std::cout << "BP1" << std::endl;
	for (int j = 0; j < x.rows(); j++) {
		i = j;
		if (bpout.F[i] > 0)
			break;
	}
	int b = bpout.F[i];
	double t = bpout.t(b);
	double dt = t - t_old;
	//std::cout << "BP2" << std::endl;

	while ( (dt_min > dt) && (i < x.rows()) ) {
		if (bpout.d(b) > 0)
			xc(b) = u(b);
		else if (bpout.d(b) < 0)
			xc(b) = l(b);
		//std::cout << "BP3" << std::endl;
		double zb = xc(b) - x(b);
		c = c + dt * p;
		double gb = g(b);
		VectorXd wbt = W.row(b);
		fp += dt * fpp + gb * gb + theta * gb * zb - gb * wbt.dot(M * c);
		fpp -= theta * gb * gb - 2.0 * gb * wbt.dot(M * p) - gb * gb * wbt.dot(M * wbt);
		fpp = std::max(DBL_EPSILON * fpp0, fpp);
		p += gb * wbt;
		bpout.d(b) = 0.0;
		dt_min = -fp / fpp;
		t_old = t;
		i++;
		if (i < x.rows()){
			b = bpout.F[i];
			t = bpout.t(b);
			dt = t - t_old;
		}
	}
	// Perform final updates
	dt_min = std::max(dt_min, 0.0);
	t_old = t_old + dt_min;
	for (int j = 0; j < xc.rows(); j++) {
		int idx = bpout.F[j];
		xc(idx) += t_old * bpout.d(idx);
	}
	c += dt_min * p;

	CauchyOutput cpout;
	cpout.c = c;
	cpout.xc = xc;
	return cpout;
}

double LBFGSB::findAlpha(VectorXd l, VectorXd u, VectorXd xc, VectorXd du,
				std::vector<int> free_vars_idx) {
	// INPUTS:
	//  l: [n,1] lower bound constraint vector.
	//  u: [n,1] upper bound constraint vector.
	//  xc: [n,1] generalized Cauchy point.
	//  du: [num_free_vars,1] solution of unconstrained minimization.
	// OUTPUTS:
	//  alpha_star: positive scaling parameter.
	double alpha_star = 1;
	int n = free_vars_idx.size();
	for (int i = 0; i < n; i++) {
		int idx = free_vars_idx[i];
		if (du(i) > 0)
			alpha_star = std::min(alpha_star, (u(idx) - xc(idx)) / du(i));
		else
			alpha_star = std::min(alpha_star, (l(idx) - xc(idx)) / du(i));
	}
	return alpha_star;
}

SubspaceMinOutput LBFGSB::subspaceMin(VectorXd x, VectorXd g, VectorXd l, VectorXd u,
								VectorXd xc, VectorXd c, double theta, MatrixXd W, MatrixXd M){
	SubspaceMinOutput subminout;
	subminout.line_search_flag = true;

	int n = x.rows();
	std::vector<int> free_vars_idx;
	std::vector<VectorXd> Z;
	for (int i = 0; i < xc.rows(); i++) {
		if ((xc(i) != u(i)) && (xc(i) != l(i))) {
			free_vars_idx.push_back(i);
			VectorXd unit = VectorXd::Constant(n,0);
			unit(i) = 1;
			Z.push_back(unit);
		}
	}

	int num_free_vars = free_vars_idx.size();
	if (num_free_vars == 0) {
		subminout.xbar = xc;
		subminout.line_search_flag = false;
		return subminout;
	}

	MatrixXd ZZ = MatrixXd::Zero(n, num_free_vars);
	for (int i = 0; i < num_free_vars; i++)
		ZZ.col(i) = Z[i];

	// compute the reduced gradient of mk restricted to free variables
	MatrixXd WTZ = W.transpose() * ZZ;
	VectorXd rr = g + theta * (xc - x) - W*M*c;
	VectorXd r = VectorXd::Constant(num_free_vars, 0.0);
	for (int i = 0; i < num_free_vars; i++)
		r(i) = rr(free_vars_idx[i]);

	// form intermediate variables
	double invtheata = 1.0 / theta;
	VectorXd v = M * WTZ * r;
	MatrixXd N = invtheata * WTZ * WTZ.transpose();
	int N_size = N.rows();
	N = MatrixXd::Identity(N_size, N_size) - M * N;
	v = N.inverse() * v;
	VectorXd du = -invtheata * r - invtheata * invtheata * WTZ.transpose()*v;

	// find alpha star
	double alpha_star = findAlpha(l, u, xc, du, free_vars_idx);
	VectorXd d_star = alpha_star * du;
	subminout.xbar = xc;
	for (int i = 0; i < num_free_vars; i++) {
		int idx = free_vars_idx[i];
		subminout.xbar(idx) += d_star(i);
	}
	return subminout;
}

double LBFGSB::strongWolfe(ObjFunc& obj, VectorXd x0, double f0, VectorXd g0, VectorXd p){
	double alpha; // return value;
	double c1 = 1e-4;
	double c2 = 0.9;
	double alpha_max = 2.5;
	double alpha_im1 = 0;
	double alpha_i = 1;
	double f_im1 = f0;
	double dphi0 = g0.dot(p);
	int i = 0;
	int max_iters = 20;
	VectorXd x;
	double f_i;
	VectorXd g_i;
	// search for alpha that satisfies strong Wolfe conditions
	while (true) {
		x = x0 + alpha_i * p;
		obj.eval(f_i, g_i, x);
		if ((f_i > f0 + c1 * dphi0) || ( (i > 1) && (f_i >= f_im1) )) {
			alpha = alphaZoom(obj, x0, f0, g0, p, alpha_im1, alpha_i);
			break;
		}
		double dphi = g_i.dot(p);
		if (fabs(dphi) <= -c2 * dphi0) {
			alpha = alpha_i;
			break;
		}
		if (dphi >= 0) {
			alpha = alphaZoom(obj, x0, f0, g0, p, alpha_i, alpha_im1);
			break;
		}

		// update
		alpha_im1 = alpha_i;
		f_im1 = f_i;
		alpha_i += 0.8 * (alpha_max - alpha_i);

		if (i > max_iters) {
			alpha = alpha_i;
			break;
		}
		i++;
	}
	return alpha;
}

double LBFGSB::alphaZoom(ObjFunc& obj, VectorXd x0, double f0, VectorXd g0, VectorXd p,
				double alpha_lo, double alpha_hi) {
	double alpha; // return value
	double c1 = 1e-4;
	double c2 = 0.9;
	int i = 0;
	int max_iters = 20;
	double dphi0 = g0.dot(p);
	double alpha_i;
	VectorXd x;
	double f_i;
	VectorXd g_i;

	while (true) {
		alpha_i = 0.5 * (alpha_lo + alpha_hi);
		alpha = alpha_i;
		x = x0 + alpha_lo * p;
		obj.eval(f_i, g_i, x);
		VectorXd x_lo = x0 + alpha_lo * p;
		double f_lo;
		VectorXd dummy_g;
		obj.eval(f_lo, dummy_g, x_lo);
		if ( (f_i > f0 + c1 * alpha_i * dphi0) || (f_i >= f_lo) )
			alpha_hi = alpha_i;
		else {
			double dphi = g_i.dot(p);
			if (fabs(dphi) <= -c2 * dphi0) {
				alpha = alpha_i;
				break;
			}
			if (dphi * (alpha_hi - alpha_lo) >= 0) {
				alpha_hi = alpha_lo;
			}
			alpha_lo = alpha_i;
		}
		i++;
		if (i > max_iters) {
			alpha = alpha_i;
			break;
		}
	}
	return alpha;
}

LBFGSB_Output LBFGSB::solve(ObjFunc& obj, VectorXd x0, VectorXd l,
							VectorXd u, LBFGSBParam params) {

    int rank=0;
    rank=MPI::COMM_WORLD.Get_rank();

	int n = x0.rows();
	MatrixXd Y(n,0);
	MatrixXd S(n,0);
	MatrixXd W = MatrixXd::Zero(n,1);
	MatrixXd M = MatrixXd::Zero(1,1);
	double theta = 1;

	// initialize obj vars
	VectorXd x = x0;
	double f;
	VectorXd g;
	obj.eval(f, g, x);
	int k = 0;
	double opt;

	if (params.display && rank == 0) {
		std::cout << "iter\t\tf(x)\t\toptimality\n";
		std::cout << "-------------------------------------\n";
		opt = getOptimality(x, g, l, u);
		std::cout << k << "\t\t" << f << "\t\t" << opt << "\t\t" << std::endl;
	}

	MatrixXd xhist(n,0);
	//std::cout << "BP0" << std::endl;
	if (params.xhistory) {
		appendVec2Mat(xhist, x0);
	}

	VectorXd x_old;
	VectorXd g_old;
	VectorXd y;
	VectorXd s;
	while ((getOptimality(x, g, l, u) > params.tol) && (k < params.max_iters) ) {
		x_old = x;
		g_old = g;
		//std::cout << "BP01" << std::endl;

		CauchyOutput cpout = getCauchyPoint(x, g, l, u, theta, W, M);
		VectorXd xc = cpout.xc;
		VectorXd c = cpout.c;
		// std::cout << "xc = " << std::endl;
		// std::cout << xc << std::endl;
		// std::cout << "c = " << std::endl;
		// std::cout << c << std::endl;
		// std::cout << "BP1" << std::endl;
		SubspaceMinOutput subminout = subspaceMin(x, g, l, u, xc, c, theta, W, M);
		//std::cout << "BP2" << std::endl;

		VectorXd xbar = subminout.xbar;
		bool line_search_flag = subminout.line_search_flag;

		double alpha = 1.0;
		if (line_search_flag) {
			alpha = strongWolfe(obj, x, f, g, xbar - x);
		}
		x += alpha * (xbar - x);
		//std::cout << "BP3" << std::endl;

		// update LBFGS data structures
		obj.eval(f, g, x);
		y = g - g_old;
		s = x - x_old;
		double curv = fabs(s.dot(y));
		if (curv < DBL_EPSILON) {
			if (params.display && rank == 0) {
    			std::cout << (" warning: negative curvature detected\n");
    			std::cout << ("          skipping L-BFGS update\n");
    		}
    		k++;
    		continue;
		}

		if (Y.cols() < params.m) {
			appendVec2Mat(Y, y);
			appendVec2Mat(S, s);
		} else {
			Y.block(0, 0, n, params.m - 1) = Y.block(0, 1, n, params.m - 1);
			S.block(0, 0, n, params.m - 1) = S.block(0, 1, n, params.m - 1);
			Y.col(params.m - 1) = y;
			S.col(params.m - 1) = s;
		}
		//std::cout << "BP4" << std::endl;
		theta = y.dot(y) / y.dot(s);
		//std::cout << "BP401" << std::endl;
		MatrixXd temp = Y;
		//std::cout << "BP402" << std::endl;
		temp.conservativeResize(temp.rows(), temp.cols()+S.cols());
		temp.block(0, Y.cols(), S.rows(), S.cols()) = theta * S;
		//std::cout << "BP403" << std::endl;
		W = temp;
		//std::cout << W << std::endl;
		//std::cout << "BP41" << std::endl;
		MatrixXd A = S.transpose() * Y;
		MatrixXd L = getTril(A);
		MatrixXd D = -1 * A.diagonal().asDiagonal();
		int D_size = D.rows();
		int L_size = L.rows();

		//std::cout << "BP42" << std::endl;

		MatrixXd MM(D_size + L_size, D_size + L_size);
		MM.block(0, 0, D_size, D_size) = D;
		MM.block(0, D_size, D_size, L_size) = L.transpose();
		MM.block(D_size, 0, L_size, D_size) = L;
		MM.block(D_size, D_size, L_size, L_size) = theta * S.transpose() * S;
		//std::cout << "BP43" << std::endl;
		M = MM.inverse();
		//std::cout << "BP5" << std::endl;
		// update the iteration
		k++;
		if (params.xhistory) {
			appendVec2Mat(xhist, x);
		}
		if (params.display && rank == 0) {
			opt = getOptimality(x, g, l, u);
			std::cout << k << "\t\t" << f << "\t\t" << opt << "\t\t" << std::endl;
		}
	}
	if (k == params.max_iters && params.display && rank == 0) {
		std::cout << " warning: maximum number of iterations reached\n";
	}

	if (getOptimality(x,g,l,u) < params.tol && params.display && rank == 0) {
		std::cout << " stopping because convergence tolerance met!\n";
	}

	LBFGSB_Output out;
	out.x = x;
	out.obj = f;
	return out;
}

//705.023 701.723 697.852 691.804  692.91 667.241 667.139 672.044 674.512 680.647

// int main(int argc, char** argv){
// 	VectorXd x(10);
// 	x = VectorXd::Constant(10,150);
// 	VectorXd l = VectorXd::Constant(10,0.0);
// 	VectorXd u = VectorXd::Constant(10,420.0);

// 	LBFGSBParam params;
// 	params.m = 10;
// 	params.tol = 1e-2;
// 	params.max_iters = 50;
// 	params.display = true;
// 	params.xhistory = false;

// 	ObjFunc obj;
// 	obj.init();
// 	VectorXd base = VectorXd::Constant(10,673);
// 	VectorXd increment(10);
// 	increment << 40, 40, 40, 40, 40, 20, 20, 20, 20, 20;
// 	obj.setTarget(base + increment);
// 	// obj.target << 50, 30, 60, 50, 60, 70, 1, 90, 50, 50;
// 	LBFGSB opt;
// 	std::cout << "starting to solve\n";
// 	x = opt.solve(obj, x, l, u, params);
// 	std::cout << x.transpose() << std::endl;

// 	VectorXd output = obj.evalANN(x);
// 	std::cout << (output - base).transpose() << std::endl;

// 	// LBFGSB optimizer;
// 	// // double opt = optimizer.getOptimality(x,g,l,u);
// 	// // BPOutput bpout = optimizer.getBreakpoints(x,g,l,u);
// 	// // std::cout<<bpout.t.transpose()<<std::endl;
// 	// // std::cout<<bpout.d.transpose()<<std::endl;
// 	// // for (int i = 0; i < bpout.F.size(); i++)
// 	// // 	std::cout<<bpout.F[i]<<"  ";
// 	// // std::cout<<std::endl;
// 	// MatrixXd W = MatrixXd::Random(10,20);
// 	// MatrixXd M = MatrixXd::Identity(20,20);
// 	// VectorXd row = M.row(10);
// 	// std::cout<<row<<std::endl;

// 	// CauchyOutput cpout = optimizer.getCauchyPoint(x,g,l,u, 0.1, W, M);
// 	// std::cout<<cpout.xc.transpose()<<std::endl;
// 	// std::cout<<cpout.c.transpose()<<std::endl;
// 	// using namespace std;
// 	// MatrixXd mat(2,0);
// 	// VectorXd vec(2);
// 	// vec << 1, 1;
// 	// cout << mat.rows() << " " << mat.cols() << endl;
// 	// mat.conservativeResize(mat.rows(), mat.cols()+1);
// 	// mat.col(mat.cols()-1) = vec;
// 	// cout << mat << endl;
// 	return 0;
// }


#include <Eigen/Dense>
#include <iostream>
#include <chrono>
#include<deformation_measures.h>

typedef std::chrono::high_resolution_clock Clock;
typedef std::vector<double> stdv;

int main(){
    /*
    ==============
    |    main    |
    ==============

    Tests different methods for computing tensor multiplications and reports the
    relative speeds.
    */

    //Deformation measures
    double J = 1.2;
    Matrix_3x3 F;
    Matrix_3x3 chi;

    //Stress measures
    Vector_9  PK2_voigt;
    Vector_9  SIGMA_voigt;
    Vector_27 M_voigt;

    Vector_9  cauchy_voigt;
    Vector_9  s_voigt;
    Vector_27 m_voigt;

    //PK2 Jacobians
    Matrix_9x9  dPK2dF;
    Matrix_9x9  dPK2dchi;
    Matrix_9x27 dPK2dgrad_chi;

    //SIGMA Jacobians
    Matrix_9x9  dSIGMAdF;
    Matrix_9x9  dSIGMAdchi;
    Matrix_9x27 dSIGMAdgrad_chi;

    //Higher order stress Jacobians
    Matrix_27x9  dMdF;
    Matrix_27x9  dMdchi;
    Matrix_27x27 dMdgrad_chi;

    //Cauchy jacobians
    Matrix_9x9  dcauchydF;
    Matrix_9x9  dcauchydchi;
    Matrix_9x27 dcauchydgrad_chi;

    //s Jacobians
    Matrix_9x9  dsdF;
    Matrix_9x9  dsdchi;
    Matrix_9x27 dsdgrad_chi;

    //m Jacobians
    Matrix_27x9  dmdF;
    Matrix_27x9  dmdchi;
    Matrix_27x27 dmdgrad_chi;

    //Set random values
    F            = Matrix_3x3::Random();
    chi          = Matrix_3x3::Random();

    PK2_voigt    = Vector_9::Random();
    SIGMA_voigt  = Vector_9::Random();
    M_voigt      = Vector_27::Random();

    cauchy_voigt = Vector_9::Random();
    s_voigt      = Vector_9::Random();
    m_voigt      = Vector_27::Random();

    dPK2dF        = Matrix_9x9::Random();
    dPK2dchi      = Matrix_9x9::Random();
    dPK2dgrad_chi = Matrix_9x27::Random();

    dSIGMAdF        = Matrix_9x9::Random();
    dSIGMAdchi      = Matrix_9x9::Random();
    dSIGMAdgrad_chi = Matrix_9x27::Random();

    dMdF        = Matrix_27x9::Random();
    dMdchi      = Matrix_27x9::Random();
    dMdgrad_chi = Matrix_27x27::Random();

    int n = 10000;

    auto t0 = Clock::now();
    auto t1 = Clock::now();

    t0 = Clock::now();
    for (int _n=0; _n<n; _n++){
        deformation_measures::map_jacobians_to_current_configuration(F, chi,
            PK2_voigt,    SIGMA_voigt, M_voigt,
            cauchy_voigt, s_voigt,     m_voigt,
            dPK2dF,       dPK2dchi,    dPK2dgrad_chi,
            dSIGMAdF,     dSIGMAdchi,  dSIGMAdgrad_chi,
            dMdF,         dMdchi,      dMdgrad_chi,
            dcauchydF,    dcauchydchi, dcauchydgrad_chi,
            dsdF,         dsdchi,      dsdgrad_chi,
            dmdF,         dmdchi,      dmdgrad_chi);
    }
    t1 = Clock::now();
    std::cout << "### Map Jacobians ###\n";
    std::cout << "   Eigen Matrices: " << std::chrono::duration_cast<std::chrono::nanoseconds>(t1 - t0).count()/((double)n) << "\n";

    t0 = Clock::now();
    for (int _n = 0; _n<n; _n++){
        deformation_measures::map_dAdgrad_chi_to_dadgrad_chi(dMdgrad_chi, J, F, chi, dmdgrad_chi);
    }
    t1 = Clock::now();
    std::cout << "### map_dAdgrad_chi_to_dadgrad_chi ###\n";
    std::cout << "    Full nested loops: " << std::chrono::duration_cast<std::chrono::nanoseconds>(t1 - t0).count()/((double)n) << "\n";

    t0 = Clock::now();
    for (int _n = 0; _n<n; _n++){
        deformation_measures::map_dAdgrad_chi_to_dadgrad_chi_alt(dMdgrad_chi, J, F, chi, dmdgrad_chi);
    }
    t1 = Clock::now();
    std::cout << "    Alternative nesting: " << std::chrono::duration_cast<std::chrono::nanoseconds>(t1 - t0).count()/((double)n) << "\n";

    return 1;
}


#ifndef MODULE_OPTIMIZER_HPP
#define MODULE_OPTIMIZER_HPP

#define _USE_MATH_DEFINES
#define M_PI 3.14159265358979323846

#include <math.h>
#include <Eigen/Core>
#include <algorithm>
#include <functional>
#include <numeric>
#include <vector>
#include "main.h"
#include "utility/io/uart.hpp"
#include "utility/math/saturation.hpp"

namespace module {
namespace MPC {
    namespace AdaptiveMPC {

        class Optimizer {
        public:
            Optimizer(int ch_max);

            template <typename T>
            float ProcessModel(const T& model,
                               int mode,
                               std::vector<float>& states,
                               std::vector<float>& setpoint,
                               std::vector<float>& input_states,
                               std::vector<float>& output_states) {

                utility::io::debug.out("Process model mode %d\n", mode);

                // Given a control action, select the correct model matrix, from that linearize the model about the
                // current state point. Calculate the next state vector. Call the control error function and return the
                // errors
                Eigen::Matrix<float, 4, 4> A_mat;

                if (mode == 1) {
                    // Negative
                    Linearise(model,
                              states,
                              -dotV_a(states[2], model.P_a, model.muscle1),
                              dotV_a(model.P_t, states[3], model.muscle2),
                              -dotV_a_P_m1(states[2], model.P_a, model.muscle1, mode),
                              dotV_a_P_m2(model.P_t, states[3], model.muscle2, mode),
                              A_mat);
                }
                else if (mode == 2) {
                    // Positive
                    Linearise(model,
                              states,
                              dotV_a(model.P_t, states[2], model.muscle1),
                              -dotV_a(states[3], model.P_a, model.muscle2),
                              dotV_a_P_m1(model.P_t, states[2], model.muscle1, mode),
                              -dotV_a_P_m2(states[3], model.P_a, model.muscle2, mode),
                              A_mat);
                }
                else if (mode == 3) {
                    // No move
                    Linearise(model,
                              states,
                              dotV_a(model.P_t, states[2], model.muscle1),
                              dotV_a(model.P_t, states[3], model.muscle2),
                              dotV_a_P_m1(model.P_t, states[2], model.muscle1, mode),
                              dotV_a_P_m2(model.P_t, states[3], model.muscle2, mode),
                              A_mat);
                }
                else {
                    Error_Handler();
                }

                utility::io::debug.out("A mat\n");
                utility::io::debug.out(
                    "%.2f\t %.2f\t %.2f\t %.2f\n", A_mat(0, 0), A_mat(0, 1), A_mat(0, 2), A_mat(0, 3));
                utility::io::debug.out(
                    "%.2f\t %.2f\t %.2f\t %.2f\n", A_mat(1, 0), A_mat(1, 1), A_mat(1, 2), A_mat(1, 3));
                utility::io::debug.out(
                    "%.2f\t %.2f\t %.2f\t %.2f\n", A_mat(2, 0), A_mat(2, 1), A_mat(2, 2), A_mat(2, 3));
                utility::io::debug.out(
                    "%.2f\t %.2f\t %.2f\t %.2f\n", A_mat(3, 0), A_mat(3, 1), A_mat(3, 2), A_mat(3, 3));

                float Sampling_time2 = 0.05;  // 0.01 T_s

                Eigen::Matrix<float, 4, 1> x_state(input_states[0], input_states[1], input_states[2], input_states[3]);

                Eigen::Matrix<float, 4, 1> x_states_update =
                    (Eigen::Matrix<float, 4, 4>::Identity() + Sampling_time2 * A_mat) * x_state;

                output_states.push_back(
                    utility::math::sat(x_states_update(0, 0), std::make_pair(M_PI / 2.0, -M_PI / 2.0)));
                output_states.push_back(utility::math::sat(x_states_update(1, 0), std::make_pair(10, -10)));
                output_states.push_back(utility::math::sat(x_states_update(2, 0), std::make_pair(413685, 0)));
                output_states.push_back(utility::math::sat(x_states_update(3, 0), std::make_pair(413685, 0)));

                return (ControlError(setpoint, output_states));
            }

            float ControlError(std::vector<float>& setpoint, std::vector<float>& output_states) {

                // Do the math to find the relative errors (state, input)

                std::vector<float> state_error = {setpoint[0] - output_states[0],
                                                  setpoint[1] - output_states[1],
                                                  setpoint[2] - output_states[2],
                                                  setpoint[3] - output_states[3]};

                std::vector<float>::const_iterator i = output_states.begin();
                std::vector<float>::const_iterator j = setpoint.begin();
                std::vector<float>::const_iterator k = state_error.begin();
                utility::io::debug.out("%f\t - %f\t = %f\n", *j, *i, *k);
                i++;
                j++;
                k++;
                utility::io::debug.out("%f\t - %f\t = %f\n", *j, *i, *k);
                i++;
                j++;
                k++;
                utility::io::debug.out("%f\t - %f\t = %f\n", *j, *i, *k);
                i++;
                j++;
                k++;
                utility::io::debug.out("%f\t - %f\t = %f\n", *j, *i, *k);
                // TODO Fix
                Eigen::Matrix<float, 4, 1> x_error(state_error[0], state_error[1], state_error[2], state_error[3]);
                return (x_error.transpose() * state_weight * x_error);
            }

            template <typename T>
            float dotV_a(float P_1, float P_2, const T& m) {
                float b       = m.critical_ratio;
                float C       = m.sonic_conductance;
                float A       = C * std::sqrt(m.T_0 / m.T_1);
                float P_b     = P_2 / P_1 - b;
                float b_1     = 1 - b;
                float P_bsqrt = std::sqrt(1 - (P_b / b_1) * (P_b / b_1));

                if (P_2 / P_1 > b) {
                    return std::isinf(P_1 * A * P_bsqrt) ? 0 : P_1 * A * P_bsqrt;
                }
                return std::isinf(P_1 * A) ? 0 : P_1 * A;
            }

            template <typename T>
            float dotV_a_P_m1(float P_1, float P_2, const T& m, int mode) {
                float b       = m.critical_ratio;
                float C       = m.sonic_conductance;
                float A       = C * std::sqrt(m.T_0 / m.T_1);
                float P_b     = P_2 / P_1 - b;
                float i       = A * P_b;
                float b_1     = 1 - b;
                float P_bsqrt = std::sqrt(1 - (P_b / b_1) * (P_b / b_1));

                if (mode == 2 || mode == 3) {
                    // P_2 == P_m1
                    if (P_2 / P_1 > b) {
                        return std::isinf(-(i / (b_1 * b_1 * P_bsqrt))) ? 0 : -(i / (b_1 * b_1 * P_bsqrt));
                    }
                    return 0;
                }
                else {
                    // P_1 == P_m1
                    if (P_2 / P_1 > b) {
                        return std::isinf(A * P_bsqrt + (i * P_2 / (b_1 * b_1 * P_1 * P_bsqrt)))
                                   ? 0
                                   : A * P_bsqrt + (i * P_2 / (b_1 * b_1 * P_1 * P_bsqrt));
                    }

                    return std::isinf(C * std::sqrt(m.T_0 / m.T_1)) ? 0 : C * std::sqrt(m.T_0 / m.T_1);
                }
                return 0;
            }

            template <typename T>
            float dotV_a_P_m2(float P_1, float P_2, const T& m, int mode) {
                float b       = m.critical_ratio;
                float C       = m.sonic_conductance;
                float A       = C * std::sqrt(m.T_0 / m.T_1);
                float P_b     = P_2 / P_1 - b;
                float i       = A * P_b;
                float b_1     = 1 - b;
                float P_bsqrt = std::sqrt(1 - (P_b / b_1) * (P_b / b_1));

                if (mode == 1 || mode == 3) {
                    // P_2 == P_m2
                    if (P_2 / P_1 > b) {
                        // utility::io::debug.out("dotV_a_P_m2 mode 13 >\n");
                        return std::isinf(-i / (b_1 * b_1 * P_bsqrt)) ? 0 : -i / (b_1 * b_1 * P_bsqrt);
                    }
                    // utility::io::debug.out("dotV_a_P_m2 mode 13 <\n");
                    return 0;
                }
                else {
                    // P_1 == P_m2
                    if (P_2 / P_1 > b) {
                        // utility::io::debug.out("dotV_a_P_m2 mode 2 >\n");
                        return std::isinf(A * P_bsqrt + ((P_2 * i) / (b_1 * b_1 * P_1 * P_bsqrt)))
                                   ? 0
                                   : A * P_bsqrt + ((P_2 * i) / (b_1 * b_1 * P_1 * P_bsqrt));
                    }

                    // utility::io::debug.out("dotV_a_P_m2 mode 2 <\n");
                    return std::isinf(C * std::sqrt(m.T_0 / m.T_1)) ? 0 : C * std::sqrt(m.T_0 / m.T_1);
                }
                return 0;
            }

            template <typename T>
            void Linearise(const T& m,
                           std::vector<float>& states,
                           float dotV_a1,
                           float dotV_a2,
                           float ddotV_a1,
                           float ddotV_a2,
                           Eigen::Matrix<float, 4, 4>& A_mat) {

                float L_10 = m.muscle1.L_0;
                float L_20 = m.muscle2.L_0;
                float k_10 = m.muscle1.K_0;
                float k_20 = m.muscle2.K_0;
                float y    = states[0];
                float dy   = states[1];
                float P_m1 = states[2];
                float P_m2 = states[3];
                float P_a  = m.P_a;
                float a1   = m.muscle1.muscle_coefficients[0];
                float b1   = m.muscle1.muscle_coefficients[1];
                float c1   = m.muscle1.muscle_coefficients[2];
                float d1   = m.muscle1.muscle_coefficients[3];
                float a2   = m.muscle2.muscle_coefficients[0];
                float b2   = m.muscle2.muscle_coefficients[1];
                float c2   = m.muscle2.muscle_coefficients[2];
                float d2   = m.muscle2.muscle_coefficients[3];
                auto F_ce1 = m.muscle1.F_ce;
                auto F_ce2 = m.muscle2.F_ce;
                float R1   = m.muscle1.damping_coefficient;
                float R2   = m.muscle2.damping_coefficient;
                float mass = m.mass;

                Eigen::Matrix<float, 6, 1> P_ce1;
                Eigen::Matrix<float, 6, 1> P_ce2;
                Eigen::Matrix<float, 6, 1> dP_ce1;
                Eigen::Matrix<float, 6, 1> dP_ce2;

                Eigen::Matrix<float, 1, 6> k_ce1;
                Eigen::Matrix<float, 1, 6> k_ce2;
                Eigen::Matrix<float, 1, 6> dk_ce1;
                Eigen::Matrix<float, 1, 6> dk_ce2;

                float F_s1;
                float F_s2;
                float F_d1;
                float F_d2;
                float V_m1;
                float V_m2;
                float dV_m1;
                float dV_m2;
                float dotV_m1;
                float dotV_m2;
                float ddotV_m1;
                float ddotV_m2;
                float ddotV_m1_dV_m1;
                float ddotV_m2_dV_m2;

                /* clang-format off */

                float k1    = k_10 + y / L_10;
                float k2    = k_20 - y / L_20;
                float dotk1 = k_10 + dy / L_10;
                float dotk2 = k_20 - dy / L_20;

                P_ce1 << 1,
                         P_m1,
                         (P_m1 * P_m1),
                         (P_m1 * P_m1 * P_m1),
                         (P_m1 * P_m1 * P_m1 * P_m1),
                         (P_m1 * P_m1 * P_m1 * P_m1 * P_m1);
                P_ce2 << 1,
                         P_m2,
                         (P_m2 * P_m2),
                         (P_m2 * P_m2 * P_m2),
                         (P_m2 * P_m2 * P_m2 * P_m2),
                         (P_m2 * P_m2 * P_m2 * P_m2 * P_m2);
                dP_ce1 << 0,
                          1,
                          (2 * P_m1),
                          (3 * P_m1 * P_m1),
                          (4 * P_m1 * P_m1 * P_m1),
                          (5 * P_m1 * P_m1 * P_m1 * P_m1);
                dP_ce2 << 0,
                          1,
                          (2 * P_m2),
                          (3 * P_m2 * P_m2),
                          (4 * P_m2 * P_m2 * P_m2),
                          (5 * P_m2 * P_m2 * P_m2 * P_m2);

                k_ce1 << 1,
                         k1,
                         (k1 * k1),
                         (k1 * k1 * k1),
                         (k1 * k1 * k1 * k1),
                         (k1 * k1 * k1 * k1 * k1);
                k_ce2 << 1,
                         k2,
                         (k2 * k2),
                         (k2 * k2 * k2),
                         (k2 * k2 * k2 * k2),
                         (k2 * k2 * k2 * k2 * k2);
                dk_ce1 << 0,
                          (1 / L_10),
                          (2 * k1 / L_10),
                          (3 * k1 * k1 / L_10),
                          (4 * k1 * k1 * k1 / L_10),
                          (5 * k1 * k1 * k1 * k1 / L_10);
                dk_ce2 << 0,
                          (-1 / L_20),
                          (-2 * k2 / L_20),
                          (-3 * k2 * k2 / L_20),
                          (-4 * k2 * k2 * k2 / L_20),
                          (-5 * k2 * k2 * k2 * k2 / L_20);

                /******************************************************************************************************/
                /********************************************** Column 1 **********************************************/
                /******************************************************************************************************/

                /********************************************* Lin_mat_0_0 ********************************************/
                float Lin_mat_0_0 = 0;
                /********************************************* Lin_mat_1_0 ********************************************/

                // Calculate our first spring force
                F_s1 = dk_ce1 * F_ce1 * P_ce1;

                // Calculate our second spring force
                F_s2 = dk_ce2 * F_ce2 * P_ce2;

                float Lin_mat_1_0 = (-F_s2 - F_s1) / mass;

                /********************************************* Lin_mat_2_0 ********************************************/
                dV_m1          = (3 * a1 * k1 * k1 + 2 * b1 * k1 + c1) / L_10;
                ddotV_m1_dV_m1 =
                    ((a1 * k1 * k1 * k1 + b1 * k1 * k1 + c1 * k1 + d1)
                        * (6 * a1 * k1 * dotk1 + 2 * b1 * dotk1)
                    - (3 * a1 * k1 * k1 * dotk1 + 2 * b1 * k1 * dotk1 + c1 * dotk1)
                        * (3 * a1 * k1 * k1 * dotk1 + 2 * b1 * k1 * dotk1 + c1 * dotk1))
                    / (L_10 * (a1 * k1 * k1 * k1 + b1 * k1 * k1 + c1 * k1 + d1)
                       * (a1 * k1 * k1 * k1 + b1 * k1 * k1 + c1 * k1 + d1));

                float Lin_mat_2_0 = (P_a * dotV_a1 / dV_m1) - (P_m1 * ddotV_m1_dV_m1);
                /********************************************* Lin_mat_3_0 ********************************************/
                dV_m2          = (3 * a2 * k2 * k2 + 2 * b2 * k2 + c2) / L_20;
                ddotV_m2_dV_m2 =
                    ((a2 * k2 * k2 * k2 + b2 * k2 * k2 + c2 * k2 + d2)
                        * (6 * a2 * k2 * dotk2 + 2 * b2 * dotk2)
                    - (3 * a2 * k2 * k2 * dotk2 + 2 * b2 * k2 * dotk2 + c2 * dotk2)
                        * (3 * a2 * k2 * k2 * dotk2 + 2 * b2 * k2 * dotk2 + c2 * dotk2))
                    / (L_20 * (a2 * k2 * k2 * k2 + b2 * k2 * k2 + c2 * k2 + d2)
                       * (a2 * k2 * k2 * k2 + b2 * k2 * k2 + c2 * k2 + d2));

                float Lin_mat_3_0 = (P_a * dotV_a2 / dV_m2) - (P_m2 * ddotV_m2_dV_m2);
                /******************************************************************************************************/
                /********************************************** Column 2 **********************************************/
                /******************************************************************************************************/

                /********************************************* Lin_mat_0_1 ********************************************/
                float Lin_mat_0_1 = 1;
                /********************************************* Lin_mat_1_1 ********************************************/
                F_d1 = -R1 * P_m1 / L_10;
                F_d2 = -R2 * P_m2 / L_20;

                float Lin_mat_1_1 = (-F_d2 - F_d1) / mass;
                /********************************************* Lin_mat_2_1 ********************************************/
                V_m1     = (a1 * k1 * k1 * k1) + (b1 * k1 * k1) + (c1 * k1) + d1;
                ddotV_m1 = (3 * a1 * k1 * k1 + 2 * b1 * k1 + c1) / L_10;

                float Lin_mat_2_1 = -P_m1 * ddotV_m1 / V_m1;
                /********************************************* Lin_mat_3_1 ********************************************/
                V_m2     = (a2 * k2 * k2 * k2) + (b2 * k2 * k2) + (c2 * k2) + d2;
                ddotV_m2 = (3 * a2 * k2 * k2 + 2 * b2 * k2 + c2) / L_20;

                float Lin_mat_3_1 = -P_m2 * ddotV_m2 / V_m2;
                /******************************************************************************************************/
                /********************************************** Column 3 **********************************************/
                /******************************************************************************************************/

                /********************************************* Lin_mat_0_2 ********************************************/
                float Lin_mat_0_2 = 0;
                /********************************************* Lin_mat_1_2 ********************************************/
                F_s1 = k_ce1 * F_ce1 * dP_ce1;
                F_d1 = -R1 * dotk1;

                utility::io::debug.out("F_s1 = %f\nk_ce1 = %f\n F_ce1 = %f\n dP_ce1 = %f\n\n", F_s1, k_ce1, F_ce1, dP_ce1);
                utility::io::debug.out("F_d1 = %f\n -R1 = %f\n dotk1 = %f\n", F_d1, -R1, dotk1);

                // TODO This is incorrect, but i don't know any better
                float Lin_mat_1_2 = (F_s1 + F_d1) / mass;
                /********************************************* Lin_mat_2_2 ********************************************/
                dotV_m1 = (3 * a1 * k1 * k1 * dotk1 + 2 * b1 * k1 * dotk1 + c1 * dotk1) ;
                V_m1    = (a1 * k1 * k1 * k1) + (b1 * k1 * k1) + (c1 * k1) + d1;

                float Lin_mat_2_2 = (P_a * ddotV_a1 - dotV_m1) / V_m1 + 0.34;

                /********************************************* Lin_mat_3_2 ********************************************/
                float Lin_mat_3_2 = 0;
                /******************************************************************************************************/
                /********************************************** Column 4 **********************************************/
                /******************************************************************************************************/

                /********************************************* Lin_mat_0_3 ********************************************/
                float Lin_mat_0_3 = 0;
                /********************************************* Lin_mat_1_3 ********************************************/
                F_s2 = k_ce2 * F_ce2 * dP_ce2;
                F_d2 = -R2 * dotk2;

                utility::io::debug.out("F_s2 = %f\nk_ce2 = %f\n F_ce2 = %f\n dP_ce2 = %f\n\n", F_s2, k_ce2, F_ce2, dP_ce2);
                utility::io::debug.out("F_d2 = %f\n -R2 = %f\n dotk2 = %f\n", F_d2, -R2, dotk2);

                float Lin_mat_1_3 = (-F_s2 - F_d2) / mass;
                /********************************************* Lin_mat_2_3 ********************************************/
                float Lin_mat_2_3 = 0;
                /********************************************* Lin_mat_3_3 ********************************************/
                dotV_m2 = (3 * a2 * k2 * k2 * dotk2 + 2 * b2 * k2 * dotk2 + c2 * dotk2) ;
                V_m2    = (a2 * k2 * k2 * k2) + (b2 * k2 * k2) + (c2 * k2) + d2;

                float Lin_mat_3_3 = (P_a * ddotV_a2 - dotV_m2) / V_m2 + 0.34;

                A_mat << Lin_mat_0_0, Lin_mat_0_1, Lin_mat_0_2, Lin_mat_0_3,
                         Lin_mat_1_0, Lin_mat_1_1, Lin_mat_1_2, Lin_mat_1_3,
                         Lin_mat_2_0, Lin_mat_2_1, Lin_mat_2_2, Lin_mat_2_3,
                         Lin_mat_3_0, Lin_mat_3_1, Lin_mat_3_2, Lin_mat_3_3;
                /* clang-format on */
            }

            template <typename T>
            std::pair<bool, bool> FirstLayer(const T& m, std::vector<float>& states, std::vector<float>& setpoint) {
                // Increment the depth (control horizon itt)
                // for (std::vector<float>::const_iterator i = states.begin(); i != states.end(); ++i) {
                //     utility::io::debug.out("%f ", *i);
                // }
                // utility::io::debug.out("\n");

                ch_itt = 1;

                // Remove our previous results
                cost_result.clear();

                // Create a cost for each root
                float cost_root_1;
                float cost_root_2;
                float cost_root_3;

                // Create a output state vector for each process
                std::vector<float> output_states_1;
                std::vector<float> output_states_2;
                std::vector<float> output_states_3;

                // Calculate the result of performing each action and add the result error to the cost vector
                cost_root_1 = ProcessModel(m, 1, states, setpoint, states, output_states_1);
                cost_root_2 = ProcessModel(m, 2, states, setpoint, states, output_states_2);
                cost_root_3 = ProcessModel(m, 3, states, setpoint, states, output_states_3);
                // utility::io::debug.out("output_states_1\n");
                // for (std::vector<float>::const_iterator i = output_states_1.begin(); i != output_states_1.end(); ++i)
                // {
                //     utility::io::debug.out("%f ", *i);
                // }
                // utility::io::debug.out("\n");
                // utility::io::debug.out("output_states_2\n");
                // for (std::vector<float>::const_iterator i = output_states_2.begin(); i != output_states_2.end(); ++i)
                // {
                //     utility::io::debug.out("%f ", *i);
                // }
                // utility::io::debug.out("\n");
                // utility::io::debug.out("output_states_3\n");
                // for (std::vector<float>::const_iterator i = output_states_3.begin(); i != output_states_3.end(); ++i)
                // {
                //     utility::io::debug.out("%f ", *i);
                // }
                // utility::io::debug.out("\n");

                // Decide if the next layer is the last or not
                if (ch_itt >= ch_max - 1) {
                    // Must be on our last layer
                    FinalLayer(m, states, setpoint, output_states_1, 1, cost_root_1);
                    FinalLayer(m, states, setpoint, output_states_2, 2, cost_root_2);
                    FinalLayer(m, states, setpoint, output_states_3, 3, cost_root_3);
                }
                else {
                    // Not our last layer lets add a general layer
                    AddLayer(m, states, setpoint, output_states_1, 1, cost_root_1);
                    AddLayer(m, states, setpoint, output_states_2, 2, cost_root_2);
                    AddLayer(m, states, setpoint, output_states_3, 3, cost_root_3);
                }

                // Now look through the cost_result vector and pick the lowest cost to perform the root action
                auto result = *std::min_element(cost_result.cbegin(),
                                                cost_result.cend(),
                                                [](const auto& lhs, const auto& rhs) { return lhs.first < rhs.first; });
                utility::io::debug.out("Cost results\n");
                for (const auto& p : cost_result) {
                    utility::io::debug.out("Mode %d, cost %f\n", p.second, p.first);
                }
                utility::io::debug.out("\n");

                if (result.second == 1) {
                    utility::io::debug.out("Optimizer result mode1, %f\n", result.first);
                    return (std::make_pair(true, false));
                }
                else if (result.second == 2) {
                    utility::io::debug.out("Optimizer result mode2, %f\n", result.first);
                    return (std::make_pair(false, true));
                }
                else if (result.second == 3) {
                    utility::io::debug.out("Optimizer result mode3, %f\n", result.first);
                    return (std::make_pair(false, false));
                }
                utility::io::debug.out("Optimizer failed\n");
                return (std::make_pair(false, false));
            }

            template <typename T>
            void AddLayer(const T& m,
                          std::vector<float>& states,
                          std::vector<float>& setpoint,
                          std::vector<float>& input_states,
                          int root,
                          float cost) {
                // We're somewhere in the middle of our recursion
                ch_itt += 1;

                // Create a output state vector for each process
                std::vector<float> output_states_1;
                std::vector<float> output_states_2;
                std::vector<float> output_states_3;

                // Calculate the result of performing each action and add the result error to the cost vector
                float cost1 = cost + ProcessModel(m, 1, states, setpoint, input_states, output_states_1);
                float cost2 = cost + ProcessModel(m, 2, states, setpoint, input_states, output_states_2);
                float cost3 = cost + ProcessModel(m, 3, states, setpoint, input_states, output_states_3);

                // Decide if the next layer is the last or not
                if (ch_itt >= ch_max - 1) {
                    // Must be on our last layer
                    FinalLayer(m, states, setpoint, output_states_1, root, cost1);
                    FinalLayer(m, states, setpoint, output_states_2, root, cost2);
                    FinalLayer(m, states, setpoint, output_states_3, root, cost3);
                }
                else {
                    // Not our last layer lets add a general layer
                    AddLayer(m, states, setpoint, output_states_1, root, cost1);
                    AddLayer(m, states, setpoint, output_states_2, root, cost2);
                    AddLayer(m, states, setpoint, output_states_3, root, cost3);
                }

                // As we're leaving a level, decrement the itterator
                ch_itt -= 1;
            }

            template <typename T>
            void FinalLayer(const T& m,
                            std::vector<float>& states,
                            std::vector<float>& setpoint,
                            std::vector<float>& input_states,
                            const int& root,
                            float cost) {
                // We're on or last layer, let's calculate the result append the cost and root

                // utility::io::debug.out("input_states\n");
                // for (std::vector<float>::const_iterator i = input_states.begin(); i != input_states.end(); ++i) {
                //     utility::io::debug.out("%f ", *i);
                // }
                // utility::io::debug.out("\n");

                // TODO This isn't needed on the last layer
                // Create a output state vector for each process
                std::vector<float> output_states_1;
                std::vector<float> output_states_2;
                std::vector<float> output_states_3;

                // Calculate the result of performing each action and add the result error to the cost vector
                float cost1 = cost + ProcessModel(m, 1, states, setpoint, input_states, output_states_1);
                float cost2 = cost + ProcessModel(m, 2, states, setpoint, input_states, output_states_2);
                float cost3 = cost + ProcessModel(m, 3, states, setpoint, input_states, output_states_3);

                // utility::io::debug.out("Final Mode %d, Cost1 %f\n", root, cost1 - cost);
                // utility::io::debug.out("Final Mode %d, Cost2 %f\n", root, cost2 - cost);
                // utility::io::debug.out("Final Mode %d, Cost3 %f\n", root, cost3 - cost);

                // Tally the cost and append it to the cost function list for each final cost
                cost_result.push_back(std::pair<float, int>(cost1, root));
                cost_result.push_back(std::pair<float, int>(cost2, root));
                cost_result.push_back(std::pair<float, int>(cost3, root));
            }

        private:
            int ch_itt;
            const int ch_max;
            std::vector<std::pair<float, int>> cost_result;
            Eigen::Matrix<float, 4, 4> state_weight;
        };

        extern Optimizer optimizer1;
    }  // namespace AdaptiveMPC
}  // namespace MPC
}  // namespace module

#endif  // MODULE_OPTIMIZER_HPP


#include <ql/quantlib.hpp>
//#include <ql/userconfig.hpp>
#include <iostream>
#include <vector>
#include <boost/foreach.hpp>

using namespace std;
using namespace QuantLib;

void testingBlackScholesCalculator() {

    Real S0  = 100.0;
    Real K   = 105.0;
    Real rd  = 0.034;
    Real rf  = 0.021;
    Real tau = 0.5;
    Real vol = 0.177;

    Real domDisc = std::exp(-rd*tau);
    Real forDisc = std::exp(-rf*tau);
    Real stdDev = vol*std::sqrt(tau);

    boost::shared_ptr<PlainVanillaPayoff> vanillaPayoffPut(new PlainVanillaPayoff(Option::Put, K));
    boost::shared_ptr<AssetOrNothingPayoff> aonPayoffCall(new AssetOrNothingPayoff(Option::Call, K));

    BlackScholesCalculator vanillaPutPricer(vanillaPayoffPut, S0, forDisc, stdDev, domDisc);
    BlackScholesCalculator aonCallPricer(aonPayoffCall, S0, forDisc, stdDev, domDisc);

    cout << "------------- Vanilla VAlues -----------------" << endl;
    cout << "Value: " << vanillaPutPricer.value() << endl;

}

int main(void) {

    testingBlackScholesCalculator();
    return 0;
}


#ifndef CALIBRATOR_MODELS_SHORTRATE_TWOFACTORMODELS_GENERALG2_HPP
#define CALIBRATOR_MODELS_SHORTRATE_TWOFACTORMODELS_GENERALG2_HPP

#include <boost/make_shared.hpp>

#include <ql/models/shortrate/twofactormodel.hpp>
#include <ql/instruments/swaption.hpp>
#include <ql/math/solver1d.hpp>
#include <ql/math/integrals/kronrodintegral.hpp>

#include <calibrator/global.hpp>
#include <calibrator/processes/generalornsteinuhlenbeckprocess.hpp>
#include <calibrator/models/shortrate/dynamics/gaussianfactordynamics.hpp>

namespace HJCALIBRATOR
{
	//! Two-additive-factor gaussian model class
	/*! This class implements a two-additive-factor model defined by
	\f[
	dr_t = \varphi(t) + x_t + y_t
	\f]
	where \f$ x_t \f$ and \f$ y_t \f$ are defined by
	\f[
	dx_t = -a x_t dt + \sigma dW^1_t, x_0 = 0
	\f]
	\f[
	dy_t = -b y_t dt + \sigma dW^2_t, y_0 = 0
	\f]
	and \f$ dW^1_t dW^2_t = \rho dt \f$.

	\bug This class was not tested enough to guarantee
	its functionality.

	\todo Tree implementation

	\ingroup shortrate
	*/
	class GeneralizedG2 : public TwoFactorModel, public AffineModel, public TermStructureConsistentModel
	{
		class Dynamics;

		shared_ptr<Gaussian2FactorDynamics> dynamics_;

	public:
		GeneralizedG2( shared_ptr<Gaussian2FactorDynamics> dynamics,
					   Real integralSignificance = 10,
					   shared_ptr<Integrator> integrator = boost::make_shared<GaussKronrodAdaptive>( GaussKronrodAdaptive( 1.e-8, 10000 ) ) );
		virtual ~GeneralizedG2() {}

		// TwoFactorModel virtual override
		shared_ptr<Lattice> tree( const TimeGrid& grid ) const override
		{
			// todo
			return shared_ptr<Lattice>();
		};

		shared_ptr<ShortRateDynamics> dynamics() const;

		virtual DiscountFactor discount( Time t ) const override
		{
			return termStructure()->discount( t );
		}

		virtual Real discountBond( Time now,
								   Time maturity,
								   Array factors ) const;

		virtual Real discountBondOption( Option::Type type,
										 Real strike,
										 Time maturity,
										 Time bondMaturity ) const override;

		virtual Real discountBondOption( Option::Type type, Real strike,
										 Time maturity, Time bondStart,
										 Time bondMaturity ) const override;

		virtual Real swaption( const Swaption::arguments& arg, Real strike ) const;

		Parameter a() const { return a_; }
		Parameter b() const { return b_; }
		Parameter sigma() const { return sigma_; }
		Parameter eta() const { return eta_; }
		Parameter rho() const { return rho_; }

	protected:
		// CalibratedModel virtual override
		virtual void generateArguments() override;

		Real A( Time t, Time T ) const;

		Parameter& a_;
		Parameter& sigma_;
		Parameter& b_;
		Parameter& eta_;
		Parameter& rho_;

		Real integralSignificance_;

		shared_ptr<Integrator> integrator_;
	};

	//! Short-rate dynamics in the time-dependent Hull-White model
	/*! The short-rate follows an time-dependent Hull-White process */
	class GeneralizedG2::Dynamics : public TwoFactorModel::ShortRateDynamics {
	public:
		Dynamics( const shared_ptr<Gaussian2FactorDynamics> dynamics )
			: ShortRateDynamics( shared_ptr<StochasticProcess1D>( new GeneralizedOrnsteinUhlenbeckProcess( dynamics->a( 0 ), dynamics->sigma( 0 ) ) ),
								 shared_ptr<StochasticProcess1D>( new GeneralizedOrnsteinUhlenbeckProcess( dynamics->a( 1 ), dynamics->sigma( 1 ) ) ),
								 dynamics->rho( 0, 1 )(0.0) )
			, dynamics_( dynamics )
		{}

		virtual Real shortRate( Time t, Real x, Real y ) const {
			return x + y + dynamics_->phi( t );
		}

		shared_ptr<Gaussian2FactorDynamics> dynamics_;
	};

	// inline definitions
	inline shared_ptr<TwoFactorModel::ShortRateDynamics>	GeneralizedG2::dynamics() const
	{
		return shared_ptr<ShortRateDynamics>( new Dynamics( dynamics_ ) );
	}
}
#endif // !CALIBRATOR_MODELS_SHORTRATE_TWOFACTORMODELS_GENERALG2_HPP


//
// Created by Robert Schmoltzi on 19.05.21.
//

#include "SpiAlgo.h"
#include <boost/math/special_functions/factorials.hpp>

static GeneralizedRfAlgo::Scalar aSmallerB(const size_t a, const size_t b, const size_t n) {
	return GeneralizedRfAlgo::factorials.lg_unrooted_dbl_fact_fast(static_cast<long>(n)) -
	       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(b)) -
	       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(n - a)) +
	       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(b - a + 1));
}

// Input: |A1|, |A2|, |X|
// Called if splits overlap
// Output:
// if(S1 == S2):
//      - lg_root(|A1|) - lg_root(|B1|) + lg_unroot(|X|)      [ <=> h(S1)]
// else: // Let S1 be the greater split
//      - lg_root(|A1|) - lg_root(|B2|) + lg_root(|A1| - |A2| + 1) + lg_unroot(|X|)  [ <=> h(S1,S2)]
GeneralizedRfAlgo::Scalar SpiAlgo::one_overlap(const size_t a, const size_t b, const size_t n) {
	if (a == b) {
		return GeneralizedRfAlgo::factorials.lg_unrooted_dbl_fact_fast(static_cast<long>(n)) -
		       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(a)) -
		       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(n - a));
	}
	if (a < b) {
		return aSmallerB(a, b, n);
	}

	return GeneralizedRfAlgo::factorials.lg_unrooted_dbl_fact_fast(static_cast<long>(n)) -
	       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(a)) -
	       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(n - b)) +
	       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(a - b + 1));
}

// Input: |A1|, |B2|, |X|
// Called if splits overlap
// Output:
// if(S1 == S2):
//      - lg_root(|A2|) - lg_root(|B2|) + lg_unroot(|X|)      [ <=> h(S1)]
// else: // Let S1 be the greater split
//      - lg_root(|A1|) - lg_root(|B2|) + lg_root(|A1| - |A2| + 1) + lg_unroot(|X|)  [ <=> h(S1,S2)]
GeneralizedRfAlgo::Scalar
SpiAlgo::one_overlap_notb(const size_t a, const size_t n_minus_b, const size_t n) {
	const size_t b = n - n_minus_b;
	if (a == b) {
		return GeneralizedRfAlgo::factorials.lg_unrooted_dbl_fact_fast(static_cast<long>(n)) -
		       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(b)) -
		       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(n_minus_b));
	}
	if (a < b) {
		return aSmallerB(a, b, n);
	}
	return GeneralizedRfAlgo::factorials.lg_unrooted_dbl_fact_fast(static_cast<long>(n)) -
	       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(a)) -
	       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(n_minus_b)) +
	       GeneralizedRfAlgo::factorials.lg_rooted_dbl_fact_fast(static_cast<long>(a - b + 1));
}

// Based on the implementation by Martin R. Smith, found at
// https://github.com/ms609/TreeDist/blob/e05ab4c9e69c9548f066b2a7b256e35f7f92067d/src/tree_distance_functions.cpp
GeneralizedRfAlgo::Scalar SpiAlgo::calc_split_score(const PllSplit &S1, const PllSplit &S2) {
	{
		bool found = false;
		for (size_t i = 0; i != PllSplit::split_len; ++i) {
			// TODO: Time direct access vs vectorized Split-Methods (less iterations vs doing less
			// in each iteration)
			if ((~S1()[i] & S2()[i])) {
				found = true;
				break;
			}
		}
		if (!found)
			return one_overlap(S1.getPrecalcPopcnt(), S2.getPrecalcPopcnt(), taxa);
	}
	{
		bool found = false;

		for (size_t i = 0; i != PllSplit::split_len; ++i) {
			if ((S1()[i] & ~S2()[i])) {
				found = true;
				break;
			}
		}
		if (!found)
			return one_overlap(S1.getPrecalcPopcnt(), S2.getPrecalcPopcnt(), taxa);
	}
	const size_t bits_too_many = GeneralizedRfAlgo::bits_too_many(taxa);
	for (size_t i = 0; i != PllSplit::split_len; ++i) {
		pll_split_base_t test = ~(S1()[i] | S2()[i]);
		if (i == PllSplit::split_len - 1 && bits_too_many) {
			test &= (static_cast<pll_split_base_t>(~0) >> bits_too_many);
		}
		if (test) {
			return 0;
		}
	}
	return one_overlap_notb(S1.getPrecalcPopcnt(), S2.getPrecalcPopcnt(), taxa);
}
SpiAlgo::SpiAlgo(size_t split_len) : GeneralizedRfAlgo(split_len) {
}

RfAlgorithmInterface::Scalar SpiAlgo::calc_split_score(const PllSplit &S1) {
	return S1.getHInfoContent();
}

// AB | DEFC
// DEF | ABC --> ABC | DEF

// ABC | DE
// ADE | BC


#define DEBUG 1
/**
 * File    : F2.cpp
 * Author  : Kazune Takahashi
 * Created : 10/7/2020, 7:10:56 PM
 * Powered by Visual Studio Code
 */
#include <algorithm>
#include <bitset>
#include <cassert>
#include <cctype>
#include <chrono>
#include <climits>
#include <cmath>
#include <complex>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <tuple>
#include <unordered_map>
#include <unordered_set>
#include <vector>
// ----- boost -----
#include <boost/integer/common_factor_rt.hpp>
#include <boost/multiprecision/cpp_int.hpp>
#include <boost/rational.hpp>
// ----- AtCoder Library -----
#include <atcoder/all>
// ----- using directives and manipulations -----
using namespace std;
using boost::rational;
using boost::integer::gcd; // for C++14 or for cpp_int
using boost::integer::lcm; // for C++14 or for cpp_int
using boost::multiprecision::cpp_int;
using ll = long long;
using ld = long double;
template <typename T>
using max_heap = priority_queue<T>;
template <typename T>
using min_heap = priority_queue<T, vector<T>, greater<T>>;
// ----- ch_max and ch_min -----
template <typename T>
bool ch_max(T &left, T right)
{
  if (left < right)
  {
    left = right;
    return true;
  }
  return false;
}
template <typename T>
bool ch_min(T &left, T right)
{
  if (left > right)
  {
    left = right;
    return true;
  }
  return false;
}
// ----- mint -----
using mint = atcoder::modint1000000007;
// using mint = atcoder::modint998244353;
// using mint = atcoder::modint; // atcoder::modint::set_mod(xxx);
// using mint = atcoder::static_modint<1000000009>;
// using mint0 = dynamic_modint<xxx>;
// using mint1 = dynamic_modint<yyy>;
istream &operator>>(istream &is, mint &x)
{
  ll t;
  is >> t;
  x = t;
  return is;
}
ostream &operator<<(ostream &os, mint const &x)
{
  return os << x.val();
}
// ----- Combination -----
template <typename Mint = mint>
class Combination
{
public:
  constexpr static ll MAX_SIZE{3'000'010LL};
  // constexpr static ll MAX_SIZE{30'000'010LL}; // if 10^7 is needed
  vector<mint> inv, fact, factinv;
  Combination() : inv(MAX_SIZE), fact(MAX_SIZE), factinv(MAX_SIZE)
  {
    inv[1] = 1;
    for (auto i{2LL}; i < MAX_SIZE; i++)
    {
      inv[i] = (-inv[Mint::mod() % i]) * (Mint::mod() / i);
    }
    fact[0] = factinv[0] = 1;
    for (auto i{1LL}; i < MAX_SIZE; i++)
    {
      fact[i] = Mint(i) * fact[i - 1];
      factinv[i] = inv[i] * factinv[i - 1];
    }
  }
  Mint operator()(int n, int k)
  {
    if (n >= 0 && k >= 0 && n - k >= 0)
    {
      return fact[n] * factinv[k] * factinv[n - k];
    }
    return 0;
  }
  Mint catalan(int x, int y)
  {
    return (*this)(x + y, y) - (*this)(x + y, y - 1);
  }
};
// ----- for C++14 -----
using combination = Combination<mint>;
// ----- for C++17 -----
template <typename T, typename enable_if<is_integral<T>::value>::type * = nullptr>
size_t popcount(T x) { return bitset<64>(x).count(); }
size_t popcount(string const &S) { return bitset<200010>{S}.count(); }
// ----- Infty -----
template <typename T>
constexpr T Infty() { return numeric_limits<T>::max(); }
template <typename T>
constexpr T mInfty() { return numeric_limits<T>::min(); }
// ----- FenwickTree -----
template <typename T = ll>
class FenwickTree : public atcoder::fenwick_tree<T>
{
public:
  using atcoder::fenwick_tree<T>::fenwick_tree;

  FenwickTree(vector<T> const &v) : FenwickTree(static_cast<int>(v.size()))
  {
    for (auto i{size_t{0}}; i < v.size(); ++i)
    {
      atcoder::fenwick_tree<T>::add(i, v[i]);
    }
  }

  T operator[](int i)
  {
    return atcoder::fenwick_tree<T>::sum(i, i + 1);
  }
};
// ----- Math -----
using atcoder::floor_sum;
// ----- UnionFind -----
using UnionFind = atcoder::dsu;
// ----- MaxFlow -----
template <typename Cap = ll>
class MaxFlow : public atcoder::mf_graph<Cap>
{
public:
  using atcoder::mf_graph<Cap>::mf_graph;

  int add_edge(int from, int to)
  {
    return atcoder::mf_graph<Cap>::add_edge(from, to, Cap{1});
  }
};
ostream &operator<<(ostream &os, typename atcoder::mf_graph<ll>::edge const &edge)
{
  return os << "from: " << edge.from << ", to: " << edge.to << ", cap: " << edge.cap << ", flow: " << edge.flow;
}
// ----- MinCostFlow -----
template <typename Cap = ll, typename Cost = ll>
class MinCostFlow : public atcoder::mcf_graph<Cap, Cost>
{
private:
  Cost infty;

public:
  using atcoder::mcf_graph<Cap, Cost>::mcf_graph;

  MinCostFlow(int n, Cost infty = Cost{0}) : atcoder::mcf_graph<Cap, Cost>::mcf_graph(n), infty{infty} {}

  int add_edge(int from, int to, Cap cap)
  {
    return atcoder::mcf_graph<Cap, Cost>::add_edge(from, to, cap, Cost{0});
  }

  int add_edge(int from, int to, Cap cap, Cost cost)
  {
    return atcoder::mcf_graph<Cap, Cost>::add_edge(from, to, cap, cost + infty);
  }

  pair<Cap, Cost> flow(int s, int t)
  {
    return flow(s, t, std::numeric_limits<Cap>::max());
  }

  pair<Cap, Cost> flow(int s, int t, Cap flow_limit)
  {
    return slope(s, t, flow_limit).back();
  }

  vector<pair<Cap, Cost>> slope(int s, int t)
  {
    return slope(s, t, numeric_limits<Cap>::max());
  }

  vector<pair<Cap, Cost>> slope(int s, int t, Cap flow_limit)
  {
    auto res{atcoder::mcf_graph<Cap, Cost>::slope(s, t, flow_limit)};
    for (auto &[cap, cost] : res)
    {
      cost -= cap * infty;
    }
    return res;
  }

  typename atcoder::mcf_graph<Cap, Cost>::edge get_edge(int i)
  {
    auto res{atcoder::mcf_graph<Cap, Cost>::get_edge(i)};
    res.cost -= infty;
    return res;
  }

  vector<typename atcoder::mcf_graph<Cap, Cost>::edge> edges()
  {
    auto res{atcoder::mcf_graph<Cap, Cost>::edges()};
    for (auto &e : res)
    {
      e.cost -= infty;
    }
    return res;
  }
};
ostream &operator<<(ostream &os, typename atcoder::mcf_graph<ll, ll>::edge const &edge)
{
  return os << "from: " << edge.from << ", to: " << edge.to << ", cap: " << edge.cap << ", flow: " << edge.flow << ", cost: " << edge.cost;
}
// ----- frequently used constexpr -----
// constexpr double epsilon{1e-10};
// constexpr ll infty{1'000'000'000'000'010LL}; // or
// constexpr int infty{1'000'000'010};
// constexpr int dx[4] = {1, 0, -1, 0};
// constexpr int dy[4] = {0, 1, 0, -1};
// ----- Yes() and No() -----
void Yes()
{
  cout << "Yes" << endl;
  exit(0);
}
void No()
{
  cout << "No" << endl;
  exit(0);
}

// ----- Solve -----

class Solve
{

public:
  Solve()
  {
  }

  void flush()
  {
  }

private:
};

// ----- main() -----

/*
int main()
{
  Solve solve;
  solve.flush();
}
*/

combination C;

int main()
{
  ll k;
  string s;
  cin >> k >> s;
  mint ans{0};
  ll y{static_cast<ll>(s.size())};
  for (auto x{0}; x <= k; ++x)
  {
    ans += C(y + k - x - 1, y - 1) * mint{26}.pow(x) * mint{25}.pow(k - x);
  }
  cout << ans << endl;
}


#ifndef LEPP3_OBSTACLE_EVALUATOR_H_
#define LEPP3_OBSTACLE_EVALUATOR_H_
#include <boost/filesystem.hpp>
#include <iostream>
#include <sstream>
#include "lepp3/FrameData.hpp"
#include "lepp3/util/util.h"
/**
 * A class that computes the volume of a given model.
 * It is a `ModelVisitor` implementation and it assumes one model as input,
 * regardless of how many sub-parts the approximation has, a.k.a the
 * `CompositeModel`.
 *
 * The volume computation is done by creating a 3D grid around the model and
 * estimating how many points on the grid are occupied.
 */
class VolumeEstimator : public ModelVisitor {
public:
  VolumeEstimator()
      : num_splits_(0) {
    min_p_.x = std::numeric_limits<int>::max();
    min_p_.y = std::numeric_limits<int>::max();
    min_p_.z = std::numeric_limits<int>::max();
    max_p_.x = std::numeric_limits<int>::min();
    max_p_.y = std::numeric_limits<int>::min();
    max_p_.z = std::numeric_limits<int>::min();
  }
  /**
   * Implementation of the `ModelVisitor` interface. It will draw the given
   * sphere onto the `PCLVisualizer` to which it holds a reference.
   */
  void visitSphere(lepp::SphereModel& sphere);
  void visitCapsule(lepp::CapsuleModel& capsule);
  int getSplitCount() { return num_splits_; }
  /**
   * Estimates the volume of the approximated model. This is achieved by
   * creating a bounding box around the min/max points of the approximation and
   * subdivide the region into small grids.
   */
  int estimateVolume();
private:
  /**
   * Checks if the current point on the 3D grid (a.k.a the bounding box around
   * the model) is inside the model.
   */
  bool isVoxelOccupied(double const& x, double const& y, double const& z);
  /**
   * Minimum/Maximum points of the bounding box surrounding the model.
   */
  Coordinate min_p_, max_p_;
  /**
   * Container to hold track of sphere models in a `CompositeModel`
   */
  std::vector<SphereModel> spheres_;
  /**
   * Container to hold track of capsule models in a `CompositeModel`
   */
  std::vector<CapsuleModel> capsules_;
  /**
   * Number of sub-models in this model. This determines how many split operations
   * have been executed
   */
  int num_splits_;
};
void VolumeEstimator::visitSphere(lepp::SphereModel& sphere) {
  ++num_splits_;
  // Store the current Sphere model.
  spheres_.push_back(sphere);
  Coordinate const center = sphere.center_point();
  double const max_x = center.x + sphere.radius();
  double const max_y = center.y + sphere.radius();
  double const max_z = center.z + sphere.radius();
  double const min_x = center.x - sphere.radius();
  double const min_y = center.y - sphere.radius();
  double const min_z = center.z - sphere.radius();
  max_p_.x = std::max(max_p_.x, max_x);
  max_p_.y = std::max(max_p_.y, max_y);
  max_p_.z = std::max(max_p_.z, max_z);
  min_p_.x = std::min(min_p_.x, min_x);
  min_p_.y = std::min(min_p_.y, min_y);
  min_p_.z = std::min(min_p_.z, min_z);
  double volume = 4/3*M_PI*sphere.radius()*sphere.radius()*sphere.radius();
  //std::cout << volume << std::endl;
}
void VolumeEstimator::visitCapsule(lepp::CapsuleModel& capsule) {
  ++num_splits_;
  // Store the current Capsule model.
  capsules_.push_back(capsule);
  Coordinate first = capsule.first();
  Coordinate second = capsule.second();
  double a = std::sqrt((first.x - second.x)*(first.x - second.x) + (first.y - second.y)*(first.y - second.y) + (first.z - second.z)*(first.z - second.z));
  double volume = M_PI*capsule.radius()*capsule.radius()*(4/3*capsule.radius() + a);
  //std::cout << volume << std::endl;
  Coordinate min, max;
  min.x = std::min(first.x, second.x);
  min.y = std::min(first.y, second.y);
  min.z = std::min(first.z, second.z);
  max.x = std::max(first.x, second.x);
  max.y = std::max(first.y, second.y);
  max.z = std::max(first.z, second.z);
  // Add the capsule radius to min/max
  min.x -= capsule.radius();
  min.y -= capsule.radius();
  min.z -= capsule.radius();
  max.x += capsule.radius();
  max.y += capsule.radius();
  max.z += capsule.radius();
  // Find the global min/max
  min_p_.x = std::min(min_p_.x, min.x);
  min_p_.y = std::min(min_p_.y, min.y);
  min_p_.z = std::min(min_p_.z, min.z);
  max_p_.x = std::max(max_p_.x, max.x);
  max_p_.y = std::max(max_p_.y, max.y);
  max_p_.z = std::max(max_p_.z, max.z);
}
bool VolumeEstimator::isVoxelOccupied(
    double const& x, double const& y, double const& z) {
  // Go through all spheres of the current model and check if the point is
  // inside one of them.
  {
    size_t sz = spheres_.size();
    for (size_t i=0; i<sz; ++i) {
      double const r = spheres_[i].radius();
      Coordinate const& center = spheres_[i].center_point();
      double const distance =
          (center.x - x) * (center.x - x) + (center.y - y) * (center.y - y) + (center.z - z) * (center.z - z);
      if (distance <= r * r)
        return true;
    }
  }
  // Go through all capsules of the current model and check if the point is
  // inside one of them.
  {
    size_t sz = capsules_.size();
    for (size_t i=0; i<sz; ++i) {
      double r = capsules_[i].radius();
      /*
       * Prepare the pcl::sqrPointToLineDistance input arguments.
       */
      // line_dir: line direction vector
      Coordinate l = capsules_[i].second() - capsules_[i].first();
      Eigen::Vector4f line_dir(l.x, l.y, l.z, 0.);
      // line_pt: a point on the line
      Eigen::Vector4f line_pt(
          capsules_[i].second().x,
          capsules_[i].second().y,
          capsules_[i].second().z,
          0.);
      Eigen::Vector4f pt(x, y, z, 0);
      double const distance = pcl::sqrPointToLineDistance(pt, line_pt, line_dir);
      if (distance < r*r)
        return true;
    }
  }
  return false;
}
int VolumeEstimator::estimateVolume() {
  // NOTE: All the values are in METERS
  // Create a 3D grid and check the distance of each point on grid to the object
  // models.
  // TODO: 3D grid creation should depend on the point cloud resolution
  double const step_size = 0.01;
  int volume = 0;
  for (double x = min_p_.x; x < max_p_.x; x += step_size) {
    for (double y = min_p_.y; y < max_p_.y; y += step_size) {
      for (double z = min_p_.z; z < max_p_.z; z += step_size) {
        if (isVoxelOccupied(x, y, z))
          ++volume;
      }
    }
  }
  return volume;
}
/**
 *
 */
class ObstacleEvaluator : public FrameDataObserver {
public:
  ObstacleEvaluator();
  ObstacleEvaluator(int vol);
  /**
   * ObstacleAggregator interface implementation: processes the current models.
   */
  void updateFrame(FrameDataPtr frameData);
private:
  void init();
  bool evaluate(ObjectModelPtr const& model, double x, double y, double z);
  std::string file_path_;
  int ref_volume_;
};
ObstacleEvaluator::ObstacleEvaluator()
    : ref_volume_(0) {
  init();
}
ObstacleEvaluator::ObstacleEvaluator(int vol)
    : ref_volume_(vol) {
  init();
}
void ObstacleEvaluator::init() {
  namespace bfs = boost::filesystem;
  std::stringstream ss;
  // Create the evaluation directory
  ss << "../evaluation/";
  std::string dir = ss.str();
  if ( !bfs::exists(bfs::path(dir)) )
    bfs::create_directory(bfs::path(dir));
  // Create the sub directory based on current timestamp
  ss << lepp::get_current_timestamp();
  dir = ss.str();
  bfs::create_directory(bfs::path(dir));
  // Prepare the csv file path
  ss  << "/eval.csv";
  file_path_ = ss.str();
  std::cout << "file_path: " << file_path_ << std::endl;
  // Create the file header
  std::ofstream tf_fout;
  tf_fout.open(file_path_.c_str(), std::ofstream::app);
  tf_fout << "model_id,"
          << "volume,"
          << "sim_veloc_x,"
          << "sim_veloc_y,"
          << "sim_veloc_z,"
          << std::endl;
  tf_fout.close();
}
bool ObstacleEvaluator::evaluate(ObjectModelPtr const& model, double x, double y, double z) {
  // TODO incorporate try-catch scheme
  VolumeEstimator estimator;
  // Find the minimum and maximum point of the approximated model
  model->accept(estimator);
  // Compute the volume of the current model
  int vol = estimator.estimateVolume();
  //std::cout << "Volume Tester " << vol << std::endl;
  // Evaluate the approximation based on the reference information
  float ratio = -1;
  if (ref_volume_ != 0)
    ratio = vol / static_cast<float>(ref_volume_);

  bool has_vel = false;
  boost::shared_ptr<CompositeModel> cm = boost::static_pointer_cast<CompositeModel>(model);
  for( auto& m : cm->models() )
  {
    boost::shared_ptr<CompositeModel> cmm = boost::static_pointer_cast<CompositeModel>(m);

    for ( auto& mm : cmm->models() )
    {
      std::cout << model->id() << ": " << mm->velocity().x << ", "
                << mm->velocity().y << ", " << mm->velocity().z << std::endl;
      if (!std::isnan(mm->velocity().x) &&
          !std::isnan(mm->velocity().y) &&
          !std::isnan(mm->velocity().z))
      {

        // Save the current approximation information
        std::stringstream ss;
        ss << model->id() << ","
           << vol << ","
           << mm->velocity().x << ","
           << mm->velocity().y << ","
           << mm->velocity().z
           << std::endl;
        std::ofstream tf_fout;
        // open the file and add the current model evaluation to the end of it
        tf_fout.open(file_path_.c_str(), std::ofstream::app);
        tf_fout << ss.str();
        tf_fout.close();
        has_vel = true;
        break;
      }
    }
  }

  // if obstacle did not have a velocity estimate, log zero
  if (!has_vel)
  {
      std::ofstream tf_out;
      tf_out.open(file_path_.c_str(), std::ofstream::app);
      tf_out << model->id() << ","
             << vol << ","
             << "0,0,0"
             << std::endl;
  }

  return true;
}
void ObstacleEvaluator::updateFrame(FrameDataPtr frameData) {
  size_t sz = frameData->obstacles.size();
  for (size_t i=0; i<sz; ++i) {
    //evaluate(frameData->obstacles[i]);
    double x = frameData->obstacles[i]->velocity().x;
    double y = frameData->obstacles[i]->velocity().y;
    double z = frameData->obstacles[i]->velocity().z;
    evaluate(frameData->obstacles[i], x, y, z);
  }
}
#endif // LEPP3_OBSTACLE_EVALUATOR_H_


/**
 * \file      miscellaneous-algorithms.hpp
 * \author    Mehdi Benallegue
 * \date       2013
 * \brief      Gathers many kinds of algorithms
 *
 *
 *
 */

#ifndef STATEOBSERVATIONTOOLSMISCELANEOUSALGORITHMS
#define STATEOBSERVATIONTOOLSMISCELANEOUSALGORITHMS
#include <boost/utility.hpp>
#include <cmath>

#include <state-observation/api.h>
#include <state-observation/tools/definitions.hpp>

namespace stateObservation
{
namespace tools
{

/// computes the square of a value of any type
template<class T>
inline T square(const T & x)
{
  return T(x * x);
}

/// derivates any type with finite differences
template<class T>
inline T derivate(const T & o1, const T & o2, double dt)
{
  T o(o2 - o1);
  return o * (1 / dt);
}

/// gives the sign of a variable (1, 0 or -1)
template<typename T>
inline int signum(T x)
{
  return (T(0) < x) - (x < T(0));
}

template<typename T>
inline std::string toString(T val)
{
  std::stringstream ss("");
  ss << val;
  return ss.str();
}

/// @brief checks if the vector is already normalized or not
///
/// @param v                  the vector to normalize
/// @return true              the vector is normalized
/// @return false             The vector is not normalized
inline bool checkIfNormalized(const Vector3 & v)
{
  if(fabs(v.squaredNorm() - 1) > cst::epsilon1)
  {
    return false;
  }
  else
  {
    return true;
  }
}

/// @brief checks if the vector is already normalized or not
///
/// @param v                  the vector to normalize
/// @param outputSquaredNorm  the squared norm as an output
/// @return true              the vector is normalized
/// @return false             The vector is not normalized
inline bool checkIfNormalized(const Vector3 & v, double & outputSquaredNorm)
{
  outputSquaredNorm = v.squaredNorm();
  if(fabs(outputSquaredNorm - 1) > cst::epsilon1)
  {
    return false;
  }
  else
  {
    return true;
  }
}

/// @brief returns the value clamped between min and max
///
/// @tparam T The type of the scalar
/// @param x the input value
/// @param max the max value
/// @param min the min value
/// @return T the calmped value
template<typename T>
inline T clampScalar(const T & x, const T & max, const T & min)
{
  if(x > max)
  {
    return max;
  }
  else
  {
    if(x < min)
    {
      return min;
    }
    else
    {
      return x;
    }
  }
}

/// @brief returns the value clamped between limit and -limit
///
/// @tparam T The type of the scalar
/// @param x the input value
/// @param limit is the maximum absolute value allowed, has to be positive
/// @return T the clamped value
template<typename T>
inline T clampScalar(const T & x, const T & limit)
{
  return clampScalar(x, limit, -limit);
}

/// @brief normalize the vector only if it is not normalized already. Useful if the vector is likely to be normalized
///
/// @param v the input vector
/// @return Vector3
inline Vector3 normalizedLazy(const Vector3 & v)
{
  double squaredNorm;
  if(checkIfNormalized(v, squaredNorm))
  {
    return v;
  }
  else
  {
    return v / sqrt(squaredNorm);
  }
}

/// provides an acceleration giving a finite time convergence to zero
/// the state is the position x and the derivative xd and the output is the
/// acceleration. The gains kp, kv must be negative
inline double STATE_OBSERVATION_DLLAPI finiteTimeAccControl(double x, double xd, double kp = -1, double kv = -1)
{
  double sax = sqrt(fabs(x));
  double xdr = kp * signum(x) * sax;
  double y = xd - xdr;
  double ydr = -kv * signum(y) * sqrt(fabs(y));
  return ydr - kp * xd / (2 * sax);
}

/// sqme as the scalar version but for every member of the vector
inline Vector STATE_OBSERVATION_DLLAPI finiteTimeAccControl(const Vector & x,
                                                            const Vector & xd,
                                                            double kp = -1,
                                                            double kv = -1)
{
  Vector xdd(x.size());
  for(Index i = 1; i < x.size(); ++i)
  {
    xdd(i) = finiteTimeAccControl(x(i), xd(i), kp, kv);
  }
  return xdd;
}

namespace Detail
{
double constexpr sqrtNewtonRaphson(double x, double curr, double prev)
{
  return curr == prev ? curr : sqrtNewtonRaphson(x, 0.5 * (curr + x / curr), curr);
}
} // namespace Detail

/// @brief Constexpr version of the square root
/// @details For a finite and non-negative value of "x", returns an approximation for the square root of "x"
///-Otherwise, returns NaN
///
/// @param x
/// @return double constexpr
double constexpr sqrt(double x)
{
  return x >= 0 && x < std::numeric_limits<double>::infinity() ? Detail::sqrtNewtonRaphson(x, x, 0)
                                                               : std::numeric_limits<double>::quiet_NaN();
}

} // namespace tools

} // namespace stateObservation

#endif // STATEOBSERVATIONTOOLSMISCELANEOUSALGORITHMS


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

#include "open3d/utility/Eigen.h"

#include <Eigen/Geometry>
#include <Eigen/Sparse>

#include "open3d/utility/Console.h"

namespace open3d {
namespace utility {

/// Function to solve Ax=b
std::tuple<bool, Eigen::VectorXd> SolveLinearSystemPSD(
        const Eigen::MatrixXd &A,
        const Eigen::VectorXd &b,
        bool prefer_sparse /* = false */,
        bool check_symmetric /* = false */,
        bool check_det /* = false */,
        bool check_psd /* = false */) {
    // PSD implies symmetric
    check_symmetric = check_symmetric || check_psd;
    if (check_symmetric && !A.isApprox(A.transpose())) {
        LogWarning("check_symmetric failed, empty vector will be returned");
        return std::make_tuple(false, Eigen::VectorXd::Zero(b.rows()));
    }

    if (check_det) {
        double det = A.determinant();
        if (fabs(det) < 1e-6 || std::isnan(det) || std::isinf(det)) {
            LogWarning("check_det failed, empty vector will be returned");
            return std::make_tuple(false, Eigen::VectorXd::Zero(b.rows()));
        }
    }

    // Check PSD: https://stackoverflow.com/a/54569657/1255535
    if (check_psd) {
        Eigen::LLT<Eigen::MatrixXd> A_llt(A);
        if (A_llt.info() == Eigen::NumericalIssue) {
            LogWarning("check_psd failed, empty vector will be returned");
            return std::make_tuple(false, Eigen::VectorXd::Zero(b.rows()));
        }
    }

    Eigen::VectorXd x(b.size());

    if (prefer_sparse) {
        Eigen::SparseMatrix<double> A_sparse = A.sparseView();
        // TODO: avoid deprecated API SimplicialCholesky
        Eigen::SimplicialCholesky<Eigen::SparseMatrix<double>> A_chol;
        A_chol.compute(A_sparse);
        if (A_chol.info() == Eigen::Success) {
            x = A_chol.solve(b);
            if (A_chol.info() == Eigen::Success) {
                // Both decompose and solve are successful
                return std::make_tuple(true, std::move(x));
            } else {
                LogWarning("Cholesky solve failed, switched to dense solver");
            }
        } else {
            LogWarning("Cholesky decompose failed, switched to dense solver");
        }
    }

    x = A.ldlt().solve(b);
    return std::make_tuple(true, std::move(x));
}

Eigen::Matrix4d TransformVector6dToMatrix4d(const Eigen::Vector6d &input) {
    Eigen::Matrix4d output;
    output.setIdentity();
    output.block<3, 3>(0, 0) =
            (Eigen::AngleAxisd(input(2), Eigen::Vector3d::UnitZ()) *
             Eigen::AngleAxisd(input(1), Eigen::Vector3d::UnitY()) *
             Eigen::AngleAxisd(input(0), Eigen::Vector3d::UnitX()))
                    .matrix();
    output.block<3, 1>(0, 3) = input.block<3, 1>(3, 0);
    return output;
}

Eigen::Vector6d TransformMatrix4dToVector6d(const Eigen::Matrix4d &input) {
    Eigen::Vector6d output;
    Eigen::Matrix3d R = input.block<3, 3>(0, 0);
    double sy = sqrt(R(0, 0) * R(0, 0) + R(1, 0) * R(1, 0));
    if (!(sy < 1e-6)) {
        output(0) = atan2(R(2, 1), R(2, 2));
        output(1) = atan2(-R(2, 0), sy);
        output(2) = atan2(R(1, 0), R(0, 0));
    } else {
        output(0) = atan2(-R(1, 2), R(1, 1));
        output(1) = atan2(-R(2, 0), sy);
        output(2) = 0;
    }
    output.block<3, 1>(3, 0) = input.block<3, 1>(0, 3);
    return output;
}

std::tuple<bool, Eigen::Matrix4d> SolveJacobianSystemAndObtainExtrinsicMatrix(
        const Eigen::Matrix6d &JTJ, const Eigen::Vector6d &JTr) {
    bool solution_exist;
    Eigen::Vector6d x;
    std::tie(solution_exist, x) = SolveLinearSystemPSD(JTJ, -JTr);

    if (solution_exist) {
        Eigen::Matrix4d extrinsic = TransformVector6dToMatrix4d(x);
        return std::make_tuple(solution_exist, std::move(extrinsic));
    }
    return std::make_tuple(false, Eigen::Matrix4d::Identity());
}

std::tuple<bool, std::vector<Eigen::Matrix4d, Matrix4d_allocator>>
SolveJacobianSystemAndObtainExtrinsicMatrixArray(const Eigen::MatrixXd &JTJ,
                                                 const Eigen::VectorXd &JTr) {
    std::vector<Eigen::Matrix4d, Matrix4d_allocator> output_matrix_array;
    output_matrix_array.clear();
    if (JTJ.rows() != JTr.rows() || JTJ.cols() % 6 != 0) {
        LogWarning(
                "[SolveJacobianSystemAndObtainExtrinsicMatrixArray] "
                "Unsupported matrix format.");
        return std::make_tuple(false, std::move(output_matrix_array));
    }

    bool solution_exist;
    Eigen::VectorXd x;
    std::tie(solution_exist, x) = SolveLinearSystemPSD(JTJ, -JTr);

    if (solution_exist) {
        int nposes = (int)x.rows() / 6;
        for (int i = 0; i < nposes; i++) {
            Eigen::Matrix4d extrinsic =
                    TransformVector6dToMatrix4d(x.block<6, 1>(i * 6, 0));
            output_matrix_array.push_back(extrinsic);
        }
        return std::make_tuple(solution_exist, std::move(output_matrix_array));
    } else {
        return std::make_tuple(false, std::move(output_matrix_array));
    }
}

template <typename MatType, typename VecType>
std::tuple<MatType, VecType, double> ComputeJTJandJTr(
        std::function<void(int, VecType &, double &, double &)> f,
        int iteration_num,
        bool verbose /*=true*/) {
    MatType JTJ;
    VecType JTr;
    double r2_sum = 0.0;
    JTJ.setZero();
    JTr.setZero();
#pragma omp parallel
    {
        MatType JTJ_private;
        VecType JTr_private;
        double r2_sum_private = 0.0;
        JTJ_private.setZero();
        JTr_private.setZero();
        VecType J_r;
        double r;
        double w = 0.0;
#pragma omp for nowait
        for (int i = 0; i < iteration_num; i++) {
            f(i, J_r, r, w);
            JTJ_private.noalias() += J_r * w * J_r.transpose();
            JTr_private.noalias() += J_r * w * r;
            r2_sum_private += r * r;
        }
#pragma omp critical
        {
            JTJ += JTJ_private;
            JTr += JTr_private;
            r2_sum += r2_sum_private;
        }
    }
    if (verbose) {
        LogDebug("Residual : {:.2e} (# of elements : {:d})",
                 r2_sum / (double)iteration_num, iteration_num);
    }
    return std::make_tuple(std::move(JTJ), std::move(JTr), r2_sum);
}

template <typename MatType, typename VecType>
std::tuple<MatType, VecType, double> ComputeJTJandJTr(
        std::function<
                void(int,
                     std::vector<VecType, Eigen::aligned_allocator<VecType>> &,
                     std::vector<double> &)> f,
        int iteration_num,
        bool verbose /*=true*/) {
    MatType JTJ;
    VecType JTr;
    double r2_sum = 0.0;
    JTJ.setZero();
    JTr.setZero();
#pragma omp parallel
    {
        MatType JTJ_private;
        VecType JTr_private;
        double r2_sum_private = 0.0;
        JTJ_private.setZero();
        JTr_private.setZero();
        std::vector<double> r;
        std::vector<VecType, Eigen::aligned_allocator<VecType>> J_r;
#pragma omp for nowait
        for (int i = 0; i < iteration_num; i++) {
            f(i, J_r, r);
            for (int j = 0; j < (int)r.size(); j++) {
                JTJ_private.noalias() += J_r[j] * J_r[j].transpose();
                JTr_private.noalias() += J_r[j] * r[j];
                r2_sum_private += r[j] * r[j];
            }
        }
#pragma omp critical
        {
            JTJ += JTJ_private;
            JTr += JTr_private;
            r2_sum += r2_sum_private;
        }
    }
    if (verbose) {
        LogDebug("Residual : {:.2e} (# of elements : {:d})",
                 r2_sum / (double)iteration_num, iteration_num);
    }
    return std::make_tuple(std::move(JTJ), std::move(JTr), r2_sum);
}

// clang-format off
template std::tuple<Eigen::Matrix6d, Eigen::Vector6d, double> ComputeJTJandJTr(
        std::function<void(int, Eigen::Vector6d &, double &, double &)> f,
        int iteration_num, bool verbose);

template std::tuple<Eigen::Matrix6d, Eigen::Vector6d, double> ComputeJTJandJTr(
        std::function<void(int,
                           std::vector<Eigen::Vector6d, Vector6d_allocator> &,
                           std::vector<double> &)> f,
        int iteration_num, bool verbose);
// clang-format on

Eigen::Matrix3d RotationMatrixX(double radians) {
    Eigen::Matrix3d rot;
    rot << 1, 0, 0, 0, std::cos(radians), -std::sin(radians), 0,
            std::sin(radians), std::cos(radians);
    return rot;
}

Eigen::Matrix3d RotationMatrixY(double radians) {
    Eigen::Matrix3d rot;
    rot << std::cos(radians), 0, std::sin(radians), 0, 1, 0, -std::sin(radians),
            0, std::cos(radians);
    return rot;
}

Eigen::Matrix3d RotationMatrixZ(double radians) {
    Eigen::Matrix3d rot;
    rot << std::cos(radians), -std::sin(radians), 0, std::sin(radians),
            std::cos(radians), 0, 0, 0, 1;
    return rot;
}

Eigen::Vector3uint8 ColorToUint8(const Eigen::Vector3d &color) {
    Eigen::Vector3uint8 rgb;
    for (int i = 0; i < 3; ++i) {
        rgb[i] = uint8_t(
                std::round(std::min(1., std::max(0., color(i))) * 255.));
    }
    return rgb;
}

Eigen::Vector3d ColorToDouble(uint8_t r, uint8_t g, uint8_t b) {
    return Eigen::Vector3d(r, g, b) / 255.0;
}

Eigen::Vector3d ColorToDouble(const Eigen::Vector3uint8 &rgb) {
    return ColorToDouble(rgb(0), rgb(1), rgb(2));
}

template <typename IdxType>
Eigen::Matrix3d ComputeCovariance(const std::vector<Eigen::Vector3d> &points,
                                  const std::vector<IdxType> &indices) {
    Eigen::Matrix3d covariance;
    Eigen::Matrix<double, 9, 1> cumulants;
    cumulants.setZero();
    for (const auto &idx : indices) {
        const Eigen::Vector3d &point = points[idx];
        cumulants(0) += point(0);
        cumulants(1) += point(1);
        cumulants(2) += point(2);
        cumulants(3) += point(0) * point(0);
        cumulants(4) += point(0) * point(1);
        cumulants(5) += point(0) * point(2);
        cumulants(6) += point(1) * point(1);
        cumulants(7) += point(1) * point(2);
        cumulants(8) += point(2) * point(2);
    }
    cumulants /= (double)indices.size();
    covariance(0, 0) = cumulants(3) - cumulants(0) * cumulants(0);
    covariance(1, 1) = cumulants(6) - cumulants(1) * cumulants(1);
    covariance(2, 2) = cumulants(8) - cumulants(2) * cumulants(2);
    covariance(0, 1) = cumulants(4) - cumulants(0) * cumulants(1);
    covariance(1, 0) = covariance(0, 1);
    covariance(0, 2) = cumulants(5) - cumulants(0) * cumulants(2);
    covariance(2, 0) = covariance(0, 2);
    covariance(1, 2) = cumulants(7) - cumulants(1) * cumulants(2);
    covariance(2, 1) = covariance(1, 2);
    return covariance;
}

template <typename IdxType>
std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
        const std::vector<Eigen::Vector3d> &points,
        const std::vector<IdxType> &indices) {
    Eigen::Vector3d mean;
    Eigen::Matrix3d covariance;
    Eigen::Matrix<double, 9, 1> cumulants;
    cumulants.setZero();
    for (const auto &idx : indices) {
        const Eigen::Vector3d &point = points[idx];
        cumulants(0) += point(0);
        cumulants(1) += point(1);
        cumulants(2) += point(2);
        cumulants(3) += point(0) * point(0);
        cumulants(4) += point(0) * point(1);
        cumulants(5) += point(0) * point(2);
        cumulants(6) += point(1) * point(1);
        cumulants(7) += point(1) * point(2);
        cumulants(8) += point(2) * point(2);
    }
    cumulants /= (double)indices.size();
    mean(0) = cumulants(0);
    mean(1) = cumulants(1);
    mean(2) = cumulants(2);
    covariance(0, 0) = cumulants(3) - cumulants(0) * cumulants(0);
    covariance(1, 1) = cumulants(6) - cumulants(1) * cumulants(1);
    covariance(2, 2) = cumulants(8) - cumulants(2) * cumulants(2);
    covariance(0, 1) = cumulants(4) - cumulants(0) * cumulants(1);
    covariance(1, 0) = covariance(0, 1);
    covariance(0, 2) = cumulants(5) - cumulants(0) * cumulants(2);
    covariance(2, 0) = covariance(0, 2);
    covariance(1, 2) = cumulants(7) - cumulants(1) * cumulants(2);
    covariance(2, 1) = covariance(1, 2);
    return std::make_tuple(mean, covariance);
}

template Eigen::Matrix3d ComputeCovariance(
        const std::vector<Eigen::Vector3d> &points,
        const std::vector<size_t> &indices);
template std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
        const std::vector<Eigen::Vector3d> &points,
        const std::vector<size_t> &indices);
template Eigen::Matrix3d ComputeCovariance(
        const std::vector<Eigen::Vector3d> &points,
        const std::vector<int> &indices);
template std::tuple<Eigen::Vector3d, Eigen::Matrix3d> ComputeMeanAndCovariance(
        const std::vector<Eigen::Vector3d> &points,
        const std::vector<int> &indices);
}  // namespace utility
}  // namespace open3d