/**
 Svi is written by Yue Ariel Wu
 UCLA
*/
#include <fstream>
#include <iostream>
#include <string>
#include <stdlib.h>
#include <vector>
#include <iomanip>

//#include <random>

#include <bits/stdc++.h>
#include <Eigen/Dense>
#include <Eigen/Core>
#include <Eigen/LU>
#include <Eigen/SVD>
#include <Eigen/QR>
#include "time.h"

#include <math.h>
#include "boost/numeric/ublas/matrix.hpp"
#include "boost/numeric/ublas/io.hpp"
#include "boost/random.hpp"
#include "boost/numeric/ublas/lu.hpp"
#include "boost/numeric/ublas/vector.hpp"
#include "boost/numeric/ublas/vector_proxy.hpp"
#include "boost/numeric/ublas/triangular.hpp"

#include "genotype.h"
#include "mailman.h"
#include "arguments.h"
#include "helper.h"
#include "storage.h"

#if SSE_SUPPORT==1
	#define fastmultiply fastmultiply_sse
	#define fastmultiply_pre fastmultiply_pre_sse
#else
	#define fastmultiply fastmultiply_normal
	#define fastmultiply_pre fastmultiply_pre_normal
#endif
namespace ukblas =boost::numeric::ublas;
using namespace Eigen;
using namespace std;

// Storing in RowMajor Form
typedef Matrix<double, Dynamic, Dynamic, RowMajor> MatrixXdr;

//Intermediate Variables
int blocksize;
double *partialsums;
double *sum_op;
double *yint_e;
double *yint_m;
double **y_e;
double **y_m;
//
//
struct timespec t0;

MatrixXdr covariate;
MatrixXdr pheno;
genotype g;
MatrixXdr geno_matrix; //(p,n)
int MAX_ITER;
int k,p,n;
int k_orig;

MatrixXdr c;
MatrixXdr x;
MatrixXdr v;
MatrixXdr means; //(p,1)
MatrixXdr stds; //(p,1)
MatrixXdr sum2;
MatrixXdr sum;

options command_line_opts;

bool debug = false;
bool check_accuracy = false;
bool var_normalize=true;
int accelerated_em=0;
double convergence_limit;
bool memory_efficient = false;
bool missing=false;
bool fast_mode = true;
bool text_version = false;
bool use_cov=false;
bool reg = true;
bool gwas=false;

vector<string> pheno_name;

std::istream& newline(std::istream& in)
{
    if ((in >> std::ws).peek() != std::char_traits<char>::to_int_type('\n')) {
        in.setstate(std::ios_base::failbit);
    }
    return in.ignore();
}
int read_cov(bool std,int Nind, std::string filename, std::string covname){
	ifstream ifs(filename.c_str(), ios::in);
	std::string line;
	std::istringstream in;
	int covIndex = 0;
	std::getline(ifs,line);
	in.str(line);
	string b;
	vector<vector<int> > missing;
	int covNum=0;
	while(in>>b)
	{
		if(b!="FID" && b !="IID"){
		missing.push_back(vector<int>()); //push an empty row
		if(b==covname && covname!="")
			covIndex=covNum;
		covNum++;
		}
	}
	vector<double> cov_sum(covNum, 0);
	if(covname=="")
	{
		covariate.resize(Nind, covNum);
		cout<< "Read in "<<covNum << " Covariates.. "<<endl;
	}
	else
	{
		covariate.resize(Nind, 1);
		cout<< "Read in covariate "<<covname<<endl;
	}

	int j=0;
	while(std::getline(ifs, line)){
		in.clear();
		in.str(line);
		string temp;
		in>>temp; in>>temp; //FID IID
		for(int k=0; k<covNum; k++){

			in>>temp;
			if(temp=="NA")
			{
				missing[k].push_back(j);
				continue;
			}
			double cur = atof(temp.c_str());
			if(cur==-9)
			{
				missing[k].push_back(j);
				continue;
			}
			if(covname=="")
			{
				cov_sum[k]= cov_sum[k]+ cur;
				covariate(j,k) = cur;
			}
			else
				if(k==covIndex)
				{
					covariate(j, 0) = cur;
					cov_sum[k] = cov_sum[k]+cur;
				}
		}
		j++;
	}
	//compute cov mean and impute
	for (int a=0; a<covNum ; a++)
	{
		int missing_num = missing[a].size();
		cov_sum[a] = cov_sum[a] / (Nind - missing_num);

		for(int b=0; b<missing_num; b++)
		{
                        int index = missing[a][b];
                        if(covname=="")
                                covariate(index, a) = cov_sum[a];
                        else if (a==covIndex)
                                covariate(index, 0) = cov_sum[a];
                }
	}
	if(std)
	{
		MatrixXdr cov_std;
		cov_std.resize(1,covNum);
		MatrixXdr sum = covariate.colwise().sum();
		MatrixXdr sum2 = (covariate.cwiseProduct(covariate)).colwise().sum();
		MatrixXdr temp;
//		temp.resize(Nind, 1);
//		for(int i=0; i<Nind; i++)
//			temp(i,0)=1;
		for(int b=0; b<covNum; b++)
		{
			cov_std(0,b) = sum2(0,b) + Nind*cov_sum[b]*cov_sum[b]- 2*cov_sum[b]*sum(0,b);
			cov_std(0,b) =sqrt((Nind-1)/cov_std(0,b)) ;
			double scalar=cov_std(0,b);
			for(int j=0; j<Nind; j++)
			{
				covariate(j,b) = covariate(j,b)-cov_sum[b];
				covariate(j,b) =covariate(j,b)*scalar;
			}
			//covariate.col(b) = covariate.col(b) -temp*cov_sum[b];

		}
	}
	return covNum;
}
int read_pheno2(int Nind, std::string filename,int pheno_idx){
//	pheno.resize(Nind,1);
	ifstream ifs(filename.c_str(), ios::in);

	std::string line;
	std::istringstream in;
	int phenocount=0;
	vector<vector<int> > missing;
//read header
	std::getline(ifs,line);
	in.str(line);
	string b;
	while(in>>b)
	{
		if(b!="FID" && b !="IID"){
			phenocount++;
			missing.push_back(vector<int>());
			pheno_name.push_back(b);
		}
	}
	if(pheno_idx !=0)
		pheno_name[0] = pheno_name[pheno_idx-1];
	vector<double> pheno_sum(phenocount,0);
	if(pheno_idx !=0)
		pheno.resize(Nind,1);
	else
		pheno.resize(Nind, phenocount);
	int i=0;
	while(std::getline(ifs, line)){
		in.clear();
		in.str(line);
		string temp;
		//fid,iid
		//todo: fid iid mapping;
		//todo: handle missing phenotype
		in>>temp; in>>temp;
		for(int j=0; j<phenocount;j++) {
			in>>temp;
			if(temp=="NA")
			{
				missing[j].push_back(i);
				continue;
			}
			double cur= atof(temp.c_str());
			if(pheno_idx==0){
				pheno(i,j)=cur;
				pheno_sum[j] = pheno_sum[j]+cur;
			}
			else if(j == (pheno_idx-1))
			{
				pheno(i,0)=cur;
				pheno_sum[j]= pheno_sum[j]+cur;
			}
		}
		i++;
	}
	for(int a=0; a<phenocount; a++)
	{
		int missing_num= missing[a].size();
		double	pheno_avg = pheno_sum[a]/(Nind- missing_num);
		//cout<<"pheno "<<a<<" avg: "<<pheno_avg<<endl;
		for(int b=0 ; b<missing_num; b++)
		{
			int index = missing[a][b];
			pheno(index, a)= pheno_avg;
		}
		if(pheno_idx!=0)
			return 1;
	}
	return phenocount;
}

void multiply_y_pre_fast(MatrixXdr &op, int Ncol_op ,MatrixXdr &res,bool subtract_means){

	for(int k_iter=0;k_iter<Ncol_op;k_iter++){
		sum_op[k_iter]=op.col(k_iter).sum();
	}

			//cout << "Nops = " << Ncol_op << "\t" <<g.Nsegments_hori << endl;
	#if DEBUG==1
		if(debug){
			print_time ();
			cout <<"Starting mailman on premultiply"<<endl;
			cout << "Nops = " << Ncol_op << "\t" <<g.Nsegments_hori << endl;
			cout << "Segment size = " << g.segment_size_hori << endl;
			cout << "Matrix size = " <<g.segment_size_hori<<"\t" <<g.Nindv << endl;
			cout << "op = " <<  op.rows () << "\t" << op.cols () << endl;
		}
	#endif

	//TODO: Memory Effecient SSE FastMultipy

	for(int seg_iter=0;seg_iter<g.Nsegments_hori-1;seg_iter++){
		mailman::fastmultiply(g.segment_size_hori,g.Nindv,Ncol_op,g.p[seg_iter],op,yint_m,partialsums,y_m);
		int p_base = seg_iter*g.segment_size_hori;
		for(int p_iter=p_base; (p_iter<p_base+g.segment_size_hori) && (p_iter<g.Nsnp) ; p_iter++ ){
			for(int k_iter=0;k_iter<Ncol_op;k_iter++)
				res(p_iter,k_iter) = y_m[p_iter-p_base][k_iter];
		}
	}

	int last_seg_size = (g.Nsnp%g.segment_size_hori !=0 ) ? g.Nsnp%g.segment_size_hori : g.segment_size_hori;
	mailman::fastmultiply(last_seg_size,g.Nindv,Ncol_op,g.p[g.Nsegments_hori-1],op,yint_m,partialsums,y_m);
	int p_base = (g.Nsegments_hori-1)*g.segment_size_hori;
	for(int p_iter=p_base; (p_iter<p_base+g.segment_size_hori) && (p_iter<g.Nsnp) ; p_iter++){
		for(int k_iter=0;k_iter<Ncol_op;k_iter++)
			res(p_iter,k_iter) = y_m[p_iter-p_base][k_iter];
	}

	#if DEBUG==1
		if(debug){
			print_time ();
			cout <<"Ending mailman on premultiply"<<endl;
		}
	#endif

	if(!subtract_means)
		return;

	for(int p_iter=0;p_iter<p;p_iter++){
 		for(int k_iter=0;k_iter<Ncol_op;k_iter++){
			res(p_iter,k_iter) = res(p_iter,k_iter) - (g.get_col_mean(p_iter)*sum_op[k_iter]);
			if(var_normalize)
				res(p_iter,k_iter) = res(p_iter,k_iter)/(g.get_col_std(p_iter));
 		}
 	}

}

void multiply_y_post_fast(MatrixXdr &op_orig, int Nrows_op, MatrixXdr &res,bool subtract_means){

	MatrixXdr op;
	op = op_orig.transpose();

	if(var_normalize && subtract_means){
		for(int p_iter=0;p_iter<p;p_iter++){
			for(int k_iter=0;k_iter<Nrows_op;k_iter++)
				op(p_iter,k_iter) = op(p_iter,k_iter) / (g.get_col_std(p_iter));
		}
	}

	#if DEBUG==1
		if(debug){
			print_time ();
			cout <<"Starting mailman on postmultiply"<<endl;
		}
	#endif

	int Ncol_op = Nrows_op;

	//cout << "ncol_op = " << Ncol_op << endl;

	int seg_iter;
	for(seg_iter=0;seg_iter<g.Nsegments_hori-1;seg_iter++){
		mailman::fastmultiply_pre(g.segment_size_hori,g.Nindv,Ncol_op, seg_iter * g.segment_size_hori, g.p[seg_iter],op,yint_e,partialsums,y_e);
	}
	int last_seg_size = (g.Nsnp%g.segment_size_hori !=0 ) ? g.Nsnp%g.segment_size_hori : g.segment_size_hori;
	mailman::fastmultiply_pre(last_seg_size,g.Nindv,Ncol_op, seg_iter * g.segment_size_hori, g.p[seg_iter],op,yint_e,partialsums,y_e);

	for(int n_iter=0; n_iter<n; n_iter++)  {
		for(int k_iter=0;k_iter<Ncol_op;k_iter++) {
			res(k_iter,n_iter) = y_e[n_iter][k_iter];
			y_e[n_iter][k_iter] = 0;
		}
	}

	#if DEBUG==1
		if(debug){
			print_time ();
			cout <<"Ending mailman on postmultiply"<<endl;
		}
	#endif

	if(!subtract_means)
		return;

	double *sums_elements = new double[Ncol_op];
 	memset (sums_elements, 0, Nrows_op * sizeof(int));

 	for(int k_iter=0;k_iter<Ncol_op;k_iter++){
 		double sum_to_calc=0.0;
 		for(int p_iter=0;p_iter<p;p_iter++)
 			sum_to_calc += g.get_col_mean(p_iter)*op(p_iter,k_iter);
 		sums_elements[k_iter] = sum_to_calc;
 	}
 	for(int k_iter=0;k_iter<Ncol_op;k_iter++){
 		for(int n_iter=0;n_iter<n;n_iter++)
 			res(k_iter,n_iter) = res(k_iter,n_iter) - sums_elements[k_iter];
 	}

}

void multiply_y_pre_naive_mem(MatrixXdr &op, int Ncol_op ,MatrixXdr &res){
	for(int p_iter=0;p_iter<p;p_iter++){
		for(int k_iter=0;k_iter<Ncol_op;k_iter++){
			double temp=0;
			for(int n_iter=0;n_iter<n;n_iter++)
				temp+= g.get_geno(p_iter,n_iter,var_normalize)*op(n_iter,k_iter);
			res(p_iter,k_iter)=temp;
		}
	}
}

void multiply_y_post_naive_mem(MatrixXdr &op, int Nrows_op ,MatrixXdr &res){
	for(int n_iter=0;n_iter<n;n_iter++){
		for(int k_iter=0;k_iter<Nrows_op;k_iter++){
			double temp=0;
			for(int p_iter=0;p_iter<p;p_iter++)
				temp+= op(k_iter,p_iter)*(g.get_geno(p_iter,n_iter,var_normalize));
			res(k_iter,n_iter)=temp;
		}
	}
}

void multiply_y_pre_naive(MatrixXdr &op, int Ncol_op ,MatrixXdr &res){
	res = geno_matrix * op;
}

void multiply_y_post_naive(MatrixXdr &op, int Nrows_op ,MatrixXdr &res){
	res = op * geno_matrix;
}

void multiply_y_post(MatrixXdr &op, int Nrows_op ,MatrixXdr &res,bool subtract_means){
    if(fast_mode)
        multiply_y_post_fast(op,Nrows_op,res,subtract_means);
    else{
		if(memory_efficient)
			multiply_y_post_naive_mem(op,Nrows_op,res);
		else
			multiply_y_post_naive(op,Nrows_op,res);
	}
}

void multiply_y_pre(MatrixXdr &op, int Ncol_op ,MatrixXdr &res,bool subtract_means){
    if(fast_mode)
        multiply_y_pre_fast(op,Ncol_op,res,subtract_means);
    else{
		if(memory_efficient)
			multiply_y_pre_naive_mem(op,Ncol_op,res);
		else
			multiply_y_pre_naive(op,Ncol_op,res);
	}
}
float compute_ELBO_svi(int j, MatrixXdr &mu_k, MatrixXdr &alpha, MatrixXdr &sigma_k, double sigma_beta, double sigma_e,double pi, MatrixXdr &XsXs, int miniB, MatrixXdr &Xs, MatrixXdr &ys)
{
//        MatrixXdr Xs =geno_matrix.block(0, j*miniB, g.Nsnp, miniB);
 //       MatrixXdr ys =pheno.block(j*miniB, 0, miniB, 1);

        float ll = -g.Nindv/2* log(sigma_e);
        MatrixXdr r = mu_k.cwiseProduct(alpha);
        MatrixXdr Xsr = ys - Xs.transpose()*r;
        MatrixXdr Xnorm = Xsr.cwiseProduct(Xsr);
        ll -= Xnorm.sum()*g.Nindv/2/sigma_e/miniB;
        for(int i =0; i<g.Nsnp ; i++)
        {
                float temp = alpha(i,0)*(sigma_k(i,0)+mu_k(i,0)*mu_k(i,0));
                temp = temp - alpha(i,0)*mu_k(i,0)*alpha(i,0)*mu_k(i,0);
                ll -= g.Nindv*XsXs(i,i)*temp /2 /sigma_e/miniB;
                if(alpha(i,0)!=0 && alpha(i,0)!=1 && pi !=1 && pi !=0){
                float temp2 = alpha(i,0)*log(alpha(i,0)/pi);
                ll -= temp2;

                float temp3 = (1-alpha(i,0))*log((1-alpha(i,0))/(1-pi));
                ll -= temp3;
                }
                float temp4 = 1+ log(sigma_k(i,0)/ sigma_beta) - (sigma_k(i,0)+ mu_k(i,0)*mu_k(i,0))/sigma_beta;
                ll += alpha(i,0)*temp4 /2;
        }
        return ll;
}

void svi_step(MatrixXdr &mu_k, MatrixXdr &alpha, MatrixXdr &sigma_k, double &sigma_beta, double &sigma_e, double &pi, int miniB, int &t, float &ll, int step_a, int step_b, int iter)
{
		//shuffle data
		//std::vector<int> perm;
		//for(int i=0; i<g.Nindv; i++)
		//	perm.push_back(i);
		//std::random_shuffle(perm.begin(), perm.end());

		//shuffle data
		//permutation matrix
		MatrixXdr perm_max = MatrixXdr::Constant(g.Nindv,g.Nindv,0);
		//random generator
		std::vector<int> perm;
		for (int i=0; i<g.Nindv; i++)
			perm.push_back(i);
		std::random_shuffle(perm.begin(), perm.end());
		for (int i=0; i<g.Nindv; i++)
		{
			int idx = perm[i];
			perm_max(i,idx)=1;

		}

		//shuffle
		MatrixXdr geno_shuffle = geno_matrix * perm_max.transpose();
		MatrixXdr pheno_shuffle = perm_max * pheno;

		 for( int j=0; j<g.Nindv/miniB; j++)
                {
                        double step_temp = pow(t+step_a, step_b);
                        double step = 1/step_temp;
			//set this if force to forget
			step=1;
                        MatrixXdr Xs = geno_shuffle.block(0, j*miniB, g.Nsnp, miniB);//take Bth block
                        MatrixXdr XsXs= Xs * Xs.transpose(); //M*M matrix
			MatrixXdr ys = pheno_shuffle.block(j*miniB, 0, miniB, 1);
                        MatrixXdr Xsys = Xs * ys;
                       //E step
			for( int i=0; i<g.Nsnp; i++)
                        {
				//sigma_m
				// 1/(g.Nindv*XsXs(i,i)/miniB/sigma_e + 1/sigma_beta) for previous parameterization
				double sigma_k_update = 1/ (g.Nindv*XsXs(i,i)/miniB/sigma_e + 1/sigma_beta);
				sigma_k(i,0) =sigma_k_update;

				//mu_m
				MatrixXdr cur_snp = XsXs.block(0,i,g.Nsnp, 1);
	                        cur_snp = cur_snp.cwiseProduct(alpha);
        	                cur_snp = cur_snp.cwiseProduct(mu_k);
                	        MatrixXdr result = cur_snp.colwise().sum();
                        	result(0,0) = result(0,0) - cur_snp(i, 0);
                        	double mu_k_update = Xsys(i,0);
                        	mu_k_update = mu_k_update - result(0,0);
                        	mu_k_update =g.Nindv*  sigma_k(i,0)* mu_k_update  / sigma_e/miniB;
	                        mu_k(i,0)= mu_k_update;
				//update alpha
				double temp = mu_k(i,0)*mu_k(i,0) / sigma_k(i,0)/2;
                        	double temp2 = sqrt(sigma_k(i,0))/sqrt(sigma_beta);
                        	double temp3 =  pi / (1-pi);
                        	double temp4 = exp(temp)*temp2 *temp3;
                        	double alpha_pre = alpha(i,0)/ (1-alpha(i,0));
                        	double update = pow(alpha_pre, 1-step);
                        	update *= pow(temp4, step);
                        	update *= pow(alpha_pre, 1-step);
				alpha(i,0) = temp4/(temp4+1);
				}
			//M-step
			//update theta
			//sigma_beta
	                MatrixXdr result = mu_k.cwiseProduct(mu_k) + sigma_k;
	                result  = result.cwiseProduct(alpha);
        	        MatrixXdr weight_sum = result.colwise().sum();
                	MatrixXdr alpha_sum = alpha.colwise().sum();
                	sigma_beta =(1-step)*sigma_beta + step* weight_sum(0,0) /alpha_sum(0,0);

                	//causal proportion pi
			pi = pi*(1-step) + step*alpha_sum(0,0) / g.Nsnp;

			//sigma_e update
              		MatrixXdr ysys = ys.transpose()*ys;
                	double sigma_e_update = ysys(0,0);
                	MatrixXdr weight_mu = alpha.cwiseProduct(mu_k);
                	MatrixXdr X_mu_alpha = Xs.transpose() * weight_mu;
                	MatrixXdr yX_mu_alpha = ys.transpose()* X_mu_alpha;
                	sigma_e_update  = sigma_e_update - yX_mu_alpha(0,0)*2;
                	MatrixXdr X_square =Xs.cwiseProduct(Xs);
                	MatrixXdr X_sigma_mu_alpha = X_square.transpose()*result;
                	MatrixXdr X_sigma_mu_alpha_sum = X_sigma_mu_alpha.colwise().sum();
                	sigma_e_update += X_sigma_mu_alpha_sum(0,0);
                	for(int b=0; b<miniB; b++)
                	{
                        	MatrixXdr ind_i = Xs.block(0, b, g.Nsnp, 1);
                        	ind_i = ind_i.cwiseProduct(mu_k);
                       		ind_i = ind_i.cwiseProduct(alpha);
                        	MatrixXdr temp = ind_i * ind_i.transpose();
                        	double temp_sum = temp.sum() - temp.trace();
                        	sigma_e_update += temp_sum;
                	}
			sigma_e_update = sigma_e_update / miniB;
               		sigma_e = (1-step)*sigma_e + step * sigma_e_update;
			ll=compute_ELBO_svi(j, mu_k, alpha, sigma_k, sigma_beta, sigma_e, pi, XsXs, miniB, Xs, ys);
			cout<<iter<<"\t"<<t<< "\t"<<ll;
                	double vg = pi * g.Nsnp * sigma_beta;
                	double h2g = vg / (vg + sigma_e);
               	 	cout<<"\t" <<pi << "\t" << vg <<"\t" <<h2g << endl;
			t++;
		}
		return;
}
void vem_step(MatrixXdr &mu_k, MatrixXdr &alpha, MatrixXdr &sigma_k, double &sigma_beta, double &sigma_e, double &pi, MatrixXdr &XX, MatrixXdr &Xy)
{
                for( int i=0; i<g.Nsnp; i++)
                {
                        sigma_k(i,0) = 1/ (XX(i,i)/sigma_e + 1/sigma_beta);
                        MatrixXdr cur_snp = XX.block(0, i, g.Nsnp, 1);
                        cur_snp = cur_snp.cwiseProduct(alpha);
                        cur_snp = cur_snp.cwiseProduct(mu_k);
                        MatrixXdr result = cur_snp.colwise().sum();
                        result(0,0) = result(0,0) - cur_snp(i, 0);
                        mu_k(i, 0) = Xy(i, 0);
                        mu_k(i, 0) = mu_k(i,0) - result(0,0);
                        mu_k(i,0) = sigma_k(i,0)* mu_k(i,0)  / sigma_e;
			//update alpha
		 	double temp = mu_k(i,0)*mu_k(i,0) / sigma_k(i,0)/2;
                        double temp2 = sqrt(sigma_k(i,0))/sqrt(sigma_beta);
                        double temp3 =  pi / (1-pi);
                        double temp4 = exp(temp)*temp2 *temp3;
                        alpha(i,0)  = temp4 / (temp4+1);
                }
                //update Theta
                MatrixXdr result = mu_k.cwiseProduct(mu_k) + sigma_k;
                result  = result.cwiseProduct(alpha);
                MatrixXdr weight_sum = result.colwise().sum();
                MatrixXdr alpha_sum = alpha.colwise().sum();
                sigma_beta = weight_sum(0,0) /alpha_sum(0,0);

                pi = alpha_sum(0,0) / g.Nsnp;
                cout<<"PI: "<<pi<<endl;

                MatrixXdr yy = pheno.transpose() * pheno;
                sigma_e = yy(0,0);
                MatrixXdr weight_mu = alpha.cwiseProduct(mu_k);
                MatrixXdr X_mu_alpha = geno_matrix.transpose() *  weight_mu;
                MatrixXdr yX_mu_alpha = pheno.transpose() * X_mu_alpha;
                sigma_e = sigma_e - yX_mu_alpha(0,0)*2;
                double elbo_temp=0;
                MatrixXdr X_square = geno_matrix.cwiseProduct(geno_matrix);
                MatrixXdr X_sigma_mu_alpha = X_square.transpose() * result;
                MatrixXdr X_sigma_mu_alpha_sum = X_sigma_mu_alpha.colwise().sum();
                sigma_e += X_sigma_mu_alpha_sum(0,0);
                elbo_temp  = X_sigma_mu_alpha_sum(0,0);
                //cout<<"sigma_e " <<sigma_e<<endl;
		 for(int i=0; i<g.Nindv; i++)
                {
                        MatrixXdr ind_i = geno_matrix.block(0,i,g.Nsnp,1);
                        ind_i  = ind_i.cwiseProduct(mu_k);
                        ind_i = ind_i.cwiseProduct(alpha);
                        MatrixXdr temp = ind_i * ind_i.transpose();
                        double temp_sum = temp.sum()- temp.trace();
                        sigma_e += temp_sum;
                        elbo_temp += temp_sum;
                }
		 sigma_e = sigma_e / g.Nindv;
		return;

}

float compute_ELBO_vem(MatrixXdr &mu_k, MatrixXdr &alpha, MatrixXdr &sigma_k, double sigma_beta, double sigma_e, double pi, MatrixXdr &XX)
{

	float ll = -g.Nindv/2* log(sigma_e);
	MatrixXdr r = mu_k.cwiseProduct(alpha);
	MatrixXdr Xr = pheno- geno_matrix.transpose()*r;
	MatrixXdr Xrnorm = Xr.cwiseProduct(Xr);
	ll -= Xrnorm.sum()/2/sigma_e;
	for(int i =0; i<g.Nsnp ; i++)
	{
		float temp = alpha(i,0)*(sigma_k(i,0)+mu_k(i,0)*mu_k(i,0));
		temp = temp - alpha(i,0)*mu_k(i,0)*alpha(i,0)*mu_k(i,0);
		ll -= XX(i,i)*temp /2 /sigma_e;
		if(alpha(i,0)!=0 && alpha(i,0)!=1){
		float temp2 = alpha(i,0)*log(alpha(i,0)/pi);
		ll -= temp2;

		float temp3 = (1-alpha(i,0))*log((1-alpha(i,0))/(1-pi));
		ll -= temp3;
		}
		float temp4 = 1+ log(sigma_k(i,0)/ sigma_beta) - (sigma_k(i,0)+ mu_k(i,0)*mu_k(i,0))/sigma_beta;
		ll += alpha(i,0)*temp4 /2;
	}
	return ll;
}
float compute_ELBO(MatrixXdr &mu_k, MatrixXdr &alpha, MatrixXdr &sigma_k, double sigma_beta, double sigma_e, double pi,double elbo_temp)
{
	MatrixXdr yy = pheno.transpose() * pheno;
	MatrixXdr mualpha = mu_k.cwiseProduct(alpha);
	MatrixXdr Xmualpha = geno_matrix.transpose() * mualpha;
	Xmualpha = pheno.cwiseProduct(Xmualpha);
	float ll = - g.Nindv/2* std::log(sigma_e) -yy(0,0)/2/sigma_e + Xmualpha.sum()/2/sigma_e;
	ll -= alpha.sum()* std::log(sigma_beta) / 2;
	MatrixXdr mu_square = mu_k.cwiseProduct(mu_k);
	MatrixXdr alpha_muS = alpha.cwiseProduct(mu_square + sigma_k);
	ll -= alpha_muS.sum() / sigma_beta/2;
	ll += std::log(pi)*alpha.sum();
	ll += std::log(1-pi)*(g.Nsnp - alpha.sum());
	ll -= alpha.sum()/2;
	for(int i=0; i<g.Nsnp; i++)
 	{
		ll -= alpha(k,0)*std::log(sigma_k(k,0))/2;
		ll -= alpha(k,0)*std::log(alpha(k,0));
		ll -= (1-alpha(k,0)) * std::log (1-alpha(k,0));
	}
	ll -= elbo_temp / 2 / sigma_e;
	return ll;

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

	pair<double,double> prev_error = make_pair(0.0,0.0);
	double prevnll=0.0;

	parse_args(argc,argv);

	//TODO: Memory effecient Version of Mailman

	memory_efficient = command_line_opts.memory_efficient;
	text_version = command_line_opts.text_version;
	fast_mode = command_line_opts.fast_mode;
	missing = command_line_opts.missing;
	reg = command_line_opts.reg;
	gwas=command_line_opts.gwas;
	if(gwas)
		fast_mode=false; //for now, gwas need the genotype matrix, and compute kinship constructed with one chrom leave out
	if(!reg)
		fast_mode=false; //force save whole genome if non randomized
	if(text_version){
		if(fast_mode)
			g.read_txt_mailman(command_line_opts.GENOTYPE_FILE_PATH,missing);
		else
			g.read_txt_naive(command_line_opts.GENOTYPE_FILE_PATH,missing);
	}
	else{
		g.read_plink(command_line_opts.GENOTYPE_FILE_PATH,missing,fast_mode);

	}

	//TODO: Implement these codes.
	if(missing && !fast_mode){
		cout<<"Missing version works only with mailman i.e. fast mode\n EXITING..."<<endl;
		exit(-1);
	}
	if(fast_mode && memory_efficient){
		cout<<"Memory effecient version for mailman EM not yet implemented"<<endl;
		cout<<"Ignoring Memory effecient Flag"<<endl;
	}
	if(missing && var_normalize){
		cout<<"Missing version works only without variance normalization\n EXITING..."<<endl;
		exit(-1);
	}

	MAX_ITER=100;
    //	MAX_ITER =  command_line_opts.max_iterations ;
	int B = command_line_opts.batchNum;
	int pheno_idx = command_line_opts.pheno_idx;
	k_orig = command_line_opts.num_of_evec ;
	debug = command_line_opts.debugmode ;
	float tr2= command_line_opts.tr2;
	check_accuracy = command_line_opts.getaccuracy;
	var_normalize = true;
	accelerated_em = command_line_opts.accelerated_em;
	k = k_orig + command_line_opts.l;
	k = (int)ceil(k/10.0)*10;
	command_line_opts.l = k - k_orig;
	p = g.Nsnp;
	n = g.Nindv;
	bool toStop=false;
		toStop=true;
	srand((unsigned int) time(0));

	means.resize(p,1);
	stds.resize(p,1);
	sum2.resize(p,1);
	sum.resize(p,1);

//	geno_matrix.resize(p,n);
//	g.generate_eigen_geno(geno_matrix, var_normalize);

	if(!fast_mode && !memory_efficient){
		geno_matrix.resize(p,n);
//		g.generate_eigen_geno(geno_matrix,var_normalize);
		g.generate_eigen_geno(geno_matrix, true);
		cout<<geno_matrix.data()<<endl;
		cout<<geno_matrix.rows();
		cout<<geno_matrix.cols();

	}

	for(int i=0;i<p;i++){
		means(i,0) = g.get_col_mean(i);
		stds(i,0) =1/ g.get_col_std(i);
		sum2(i,0) =g.get_col_sum2(i);
		sum(i,0)= g.get_col_sum(i);
	}

//	cout<<"printing means: "<<endl<<means<<endl;
//	cout<<"printing std: "<<endl<<stds<<endl;
	ofstream c_file;

	cout<<"Running on Dataset of "<<g.Nsnp<<" SNPs and "<<g.Nindv<<" Individuals"<<endl;

	#if SSE_SUPPORT==1
		if(fast_mode)
			cout<<"Using Optimized SSE FastMultiply"<<endl;
	#endif

	//get geno
	//cout<<g.get_geno(0,0,false);
	//read phenotype
	//
	//
	std::string filename=command_line_opts.PHENOTYPE_FILE_PATH;
	int pheno_num= read_pheno2(g.Nindv, filename, pheno_idx);
	int cov_num=0 ;
	if(filename=="")
	{
		cout<<"No Phenotype File Specified"<<endl;
		return 0 ;
	}
	cout<< "Read in "<<pheno_num << " phenotypes"<<endl;
	if(pheno_idx!=0)
		cout<<"Using phenotype "<<pheno_name[pheno_idx-1]<<endl;

	MatrixXdr y_sum=pheno.colwise().sum();
	MatrixXdr y_mean = y_sum/g.Nindv;
	for(int i=0; i<g.Nindv; i++)
		pheno.block(i,0,1,pheno_num) =pheno.block(i,0,1,pheno_num) - y_mean; //center phenotype
	y_sum=pheno.colwise().sum();

//model
	float hyper[4];
	for(int i=0; i<4; i++)
		hyper[i]=1e-3;
	float tol = 1e-6;
	MatrixXdr mu_k = MatrixXdr::Constant(g.Nsnp, 1, 0);
	MatrixXdr sigma_k=MatrixXdr::Constant(g.Nsnp, 1, 1);
	MatrixXdr alpha= MatrixXdr::Constant(g.Nsnp,1, 0.5);

	MatrixXdr beta(g.Nsnp, 1);
	double pi=0.5;
	//exponential distribution
	boost::mt19937 gen;
	boost::exponential_distribution<> dist(1);
	boost::variate_generator<boost::mt19937&, boost::exponential_distribution<> > init(gen,dist);

	double sigma_beta = 1/ (init());
	double sigma_e = 1/(init());
//loat l=init();
//	cout<<"Initializing prarameters..."<<endl<<"a: "<<a <<endl<<"l: "<<l <<endl;

	float rho;
	int iter=0;

	float ll_det=1;
	float prevll=0 ;
//	MatrixXdr XX= geno_matrix*geno_matrix.transpose();
//	MatrixXdr Xy = geno_matrix * pheno;

	double step_a =1;
	double step_b =1;
	int t=0;
	int miniB=g.Nindv/B;
	cout<< "iter\tt\tELBO\tpi\tvg\th2g"<<endl;
	while(iter<=MAX_ITER ){
//	while(iter<= MAX_ITER && ll_det > tol){

		/* one sample
		MatrixXdr weight_mu = alpha.cwiseProduct(mu_k);
		MatrixXdr Xsmualpha = Xs.transpose()*weight_mu;
		double y_Xmualpha = pheno(j,0) - Xsmualpha(0,0);
		double sigma_e_update = y_Xmualpha*y_Xmualpha;

		MatrixXdr XsXs = Xs.cwiseProduct(Xs);
		MatrixXdr varB = weight_mu.cwiseProduct(weight_mu);
		MatrixXdr mumu = mu_k.cwiseProduct(mu_k)+sigma_k;
		MatrixXdr varB_temp = alpha.cwiseProduct(mumu);
		varB = varB_temp - varB;
		MatrixXdr varB_result = XsXs.cwiseProduct(varB);
		sigma_e_update += varB_result(0,0);
		sigma_e = (1-step)*sigma_e+ step*sigma_e_update;

				*/

	/*
		//variational EM step
		vem_step(mu_k, alpha, sigma_k, sigma_beta, sigma_e, pi, XX, Xy);
	//	variational EM ELBO
		float ll = compute_ELBO_vem(mu_k, alpha,sigma_k, sigma_beta, sigma_e, pi,XX);
		cout<<iter<<"\t"<<iter<<"\t"<<ll<<"\t"<<pi ;
		double vg = sigma_beta * pi * g.Nsnp ;
		double h2g  = vg / (vg+sigma_e);
		cout<<"\t"<<vg<<"\t"<<h2g<<endl;
	*/
		//svi step,ELBO
		float ll;
		svi_step(mu_k, alpha, sigma_k, sigma_beta, sigma_e, pi, miniB,t, ll,step_a, step_b,iter);
		if(iter>0)
			ll_det = ll - prevll;
		prevll = ll;
		iter++;

	}
	return 0;
}


#ifndef QST_WAVEFUNCTIONPOSITIVE_HPP
#define QST_WAVEFUNCTIONPOSITIVE_HPP
#include <iostream>
#include <Eigen/Dense>
#include <random>
#include <fstream>
#include <iostream>

namespace qst{

class WavefunctionPositive{

    int N_;            // Number of degrees of freedom (visible units)
    int npar_;         // Number of parameters
    int nparLambda_;   // Number of amplitude parameters
    Rbm rbmAm_;        // RBM for the amplitude

    const std::complex<double> I_; // Imaginary unit

    //Random number generator
    std::mt19937 rgen_;

public:
    // Constructor
    WavefunctionPositive(Parameters &par):rbmAm_(par),
                                  I_(0,1){
        npar_ = rbmAm_.Npar();  // Total number of parameters
        nparLambda_ = npar_;
        N_ = rbmAm_.Nvisible();
        std::random_device rd;
        //rgen_.seed(rd());
        rgen_.seed(13579);
    }

    // Private members access functions
    inline int N()const{
        return N_;
    }
    inline int Npar()const{
        return npar_;
    }
    inline int NparLambda()const{
        return nparLambda_;
    }
    inline int NparMu()const{
        return 0;
    }
    inline int Nchains(){
        return rbmAm_.Nchains();
    }
    inline Eigen::VectorXd VisibleStateRow(int s){
        return rbmAm_.VisibleStateRow(s);
    }

    // Set the state of the wavefunction's degrees of freedom
    inline void SetVisibleLayer(Eigen::MatrixXd v){
        rbmAm_.SetVisibleLayer(v);
    }

    // Initialize the wavefunction parameters
    void InitRandomPars(int seed,double sigma){
        rbmAm_.InitRandomPars(seed,sigma);
    }

    // Amplitude
    double amplitude(const Eigen::VectorXd & v){
        return std::sqrt(rbmAm_.prob(v));//exp(0.5*rbmAm_.LogVal(v));
    }
    // Phase
    double phase(const Eigen::VectorXd & v){
        return 0.0;
    }
    // Psi
    std::complex<double> psi(const Eigen::VectorXd & v){
        return amplitude(v);
    }

    //---- SAMPLING ----/

    // Perform k steps of Gibbs sampling
    void Sample(int steps){
        rbmAm_.Sample(steps);
    }

    //---- DERIVATIVES ----//

    //Compute gradient of effective energy wrt Lambda
    Eigen::VectorXd LambdaGrad(const Eigen::VectorXd & v){
        return rbmAm_.VisEnergyGrad(v);
    }
    //Compute gradient of effective energy wrt all parameters
    Eigen::VectorXd Grad(const Eigen::VectorXd & v){
        return rbmAm_.VisEnergyGrad(v);
    }
    //Dummy function
    void rotatedGrad(const std::vector<std::string> & basis,
                            const Eigen::VectorXd & state,//VectorRbmT & gradR){
                            std::map<std::string,Eigen::MatrixXcd> & Unitaries,
                            Eigen::VectorXcd &gradR ){
    }

    //---- UTILITIES ----//

    //Get RBM parameters
    Eigen::VectorXd GetParameters(){
        return rbmAm_.GetParameters();
    }

    // Set RBM parameters
    void SetParameters(const Eigen::VectorXd & pars){
        rbmAm_.SetParameters(pars);
    }

    void LoadWeights(std::string &fileName){
        std::ifstream fin(fileName);
        rbmAm_.LoadWeights(fin);
    }

    void SaveWeights(std::string &fileName){
        std::ofstream fout(fileName);
        rbmAm_.SaveWeights(fout);
    }
};
}

#endif


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

#include <rclcpp/node.hpp>
#include <rclcpp/executors.hpp>

#include <rmf_fleet_msgs/msg/fleet_state.hpp>

#include <cmath>

#include <Eigen/Geometry>

using Location = rmf_fleet_msgs::msg::Location;
using RobotState = rmf_fleet_msgs::msg::RobotState;
using RobotMode = rmf_fleet_msgs::msg::RobotMode;
using FleetState = rmf_fleet_msgs::msg::FleetState;

Location make_location(float x, float y, float yaw)
{
  Location l;
  l.x = x;
  l.y = y;
  l.yaw = yaw;
  l.level_name = "test_level";
  return l;
}

Eigen::Vector3d to_eigen(Location l)
{
  return Eigen::Vector3d(l.x, l.y, l.yaw);
}

Location to_location(Eigen::Vector3d v)
{
  return make_location(v[0], v[1], v[2]);
}

class DumpFleetStates : public rclcpp::Node
{
public:

  DumpFleetStates()
  : rclcpp::Node("dump_test_fleet_states")
  {
    fleet_state_publisher = create_publisher<FleetState>(
      "/fleet_states", rclcpp::SystemDefaultsQoS());

    timer = create_wall_timer(
      std::chrono::milliseconds(period_ms), [&]() { update(); });

    current_location = Eigen::Vector3d(0.0, 0.0, 0.0);

    path.emplace_back(make_location(distance/2.0, 0.0, 0.0));
    path.emplace_back(make_location(distance, 0.0, 0.0));
    path.emplace_back(make_location(distance, 0.0, 90.0*M_PI/180.0));
    path.emplace_back(make_location(distance, distance/4.0, 90.0*M_PI/180.0));
    path.emplace_back(make_location(distance, distance, 90.0*M_PI/180.0));
  }

private:

  void update()
  {
    std::cout << "Current location: " << current_location.transpose() <<
      std::endl;

    if (path.empty())
      return publish();

    const Eigen::Vector3d next_location = to_eigen(path.front());

    const Eigen::Vector2d p = current_location.block<2, 1>(0, 0);
    const Eigen::Vector2d p_next = next_location.block<2, 1>(0, 0);

    const Eigen::Vector2d n = p_next - p;
    const double dist = n.norm();
    if (dist <= velocity*timestep)
    {
      const double yaw = current_location[2];
      const double next_yaw = next_location[2];
      const double d_yaw = next_yaw - yaw;
      if (std::abs(d_yaw) <= velocity*timestep)
      {
        current_location = next_location;
        path.erase(path.begin());
        return publish();
      }

      current_location.block<2, 1>(0, 0) = p_next;
      current_location[2] += d_yaw;

      return publish();
    }

    const Eigen::Vector2d dir = n/dist;
    current_location.block<2, 1>(0, 0) += velocity*timestep*dir;
    return publish();
  }

  void publish()
  {
    FleetState fleet;
    fleet.name = "test_fleet";

    RobotState robot;
    robot.name = "test_robot";
    robot.mode.mode = RobotMode::MODE_MOVING;
    robot.location = to_location(current_location);
    robot.location.t = now();
    robot.path = path;
    robot.task_id = "test";
    fleet.robots.push_back(robot);

    fleet_state_publisher->publish(fleet);
  }

  const double distance = 10.0;
  const double velocity = 0.4;
  const int64_t period_ms = 200;
  const double timestep = static_cast<double>(period_ms)/1000.0;

  Eigen::Vector3d current_location;

  std::vector<Location> path;

  rclcpp::Publisher<FleetState>::SharedPtr fleet_state_publisher;

  rclcpp::TimerBase::SharedPtr timer;

};

int main(int argc, char* argv[])
{
  rclcpp::init(argc, argv);
  rclcpp::spin(std::make_shared<DumpFleetStates>());
  rclcpp::shutdown();
}


#include <cstdio>
#include <vector>

#include <Eigen/Core>
#include <Eigen/StdVector>
#include <maplab-common/test/testing-entrypoint.h>
#include <vocabulary-tree/distance.h>
#include <vocabulary-tree/simple-kmeans.h>
#include <vocabulary-tree/types.h>

#include "./floating-point-test-helpers.h"

TEST(VocabularyTree, SimpleKMeans_AcceleratedKMeansCluster) {
  std::mt19937 generator(40);
  static const size_t kNumfeaturesPerCluster = 100;
  static const size_t kNumClusters = 100;
  DescriptorVector gt_centers;
  DescriptorVector descriptors;
  std::vector<unsigned int> membership;
  std::vector<unsigned int> gt_membership;

  GenerateTestData(
      kNumfeaturesPerCluster, kNumClusters, generator(), &gt_centers,
      &descriptors, &membership, &gt_membership);

  // Init with ground-truth.
  std::shared_ptr<DescriptorVector> centers =
      aligned_shared<DescriptorVector>(gt_centers);

  DescriptorType descriptor_zero;
  descriptor_zero.setConstant(
      kDescriptorDimensionality, 1, static_cast<Scalar>(0));

  loop_closure::SimpleKmeans<DescriptorType,
                             loop_closure::distance::L2<DescriptorType> >
      kmeans(descriptor_zero);

  kmeans.SetInitMethod(
      loop_closure::InitGiven<DescriptorType>(descriptor_zero));

  kmeans.Cluster(descriptors, kNumClusters, generator(), &membership, &centers);

  std::vector<unsigned int> membercnt;
  membercnt.resize(centers->size(), 0);
  for (size_t i = 0; i < membership.size(); ++i) {
    unsigned int member = membership[i];
    EXPECT_EQ(membership[i], gt_membership[i]);
    ++membercnt[member];
  }
  for (size_t i = 0; i < membercnt.size(); ++i) {
    EXPECT_NE(membercnt[i], static_cast<unsigned int>(0));
  }

  loop_closure::distance::L2<DescriptorType> l2_distance;

  for (size_t descriptor_idx = 0; descriptor_idx < descriptors.size();
       ++descriptor_idx) {
    DescriptorType& descriptor = descriptors.at(descriptor_idx);
    int closest_center = -1;
    unsigned int closest_distance = std::numeric_limits<unsigned int>::max();
    unsigned int second_closest_distance =
        std::numeric_limits<unsigned int>::max();
    for (size_t center_idx = 0; center_idx < centers->size(); ++center_idx) {
      unsigned int distance = l2_distance(descriptor, centers->at(center_idx));
      if (distance < closest_distance) {
        second_closest_distance = closest_distance;
        closest_distance = distance;
        closest_center = center_idx;
      }
    }
    // Check that we don't have a trivial solution.
    EXPECT_NE(second_closest_distance, closest_distance);
    EXPECT_NE(closest_center, -1);
  }
}

MAPLAB_UNITTEST_ENTRYPOINT


/*
 * ARL_Topologies - An extensible topology optimization program
 *
 * Written in 2017 by Raymond A. Wildman <raymond.a.wildman.civ@mail.mil>
 * This project constitutes a work of the United States Government and is not
 * subject to domestic copyright protection under 17 USC Sec. 105.
 * Release authorized by the US Army Research Laboratory
 *
 * To the extent possible under law, the author(s) have dedicated all copyright
 * and related and neighboring rights to this software to the public domain
 * worldwide. This software is distributed without any warranty.
 *
 * You should have received a copy of the CC0 Public Domain Dedication along
 * with this software. If not, see <http://creativecommons.org/publicdomain/zero/1.0/>.
 *
 */

#include <Eigen/IterativeLinearSolvers>
//#include <Eigen/SparseLU>
#include "femproblem.h"
#include "element.h"
#include "lintri.h"
#include "linquad.h"
#include "lintetra.h"
#include "trilinhex.h"
#include "mesh2d.h"
#include "mesh3d.h"
#include "REP/tomesh.h"

FEMProblem::FEMProblem(const Topologies::TOMesh* const inMesh, const Topologies::GenericMaterial& baseMat) :
	numFreeDOFs(0),
	invalid(false),
	itTol(1e-6)
{
	assert(inMesh);
	assert(baseMat.getNumParameters() > 2); // Assumes 3 material properties: density and 2 Lame parameters
	dim = inMesh->dimNum();
	if(dim == 2)
		probMesh = std::unique_ptr<FEMMesh>(new Mesh2D(inMesh, baseMat));
	else if(dim == 3)
		probMesh = std::unique_ptr<FEMMesh>(new Mesh3D(inMesh, baseMat));
	// Check element types (tris and tets need lower tolerance on iterative solver)
	if(checkForSimplex())
		itTol = 1e-12;
}

FEMProblem::~FEMProblem()
{
}

FEMProblem::FEMProblem(const FEMProblem& copy) :
	probMesh(std::unique_ptr<FEMMesh>(copy.probMesh->clone())),
	pFEMMatrix(copy.pFEMMatrix ? std::unique_ptr<EigenSparseMat>(new EigenSparseMat(*copy.pFEMMatrix)) : nullptr),
	fixedDOFs(copy.fixedDOFs),
	bfRemapVec(copy.bfRemapVec),
	loadVec(copy.loadVec),
	pVVec(copy.pVVec ? std::unique_ptr<EigenVector>(new EigenVector(*copy.pVVec)) : nullptr),
	pForce(copy.pForce ? std::unique_ptr<EigenVector>(new EigenVector(*copy.pForce)) : nullptr),
	dim(copy.dim),
	numFreeDOFs(copy.numFreeDOFs),
	invalid(copy.invalid),
	itTol(copy.itTol)
{
}

FEMProblem::FEMProblem(FEMProblem&& copy)
{
	swap(copy);
}

FEMProblem FEMProblem::operator=(FEMProblem copy)
{
	swap(copy);
	return *this;
}

void FEMProblem::swap(FEMProblem& arg2)
{
	probMesh.swap(arg2.probMesh);
	pFEMMatrix.swap(arg2.pFEMMatrix);
	pVVec.swap(arg2.pVVec);
	pForce.swap(arg2.pForce);
	std::swap(fixedDOFs, arg2.fixedDOFs);
	std::swap(bfRemapVec, arg2.bfRemapVec);
	std::swap(loadVec, arg2.loadVec);
	std::swap(dim, arg2.dim);
	std::swap(numFreeDOFs, arg2.numFreeDOFs);
	std::swap(invalid, arg2.invalid);
	std::swap(itTol, arg2.itTol);
}

bool FEMProblem::checkForSimplex() const
{
	bool foundSimplex = false;
	for(std::size_t k = 0; k < probMesh->getNumElements() && !foundSimplex; ++k)
		foundSimplex |= probMesh->getNumElementNodes(k) == (dim + 1);
	return foundSimplex;
}

Point3D FEMProblem::getMeshPoint(std::size_t kn) const
{
	if(probMesh->getDim() == 2)
	{
		Point2D pt = *dynamic_cast<Mesh2D*>(probMesh.get())->getNode(kn);
		return Point3D(pt.x, pt.y, 0.);
	}
	return *dynamic_cast<Mesh3D*>(probMesh.get())->getNode(kn);
}

void FEMProblem::changeBoundaryConditionsTo(const std::vector<ExoBC>& bcVec)
{
	std::size_t nnodes = probMesh->getNumUnknowns();
	fixedDOFs = std::vector<bool>(dim*nnodes, false);
	loadVec.clear();
	for(auto const& curBC : bcVec)
	{
		for(auto kn : curBC.nodeIDVec)
		{
			// Set up supports
			if(curBC.isSupport)
			{
				fixedDOFs[kn] = curBC.xsup;
				fixedDOFs[kn + nnodes] = curBC.ysup;
				if(dim == 3)
					fixedDOFs[kn + 2*nnodes] = curBC.zsup;
			}
			else // Set up loads
			{
				Point3D cartVec = CoordinateSystem::convertVector(curBC.loadVec, getMeshPoint(kn), curBC.ct);
				loadVec[kn] = cartVec.x;
				loadVec[kn + nnodes] = cartVec.y;
				if(dim == 3)
					loadVec[kn + 2*nnodes] = cartVec.z;
			}
		}
	}
	// Save global row numbers for compressed system (without fixed dofs)
	bfRemapVec.resize(dim*nnodes);
	std::size_t curBF = 0;
	for(std::size_t k = 0; k < fixedDOFs.size(); ++k)
	{
		if(!fixedDOFs[k])
			bfRemapVec[k] = curBF++;
	}
	numFreeDOFs = curBF;
	solveProblem();
}

void FEMProblem::solveProblem()
{
	// Set up vectors and matrices
	setMatrix();
	setVector();
	// Matrix stuff:
	// Save RHS to use for computation of compliance
	pForce = std::unique_ptr<EigenVector>(new EigenVector(*pVVec));
	invalid = false;
	// Solve
	Eigen::ConjugateGradient<EigenSparseMat,Eigen::Lower|Eigen::Upper> solver;
	solver.compute(*pFEMMatrix);
	solver.setTolerance(itTol);
	unsigned niters = 10000;
	solver.setMaxIterations(niters);
	*pVVec = solver.solve(*pForce);
/* Direct solver
	Eigen::SparseLU<EigenSparseMat> solver;
	pFEMMatrix->makeCompressed();
	solver.analyzePattern(*pFEMMatrix);
	solver.factorize(*pFEMMatrix);
	*pVVec = solver.solve(*pForce);*/

	if(solver.iterations() >= niters)
		invalid = true;
}

void FEMProblem::setMatrix()
{
	std::size_t numNodes = probMesh->getNumUnknowns();
	std::vector<EigenT> sparseVec;
	sparseVec.reserve(10*numFreeDOFs); //Size of this depends on matrix bandwidth, which is unknown as it depends on mesh connectivity
  std::size_t nelems = probMesh->getNumElements();
  for(std::size_t ielem = 0; ielem < nelems; ++ielem)
  {
    Eigen::MatrixXd elemMat = probMesh->getElementMatrix(ielem);
    assembleMatrix(sparseVec, ielem, elemMat, numNodes);
  }
	pFEMMatrix = std::unique_ptr<EigenSparseMat>(new EigenSparseMat(numFreeDOFs,numFreeDOFs));
	pFEMMatrix->setFromTriplets(sparseVec.begin(), sparseVec.end());
}

void FEMProblem::assembleMatrix(std::vector<EigenT>& rseMat, const std::size_t kelem, const Eigen::MatrixXd& elemMat, std::size_t numUnk) const
{
	if(dim == 2)
		assembleMatrix2D(rseMat, kelem, elemMat, numUnk);
	else if(dim == 3)
		assembleMatrix3D(rseMat, kelem, elemMat, numUnk);
}

void FEMProblem::assembleMatrix2D(std::vector<EigenT>& rseMat, const std::size_t kelem, const Eigen::MatrixXd& elemMat, std::size_t numUnk) const
{
	const unsigned numElemNodes = probMesh->getNumElementNodes(kelem);
	for(unsigned iTst = 0; iTst < numElemNodes; iTst++)
	{
		std::size_t gTst = probMesh->getGlobalBF(kelem, iTst);
		for(unsigned iBas = 0; iBas < numElemNodes; iBas++)
		{
			std::size_t gBas = probMesh->getGlobalBF(kelem, iBas);
			for(unsigned tstdof = 0; tstdof < dim; ++tstdof)
			{
				for(unsigned basdof = 0; basdof < dim; ++basdof)
				{
					bool isFree = !(fixedDOFs[gTst + tstdof*numUnk] || fixedDOFs[gBas + basdof*numUnk]);
					if(fabs(elemMat(iTst + tstdof*numElemNodes, iBas + basdof*numElemNodes)) > 1.e-16 && isFree)
					{
						rseMat.push_back(EigenT(bfRemapVec[gTst + tstdof*numUnk],
																		bfRemapVec[gBas + basdof*numUnk],
																		elemMat(iTst + tstdof*numElemNodes, iBas + basdof*numElemNodes)));
					}
				}
			}
		}
	}
}

void FEMProblem::assembleMatrix3D(std::vector<EigenT>& rseMat, const std::size_t kelem, const Eigen::MatrixXd& elemMat, std::size_t numUnk) const
{
	const unsigned numElemNodes = probMesh->getNumElementNodes(kelem);
	for(unsigned iTst = 0; iTst < numElemNodes; iTst++)
	{
		std::size_t gTst = probMesh->getGlobalBF(kelem, iTst);
		for(unsigned iBas = 0; iBas < numElemNodes; iBas++)
		{
			std::size_t gBas = probMesh->getGlobalBF(kelem, iBas);
			for(unsigned tstOff = 0; tstOff < dim; ++tstOff)
			{
				for(unsigned basOff = 0; basOff < dim; ++basOff)
				{
					bool isFree = !(fixedDOFs[gTst + tstOff*numUnk] || fixedDOFs[gBas + basOff*numUnk]);
					if(fabs(elemMat(dim*iTst + tstOff, dim*iBas + basOff)) > 1.e-16 && isFree)
					{
						rseMat.push_back(EigenT(bfRemapVec[gTst + tstOff*numUnk],
																		bfRemapVec[gBas + basOff*numUnk],
																		elemMat(dim*iTst + tstOff, dim*iBas + basOff)));
					}
				}
			}
		}
	}
}

void FEMProblem::setVector()
{
  pVVec = std::unique_ptr<EigenVector>(new EigenVector(numFreeDOFs));
	pVVec->setZero();
	EigenVector& rVVec = *pVVec;
	for(std::map<std::size_t,double>::const_iterator cit = loadVec.begin(); cit != loadVec.end(); ++cit)
		rVVec(bfRemapVec[cit->first]) = cit->second;
}

std::pair<double, bool> FEMProblem::computeCompliance()
{
	if(invalid)
	{
		std::cout << "Warning, FEM OFV not valid" << std::endl;
		return std::pair<double, bool>(1e6, false);
	}
	double sum = pForce->adjoint()*(*pVVec);
	return std::pair<double, bool>(sum, true);
}

double FEMProblem::elementCompliance(const std::size_t kelem, const EigenDenseMat& elemMat, std::size_t numUnk) const
{
	if(dim == 2)
		return elementCompliance2D(kelem, elemMat, numUnk);
	return elementCompliance3D(kelem, elemMat, numUnk);
}

double FEMProblem::elementCompliance2D(const std::size_t kelem, const EigenDenseMat& elemMat, std::size_t numUnk) const
{
	double res = 0.;
	const unsigned numElemNodes = probMesh->getNumElementNodes(kelem);
	for(unsigned iTst = 0; iTst < numElemNodes; iTst++)
	{
		std::size_t gTst = probMesh->getGlobalBF(kelem, iTst);
		for(unsigned iBas = 0; iBas < numElemNodes; iBas++)
		{
			std::size_t gBas = probMesh->getGlobalBF(kelem, iBas);
			for(unsigned tstdof = 0; tstdof < dim; ++tstdof)
			{
				std::size_t m = gTst + tstdof*numUnk;
				for(unsigned basdof = 0; basdof < dim; ++basdof)
				{
					std::size_t n = gBas + basdof*numUnk;
					if(!(fixedDOFs[m] || fixedDOFs[n]))
						res += elemMat(iTst + tstdof*numElemNodes, iBas + basdof*numElemNodes)*(*pVVec)(bfRemapVec[m])*(*pVVec)(bfRemapVec[n]);
				}
			}
		}
	}
	return res;
}

double FEMProblem::elementCompliance3D(const std::size_t kelem, const EigenDenseMat& elemMat, std::size_t numUnk) const
{
	double res = 0.;
	const unsigned numElemNodes = probMesh->getNumElementNodes(kelem);
	for(unsigned iTst = 0; iTst < numElemNodes; iTst++)
	{
		std::size_t gTst = probMesh->getGlobalBF(kelem, iTst);
		for(unsigned iBas = 0; iBas < numElemNodes; iBas++)
		{
			std::size_t gBas = probMesh->getGlobalBF(kelem, iBas);
			for(unsigned tstOff = 0; tstOff < dim; ++tstOff)
			{
				std::size_t m = gTst + tstOff*numUnk;
				for(unsigned basOff = 0; basOff < dim; ++basOff)
				{
					std::size_t n = gBas + basOff*numUnk;
					if(!(fixedDOFs[m] || fixedDOFs[n]))
						res += elemMat(dim*iTst + tstOff, dim*iBas + basOff)*(*pVVec)(bfRemapVec[m])*(*pVVec)(bfRemapVec[n]);
				}
			}
		}
	}
	return res;
}

std::vector<double> FEMProblem::gradCompliance(const Topologies::TOMesh* const inMesh) const
{
	std::size_t numNodes = probMesh->getNumUnknowns();
	std::size_t nelems = probMesh->getNumElements();
	// Precompute uku values
	std::vector<double> ukuVec(nelems);
	for(std::size_t ielem = 0; ielem < nelems; ++ielem)
  {
		Eigen::MatrixXd elemMat = probMesh->getElementMatrix(ielem);
		ukuVec[ielem] = -elementCompliance(ielem, elemMat, numNodes);
		ukuVec[ielem] /= inMesh->getOptVal(ielem);
	}
	return ukuVec;
}


//
// Copyright (c) 2015-2020 CNRS INRIA
// Copyright (c) 2015 Wandercraft, 86 rue de Paris 91400 Orsay, France.
//

#include "pinocchio/math/fwd.hpp"
#include "pinocchio/multibody/joint/joints.hpp"
#include "pinocchio/algorithm/rnea.hpp"
#include "pinocchio/algorithm/aba.hpp"
#include "pinocchio/algorithm/crba.hpp"
#include "pinocchio/algorithm/jacobian.hpp"
#include "pinocchio/algorithm/compute-all-terms.hpp"

#include <boost/test/unit_test.hpp>
#include <iostream>

using namespace pinocchio;

BOOST_AUTO_TEST_SUITE(JointFreeFlyer)

BOOST_AUTO_TEST_CASE(spatial)
{
  Motion v(Motion::Random());

  ConstraintIdentityTpl<double,0> constraint;
  Motion Sv = constraint * v.toVector();

  BOOST_CHECK(Sv == v);
}

BOOST_AUTO_TEST_SUITE_END()


/*
 * A c++ version of @RigidBodyManipulator/approximateIK.m
 */
#include <math.h>
#include <set>
#include <mex.h>
#include <Eigen/Dense>
#include "RigidBodyManipulator.h"
#include <iostream>
#include <Eigen/Cholesky>
#include <Eigen/LU>
#include <Eigen/SVD>

#include "drakeUtil.h"
#include "fastQP.h"

#define USE_EIQUADPROG_BACKUP 1

#if USE_EIQUADPROG_BACKUP
#include "eiquadprog.hpp"
#endif

#define _USE_MATH_DEFINES

#define MAX_CONSTRS 1000
#define MAX_ITER    20
using namespace std;

void angleDiff(VectorXd phi1, VectorXd phi2, VectorXd* d) {
  *d = phi2 - phi1;

  for (int i = 0; i < phi1.size(); i++) {
    if ((*d)(i) < -M_PI) {
      (*d)(i) = fmod((*d)(i) + M_PI, 2*M_PI) + M_PI;
    } else {
      (*d)(i) = fmod((*d)(i) + M_PI, 2*M_PI) - M_PI;
    }
  }
}

void angleDiff(MatrixXd phi1, MatrixXd phi2, MatrixXd* d) {
  *d = phi2 - phi1;

  for (int i = 0; i < phi1.rows(); i++) {
    for (int j = 0; j < phi1.cols(); j++) {
      if ((*d)(i, j) < -M_PI) {
        (*d)(i, j) = fmod((*d)(i, j) + M_PI, 2*M_PI) + M_PI;
      } else {
        (*d)(i, j) = fmod((*d)(i, j) + M_PI, 2*M_PI) - M_PI;
      }
    }
  }
}

using namespace std;

/**
 * Use Frank's fastQP code (mexed)
 * [q,info] = approximateIKEIQPmex(objgetMexModelPtr, q0, q_nom, Q, varargin)
 * info = 0 on success, 1 on failure
 **/
void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {
  if (nrhs < 4) {
    mexErrMsgIdAndTxt("Drake:approximateIKmex:NotEnoughInputs", "Usage approximateIKmex(model_ptr,q0,q_nom,Q,...)");
  }

  if (nlhs<1) return;

  // first get the model_ptr back from matlab
  RigidBodyManipulator *model= (RigidBodyManipulator*) getDrakeMexPointer(prhs[0]);

  int i, j, error, nq = model->num_positions;

  static RigidBodyManipulator* lastModel=NULL;
  static int lastNumJointLimits = 0;

  int equality_ind = 0;
  int inequality_ind = 0;

  //Constraint preallocation
  Matrix<double,-1,1,0,MAX_CONSTRS,1> beq(MAX_CONSTRS);
  Matrix<double,-1,1,0,MAX_CONSTRS,1> bin(MAX_CONSTRS);
  Matrix<double,-1,-1,RowMajor,MAX_CONSTRS,-1> Aeq(MAX_CONSTRS,nq);
  Matrix<double,-1,-1,RowMajor,MAX_CONSTRS,-1> Ain(MAX_CONSTRS,nq);

  //Add joint limits
  if (lastModel != model || true) {
    lastModel = model;
    for (i=0;i<nq;i++) {
      if(!mxIsInf(model->joint_limit_max[i])) {
        bin(inequality_ind) = model->joint_limit_max[i];
        Ain.row(inequality_ind) = MatrixXd::Zero(1,nq);
        Ain(inequality_ind,i) = 1;

//         cout << bin(inequality_ind) << endl << endl;
        inequality_ind++;
      }

      if(!mxIsInf(model->joint_limit_min[i])) {
        bin(inequality_ind) = -model->joint_limit_min[i];
        Ain.row(inequality_ind) << MatrixXd::Zero(1,nq);
        Ain(inequality_ind,i) = -1;
//         cout << bin(inequality_ind) << endl << endl;
        inequality_ind++;
      }
    }
    lastNumJointLimits = inequality_ind;
  } else {
    inequality_ind = lastNumJointLimits;
  }

  // cost:  (q-q_nom)'*Q*(q-q_nom)  \equiv q'*Q*q - 2*q_nom'*Q*q  (const term doesn't matter)
  Map<VectorXd> q_nom(mxGetPrSafe(prhs[2]), nq);
  Map<MatrixXd> Q(mxGetPrSafe(prhs[3]), nq, nq);
  VectorXd c = -Q*q_nom;

  double *q0 = mxGetPrSafe(prhs[1]);
  model->doKinematics(q0);

  VectorXd q0vec = Map<VectorXd>(q0, nq);

  i=4;
  while (i<nrhs) {
  	MatrixXd body_pos;
    Map<MatrixXd>* world_pos = NULL;
    int rows;
    int body_ind = mxGetScalar(prhs[i]) - 1;
    mxArray *min = NULL;
    mxArray *max = NULL;
    int n_pts;
    if (body_ind==-1) {
      int n_pts = mxGetN(prhs[i+1]);
      body_pos.resize(4, 1);
      body_pos << 0, 0, 0, 1;
      world_pos = new Map<MatrixXd>(mxGetPrSafe(prhs[i+1]), 3, n_pts);

      rows = 3;
      i+=2;
    } else {
      if (mxIsClass(prhs[i+2], "struct"))  {//isstruct(worldpos)
        min = mxGetField(prhs[i+2], 0, "min");
        max = mxGetField(prhs[i+2], 0, "max");

        if (min == NULL || max == NULL) {
          mexErrMsgIdAndTxt("Drake:approximateIKMex:BadInputs", "if world_pos is a struct, it must have fields .min and .max");
        }

        rows = mxGetM(min);

        if (rows != 3 && rows != 6) {
          mexErrMsgIdAndTxt("Drake:approximateIKMex:BadInputs", "world_pos.min must have 3 or 6 rows");
        }

        if (mxGetM(max) != rows) {
          mexErrMsgIdAndTxt("Drake:approximateIKMex:BadInputs", "world_pos.max must have the same number of rows as world_pos.min");
        }
      } else {
        rows = mxGetM(prhs[i+2]);
        int n_pts = mxGetN(prhs[i+2]);
        world_pos = new Map<MatrixXd>(mxGetPrSafe(prhs[i+2]), rows, n_pts);
      }

      if (mxIsClass(prhs[i+1], "char") || mxGetM(prhs[i+1]) == 1) {
        mexErrMsgIdAndTxt("Drake:approximateIKMex:NotImplemented", "collision group not implemented in mex (mposa)");
      } else {
        if (mxGetM(prhs[i+1]) !=3) {
          mexErrMsgIdAndTxt("Drake:approximateIKMex:BadInputs", "bodypos must be 3xmi");
        }
        n_pts = mxGetN(prhs[i+1]);

        Map<MatrixXd> pts_tmp(mxGetPrSafe(prhs[i+1]), 3, n_pts);
        body_pos.resize(4, n_pts);
        body_pos << pts_tmp, MatrixXd::Ones(1, n_pts);
      }
      i += 3;
    }

    MatrixXd x;
    MatrixXd J;
    if (body_ind == -1) {
      x = VectorXd::Zero(3, 1);
      Vector3d x_com;

      model->getCOM(x_com);
      model->getCOMJac(J);
      x.resize(3, 1);
      x << x_com;
    } else {

      J.resize(rows*n_pts, nq);

      model->forwardKin(body_ind, body_pos, (int) rows==6, x);
      model->forwardJac(body_ind, body_pos, (int) rows==6, J);

      if (rows == 6 && min == NULL && max == NULL) {
        VectorXd delta;
        angleDiff(x.block(3, 0, 3, n_pts), (*world_pos).block(3, 0, 3, n_pts), &delta);
        (*world_pos).block(3,0,3,n_pts) = x.block(3,0,3,n_pts)+delta;
      }
    }

    if (max != NULL) {
      double *val = mxGetPrSafe(max);

      // add inequality constraint
      for (j = 0; j < rows; j++) {
        if (!mxIsNaN(val[j])) {
          VectorXd rowVec = J.row(j);
          double rhs = val[j] - x(j) + J.row(j)*q0vec;
          if (mxIsInf(rhs)) {
            if (rhs < 0) {
              mexErrMsgIdAndTxt("Drake:approximateIKMex:BadInputs", "RHS of a constraint evaluates to -inf");
            }
          } else {
            Ain.row(inequality_ind) << J.row(j);
            bin(inequality_ind) = rhs;
            inequality_ind++;
          }
        }
      }
    }

    if (min != NULL) {
      double *val = mxGetPrSafe(min);
      // add inequality constraint
      for (j = 0; j < rows; j++) {
        if (!mxIsNaN(val[j])) {
          VectorXd rowVec = J.row(j);
          double rhs = -(val[j] - x(j) + J.row(j)*q0vec);
          if (mxIsInf(rhs)) {
            if (rhs < 0) {
              mexErrMsgIdAndTxt("Drake:approximateIKMex:BadInputs", "RHS of a constraint evaluates to -inf");
            }
          } else {
            Ain.row(inequality_ind) << J.row(j);
            bin(inequality_ind) = rhs;
            inequality_ind++;
          }
        }
      }
    }

    if (min == NULL && max == NULL) {
      for (j = 0; j < rows; j++) {
        if (!mxIsNaN((*world_pos)(j))) {
          double rhs =  (*world_pos)(j) - x(j) + J.row(j)*q0vec;

          Aeq.row(equality_ind) << J.row(j);
          beq(equality_ind) = rhs;
          equality_ind++;
        }
      }
    }
    delete world_pos;
    if (min) mxDestroyArray(min);
    if (max) mxDestroyArray(max);
  }

  Aeq.conservativeResize(equality_ind, nq);
  Ain.conservativeResize(inequality_ind, nq);
  beq.conservativeResize(equality_ind);
  bin.conservativeResize(inequality_ind);

  cout << "c is " << c.rows() << endl;
  cout << "Aeq is " << Aeq.rows() << " by " << Aeq.cols() << endl;

  VectorXd q = VectorXd::Zero(nq);
//   double result = solve_quadprog(Q, c, -Aeq, beq, -Ain, bin, q);

  VectorXd Qdiag = Q.diagonal();
  vector< MatrixBase<VectorXd > > blkQ;  blkQ.push_back(Qdiag);
  set<int> active;
  double result = fastQP(blkQ, c, Aeq, beq, Ain, bin, active, q);

  //Enable the block below to use stephens' qp code
  #if USE_EIQUADPROG_BACKUP
  if(result == 1) {
    Q += 1e-8*MatrixXd::Identity(nq,nq);
    result = solve_quadprog(Q, c, -Aeq.transpose(), beq, -Ain.transpose(), bin, q);

    if (mxIsInf(result)) {
      result = 1;
    } else {
      result = 0;
    }
  }
  #endif

  plhs[0] = mxCreateDoubleMatrix(nq, 1, mxREAL);
  memcpy(mxGetPrSafe(plhs[0]), q.data(), sizeof(double)*nq);

  if (nlhs>1) {
    plhs[1] = mxCreateDoubleScalar(result);
  }

  return;
}


/*
  [auto_generated]
  boost/numeric/odeint/stepper/extrapolation_stepper.hpp

  [begin_description]
  extrapolation stepper
  [end_description]

  Copyright 2009-2015 Mario Mulansky

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE_1_0.txt or
  copy at http://www.boost.org/LICENSE_1_0.txt)
*/

#ifndef BOOST_NUMERIC_ODEINT_STEPPER_EXTRAPOLATION_STEPPER_HPP_INCLUDED
#define BOOST_NUMERIC_ODEINT_STEPPER_EXTRAPOLATION_STEPPER_HPP_INCLUDED

#include <iostream>

#include <algorithm>

#include <boost/config.hpp> // for min/max guidelines
#include <boost/static_assert.hpp>

#include <boost/numeric/odeint/util/bind.hpp>
#include <boost/numeric/odeint/util/unwrap_reference.hpp>

#include <boost/numeric/odeint/stepper/base/explicit_error_stepper_base.hpp>
#include <boost/numeric/odeint/stepper/modified_midpoint.hpp>
#include <boost/numeric/odeint/stepper/controlled_step_result.hpp>
#include <boost/numeric/odeint/algebra/range_algebra.hpp>
#include <boost/numeric/odeint/algebra/default_operations.hpp>
#include <boost/numeric/odeint/algebra/algebra_dispatcher.hpp>
#include <boost/numeric/odeint/algebra/operations_dispatcher.hpp>

#include <boost/numeric/odeint/util/state_wrapper.hpp>
#include <boost/numeric/odeint/util/is_resizeable.hpp>
#include <boost/numeric/odeint/util/resizer.hpp>
#include <boost/numeric/odeint/util/unit_helper.hpp>
#include <boost/numeric/odeint/util/detail/less_with_sign.hpp>

namespace boost
{
namespace numeric
{
namespace odeint
{

template < unsigned short Order, class State, class Value = double,
           class Deriv = State, class Time = Value,
           class Algebra = typename algebra_dispatcher< State >::algebra_type,
           class Operations =
               typename operations_dispatcher< State >::operations_type,
           class Resizer = initially_resizer >
#ifndef DOXYGEN_SKIP
class extrapolation_stepper
    : public explicit_error_stepper_base<
          extrapolation_stepper< Order, State, Value, Deriv, Time, Algebra,
                                 Operations, Resizer >,
          Order, Order, Order - 2, State, Value, Deriv, Time, Algebra,
          Operations, Resizer >
#else
class extrapolation_stepper : public explicit_error_stepper_base
#endif
{

  private:
    // check for Order being odd
    BOOST_STATIC_ASSERT_MSG(
        ( ( Order % 2 ) == 0 ) && ( Order > 2 ),
        "extrapolation_stepper requires even Order larger than 2" );

  public:
#ifndef DOXYGEN_SKIP
    typedef explicit_error_stepper_base<
        extrapolation_stepper< Order, State, Value, Deriv, Time, Algebra,
                               Operations, Resizer >,
        Order, Order, Order - 2, State, Value, Deriv, Time, Algebra, Operations,
        Resizer > stepper_base_type;
#else
    typedef explicit_error_stepper_base< extrapolation_stepper< ... >, ... >
    stepper_base_type;
#endif

    typedef typename stepper_base_type::state_type state_type;
    typedef typename stepper_base_type::value_type value_type;
    typedef typename stepper_base_type::deriv_type deriv_type;
    typedef typename stepper_base_type::time_type time_type;
    typedef typename stepper_base_type::algebra_type algebra_type;
    typedef typename stepper_base_type::operations_type operations_type;
    typedef typename stepper_base_type::resizer_type resizer_type;

#ifndef DOXYGEN_SKIP
    typedef typename stepper_base_type::stepper_type stepper_type;
    typedef typename stepper_base_type::wrapped_state_type wrapped_state_type;
    typedef typename stepper_base_type::wrapped_deriv_type wrapped_deriv_type;

    typedef std::vector< value_type > value_vector;
    typedef std::vector< value_vector > value_matrix;
    typedef std::vector< size_t > int_vector;
    typedef std::vector< wrapped_state_type > state_table_type;
    typedef modified_midpoint< state_type, value_type, deriv_type, time_type,
                               algebra_type, operations_type,
                               resizer_type > midpoint_stepper_type;

#endif // DOXYGEN_SKIP

    typedef unsigned short order_type;
    static const order_type order_value = stepper_base_type::order_value;
    static const order_type stepper_order_value =
        stepper_base_type::stepper_order_value;
    static const order_type error_order_value =
        stepper_base_type::error_order_value;

    const static size_t m_k_max = ( order_value - 2 ) / 2;

    extrapolation_stepper( const algebra_type &algebra = algebra_type() )
        : stepper_base_type( algebra ), m_interval_sequence( m_k_max + 1 ),
          m_coeff( m_k_max + 1 ), m_table( m_k_max )
    {
        for ( unsigned short i = 0; i < m_k_max + 1; i++ )
        {
            m_interval_sequence[i] = 2 * ( i + 1 );
            m_coeff[i].resize( i );
            for ( size_t k = 0; k < i; ++k )
            {
                const value_type r =
                    static_cast< value_type >( m_interval_sequence[i] ) /
                    static_cast< value_type >( m_interval_sequence[k] );
                m_coeff[i][k] =
                    static_cast< value_type >( 1 ) /
                    ( r * r - static_cast< value_type >(
                                  1 ) ); // coefficients for extrapolation
            }
        }
    }

    template < class System, class StateIn, class DerivIn, class StateOut,
               class Err >
    void do_step_impl( System system, const StateIn &in, const DerivIn &dxdt,
                       time_type t, StateOut &out, time_type dt, Err &xerr )
    {
        // std::cout << "dt: " << dt << std::endl;
        // normal step
        do_step_impl( system, in, dxdt, t, out, dt );

        static const value_type val1( 1.0 );
        // additionally, perform the error calculation
        stepper_base_type::m_algebra.for_each3(
            xerr, out, m_table[0].m_v,
            typename operations_type::template scale_sum2<
                value_type, value_type >( val1, -val1 ) );
    }

    template < class System, class StateInOut, class DerivIn, class Err >
    void do_step_impl_io( System system, StateInOut &inout, const DerivIn &dxdt,
                          time_type t, time_type dt, Err &xerr )
    {
        // normal step
        do_step_impl_io( system, inout, dxdt, t, dt );

        static const value_type val1( 1.0 );
        // additionally, perform the error calculation
        stepper_base_type::m_algebra.for_each3(
            xerr, inout, m_table[0].m_v,
            typename operations_type::template scale_sum2<
                value_type, value_type >( val1, -val1 ) );
    }

    template < class System, class StateIn, class DerivIn, class StateOut >
    void do_step_impl( System system, const StateIn &in, const DerivIn &dxdt,
                       time_type t, StateOut &out, time_type dt )
    {
        m_resizer.adjust_size(
            in, detail::bind( &stepper_type::template resize_impl< StateIn >,
                              detail::ref( *this ), detail::_1 ) );
        size_t k = 0;
        m_midpoint.set_steps( m_interval_sequence[k] );
        m_midpoint.do_step( system, in, dxdt, t, out, dt );
        for ( k = 1; k <= m_k_max; ++k )
        {
            m_midpoint.set_steps( m_interval_sequence[k] );
            m_midpoint.do_step( system, in, dxdt, t, m_table[k - 1].m_v, dt );
            extrapolate( k, m_table, m_coeff, out );
        }
    }

    template < class System, class StateInOut, class DerivIn >
    void do_step_impl_io( System system, StateInOut &inout, const DerivIn &dxdt,
                          time_type t, time_type dt )
    {
        // special care for inout
        m_xout_resizer.adjust_size(
            inout,
            detail::bind( &stepper_type::template resize_m_xout< StateInOut >,
                          detail::ref( *this ), detail::_1 ) );
        do_step_impl( system, inout, dxdt, t, m_xout.m_v, dt );
        boost::numeric::odeint::copy( m_xout.m_v, inout );
    }

    template < class System, class StateInOut, class DerivIn >
    void do_step_dxdt_impl( System system, StateInOut &x, const DerivIn &dxdt,
                            time_type t, time_type dt )
    {
        do_step_impl_io( system , x , dxdt , t , dt );
    }

    template < class System, class StateIn, class DerivIn, class StateOut >
    void do_step_dxdt_impl( System system, const StateIn &in,
                            const DerivIn &dxdt, time_type t, StateOut &out,
                            time_type dt )
    {
        do_step_impl( system , in , dxdt , t , out , dt );
    }

    template < class StateIn > void adjust_size( const StateIn &x )
    {
        resize_impl( x );
        m_midpoint.adjust_size( x );
    }

  private:
    template < class StateIn > bool resize_impl( const StateIn &x )
    {
        bool resized( false );
        for ( size_t i = 0; i < m_k_max; ++i )
            resized |= adjust_size_by_resizeability(
                m_table[i], x, typename is_resizeable< state_type >::type() );
        return resized;
    }

    template < class StateIn > bool resize_m_xout( const StateIn &x )
    {
        return adjust_size_by_resizeability(
            m_xout, x, typename is_resizeable< state_type >::type() );
    }

    template < class StateInOut >
    void extrapolate( size_t k, state_table_type &table,
                      const value_matrix &coeff, StateInOut &xest )
    /* polynomial extrapolation, see http://www.nr.com/webnotes/nr3web21.pdf
       uses the obtained intermediate results to extrapolate to dt->0
    */
    {
        static const value_type val1 = static_cast< value_type >( 1.0 );

        for ( int j = k - 1; j > 0; --j )
        {
            stepper_base_type::m_algebra.for_each3(
                table[j - 1].m_v, table[j].m_v, table[j - 1].m_v,
                typename operations_type::template scale_sum2<
                    value_type, value_type >( val1 + coeff[k][j],
                                              -coeff[k][j] ) );
        }
        stepper_base_type::m_algebra.for_each3(
            xest, table[0].m_v, xest,
            typename operations_type::template scale_sum2<
                value_type, value_type >( val1 + coeff[k][0], -coeff[k][0] ) );
    }

  private:
    midpoint_stepper_type m_midpoint;

    resizer_type m_resizer;
    resizer_type m_xout_resizer;

    int_vector m_interval_sequence; // stores the successive interval counts
    value_matrix m_coeff;

    wrapped_state_type m_xout;
    state_table_type m_table; // sequence of states for extrapolation
};

/******** DOXYGEN *******/

/**
 * \class extrapolation_stepper
 * \brief Extrapolation stepper with configurable order, and error estimation.
 *
 * The extrapolation stepper is a stepper with error estimation and configurable
 * order. The order is given as template parameter and needs to be an _odd_
 * number. The stepper is based on several executions of the modified midpoint
 * method and a Richardson extrapolation. This is essentially the same technique
 * as for bulirsch_stoer, but without the variable order.
 *
 * \note The Order parameter has to be an even number greater 2.
 */
}
}
}
#endif


/*
 Copyright (c) 2015 Siyu Lei, Silviu Maniu, Luyi Mo (University of Hong Kong)

 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
 in the Software without restriction, including without limitation the rights
 to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 copies of the Software, and to permit persons to whom the Software is
 furnished to do so, subject to the following conditions:

 The above copyright notice and this permission notice shall be included in
 all copies or substantial portions of the Software.

 THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 */

#ifndef __oim__DiscountDegreeEvaluator__
#define __oim__DiscountDegreeEvaluator__

#include "common.hpp"
#include "Evaluator.hpp"

#include <boost/heap/fibonacci_heap.hpp>

class DiscountDegreeEvaluator : public Evaluator {
 public:
  std::unordered_set<unode_int> select(
        const Graph& graph, Sampler& sampler,
        const std::unordered_set<unode_int>& activated, unsigned int k) {
    std::unordered_set<unode_int> set;
    unsigned int type = sampler.get_type();
    boost::heap::fibonacci_heap<NodeType> queue;
    std::unordered_map<unode_int,
        boost::heap::fibonacci_heap<NodeType>::handle_type> queue_nodes;
    for (unode_int node : graph.get_nodes()) {
      NodeType nstruct;
        nstruct.id = node;
        nstruct.deg = activated.find(node) == activated.end() ? 1.0f : 0.0f;
        if (graph.has_neighbours(node)) {
          for (auto edge : graph.get_neighbours(node)) {
            if(activated.find(edge.target) == activated.end()) {
              nstruct.deg += edge.dist->sample(type);
            }
          }
        }
        queue_nodes[node] = queue.push(nstruct);
    }
    while (set.size() < k && (!queue.empty())) {
      NodeType nstruct = queue.top();
      set.insert(nstruct.id);
      for (auto edge : graph.get_neighbours(nstruct.id)) {
        if (activated.find(edge.target) == activated.end() &&
            set.find(edge.target) == set.end()) {
          NodeType newnstruct = *queue_nodes[edge.target];
          newnstruct.id = edge.target;
          newnstruct.deg = newnstruct.deg*(1.0f - edge.dist->sample(type));
          queue.update(queue_nodes[edge.target], newnstruct);
        }
      }
      queue.pop();
    }
    queue_nodes.clear();
    return set;
  }
};

#endif /* defined(__oim__DiscountDegreeEvaluator__) */


//
// Created by yche on 12/23/17.
//

#include <boost/program_options.hpp>

#include "../yche_refactor/flpmc_yche.h"
#include "../yche_refactor/simrank.h"

using namespace std;
using namespace boost::program_options;

void test_FLPMC(string data_name, double c, double epsilon, double delta, int x, int y) {
    string path = get_edge_list_path(data_name);
    GraphYche g(path);
    NodePair q{x, y};

    // 1st: init
    FLPMC flpmc(data_name, g, c, epsilon, delta, 100);

    // 2nd: query
    TruthSim ts(data_name, g, c, epsilon);
    auto max_err = 0.0;
    auto min_err = 1.0;
    auto failure_count = 0;

    auto start = std::chrono::high_resolution_clock::now();

    for (auto i = 0; i < 2000; i++) {
        double result = flpmc.query_one2one(q);
        auto cur_err = abs(result - ts.sim(x, y));
        max_err = max(max_err, cur_err);
        min_err = min(min_err, cur_err);
        if (cur_err > 0.01) {
            cout << result << " ," << ts.sim(x, y) << endl;
            failure_count++;
        }
    }
    auto end = std::chrono::high_resolution_clock::now();
    std::chrono::duration<double> elapsed = end - start;

    cout << "failure count:" << failure_count << endl;
    cout << "max err:" << max_err << endl;
    cout << "min err:" << min_err << endl;
}

int main(int args, char *argv[]) {
    try {
        options_description desc{"Options"};
        double c = 0.6;
        double epsilon = 0.01;
        double delta = 0.01;
        string data_name;
        int x, y;
        desc.add_options()
                ("help,h", "Help Screen")
                ("DataName,d", value<string>(&data_name), "Graph Name")
                ("decay,c", value<double>(&c), "Decay Facrot c")
                ("x,x", value<int>(&x), "x: the query node")
                ("y,y", value<int>(&y), "y: the query node")
                ("epsilon,e", value<double>(&epsilon), "Error bound");
        variables_map vm;
        store(parse_command_line(args, argv, desc), vm);
        notify(vm);

        if (vm.count("help")) {
            cout << desc << endl;
        } else {
            test_FLPMC(data_name, c, epsilon, delta, x, y);
        }

    }
    catch (const error &ex) {
        cerr << ex.what() << endl;
    }
    return 0;
}


#ifndef CAFFE_RELU_LAYER_HPP_
#define CAFFE_RELU_LAYER_HPP_

#include <vector>

#include <boost/random.hpp>
#include <boost/random/uniform_int_distribution.hpp>
#include <boost/math/distributions/normal.hpp>

#include "caffe/blob.hpp"
#include "caffe/layer.hpp"
#include "caffe/proto/caffe.pb.h"

#include "caffe/layers/neuron_layer.hpp"

namespace caffe {

/**
 * @brief Perturbs incoming values by a random amount
 */
template <typename Dtype> class RandomLayer : public NeuronLayer<Dtype> {
public:
  explicit RandomLayer(const LayerParameter &param)
      : NeuronLayer<Dtype>(param) {}

  virtual void LayerSetUp(const vector<Blob<Dtype> *> &bottom,
                          const vector<Blob<Dtype> *> &top);

  virtual inline const char *type() const { return "Random"; }

protected:
  boost::mt19937 rng;
  boost::random::uniform_int_distribution<> dist;

  virtual void Forward_cpu(const vector<Blob<Dtype> *> &bottom,
                           const vector<Blob<Dtype> *> &top);

  virtual void Backward_cpu(const vector<Blob<Dtype> *> &top,
                            const vector<bool> &propagate_down,
                            const vector<Blob<Dtype> *> &bottom);
};

} // namespace caffe

#endif // CAFFE_RELU_LAYER_HPP_


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2007-2015 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2015 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2015 Mateusz Loskot, London, UK.
// Copyright (c) 2013-2019 Adam Wulkiewicz, Lodz, Poland.

// This file was modified by Oracle on 2020.
// Modifications copyright (c) 2020 Oracle and/or its affiliates.
// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

// Parts of Boost.Geometry are redesigned from Geodan's Geographic Library
// (geolib/GGL), copyright (c) 1995-2010 Geodan, Amsterdam, the Netherlands.

// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_GEOMETRY_EXTENSIONS_NSPHERE_ALGORITHMS_WITHIN_HPP
#define BOOST_GEOMETRY_EXTENSIONS_NSPHERE_ALGORITHMS_WITHIN_HPP

#include <boost/core/ignore_unused.hpp>

#include <boost/geometry/algorithms/distance.hpp>
#include <boost/geometry/algorithms/make.hpp>
#include <boost/geometry/algorithms/within.hpp>
#include <boost/geometry/strategies/distance.hpp>

#include <boost/geometry/core/tags.hpp>

#include <boost/geometry/extensions/nsphere/core/access.hpp>
#include <boost/geometry/extensions/nsphere/core/radius.hpp>
#include <boost/geometry/extensions/nsphere/core/tags.hpp>
#include <boost/geometry/extensions/nsphere/algorithms/assign.hpp>

#include <boost/geometry/extensions/nsphere/strategies/relate/cartesian.hpp>

namespace boost { namespace geometry
{

#ifndef DOXYGEN_NO_DETAIL
namespace detail { namespace within
{

// TODO: This is not fully correct for all geometries
//   For more complex geometries than Point if at least one point is in the
//   interior of nsphere within should return true even if other points are
//   on the boudary

template<typename B, typename C, typename S>
inline bool box_in_circle(B const& b, C const& c, S const& strategy)
{
    typedef typename point_type<B>::type point_type;

    // Currently only implemented for 2d geometries
    assert_dimension<point_type, 2>();
    assert_dimension<C, 2>();

    // Box: all four points must lie within circle
    point_type const p0 = geometry::make<point_type>(get<min_corner, 0>(b), get<min_corner, 1>(b));
    point_type const p1 = geometry::make<point_type>(get<max_corner, 0>(b), get<max_corner, 1>(b));
    point_type const p2 = geometry::make<point_type>(get<min_corner, 0>(b), get<max_corner, 1>(b));
    point_type const p3 = geometry::make<point_type>(get<max_corner, 0>(b), get<min_corner, 1>(b));

    auto const s = strategy.within(p0, c);

    // Check points lower-left and upper-right, then lower-right and upper-left
    return s.apply(p0, c) && s.apply(p1, c) && s.apply(p2, c) && s.apply(p3, c);
}

// Generic "range-in-circle", true if all points within circle
template<typename R, typename C, typename S>
inline bool range_in_circle(R const& range, C const& c, S const& s)
{
    ::boost::ignore_unused(s);

    for (typename boost::range_iterator<R const>::type it = boost::begin(range);
         it != boost::end(range); ++it)
    {
        auto const& p = *it;

        // TODO: This is not fully correct since some of the points can lie on the boundary
        if (! s.within(p, c).apply(p, c))
        {
            return false;
        }
    }

    return true;
}

template<typename Y, typename C, typename S>
inline bool polygon_in_circle(Y const& poly, C const& c, S const& s)
{
    return range_in_circle(exterior_ring(poly), c, s);
}

template<typename I, typename C, typename S>
inline bool multi_polygon_in_circle(I const& m, C const& c, S const& s)
{
    for (typename I::const_iterator i = m.begin(); i != m.end(); i++)
    {
        if (! polygon_in_circle(*i, c, s))
        {
            return false;
        }
    }
    return true;
}

}} // namespace detail::within
#endif // DOXYGEN_NO_DETAIL

#ifndef DOXYGEN_NO_DISPATCH
namespace dispatch
{

template <typename P, typename Circle>
struct within<P, Circle, point_tag, nsphere_tag>
{
    template <typename Strategy>
    static inline bool apply(P const& p, Circle const& c, Strategy const& strategy)
    {
        return strategy.within(p, c).apply(p, c);
    }
};

template <typename Box, typename Circle>
struct within<Box, Circle, box_tag, nsphere_tag>
{
    template <typename Strategy>
    static inline bool apply(Box const& b, Circle const& c, Strategy const& s)
    {
        return detail::within::box_in_circle(b, c, s);
    }
};

template <typename Linestring, typename Circle>
struct within<Linestring, Circle, linestring_tag, nsphere_tag>
{
    template <typename Strategy>
    static inline bool apply(Linestring const& ln, Circle const& c, Strategy const& s)
    {
        return detail::within::range_in_circle(ln, c, s);
    }
};

template <typename Ring, typename Circle>
struct within<Ring, Circle, ring_tag, nsphere_tag>
{
    template <typename Strategy>
    static inline bool apply(Ring const& r, Circle const& c, Strategy const& s)
    {
        return detail::within::range_in_circle(r, c, s);
    }
};

template <typename Polygon, typename Circle>
struct within<Polygon, Circle, polygon_tag, nsphere_tag>
{
    template <typename Strategy>
    static inline bool apply(Polygon const& poly, Circle const& c, Strategy const& s)
    {
        return detail::within::polygon_in_circle(poly, c, s);
    }
};

template <typename M, typename C>
struct within<M, C, multi_polygon_tag, nsphere_tag>
{
    template <typename Strategy>
    static inline bool apply(M const& m, C const& c, Strategy const& s)
    {
        return detail::within::multi_polygon_in_circle(m, c, s);
    }
};

template <typename NSphere, typename Box>
struct within<NSphere, Box, nsphere_tag, box_tag>
{
    template <typename Strategy>
    static inline bool apply(NSphere const& nsphere, Box const& box, Strategy const& strategy)
    {
        assert_dimension_equal<NSphere, Box>();

        return strategy.within(nsphere, box).apply(nsphere, box);
    }
};

} // namespace dispatch
#endif // DOXYGEN_NO_DISPATCH

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_EXTENSIONS_NSPHERE_ALGORITHMS_WITHIN_HPP


/*! \file hydrogen_fem_cuda_main.cpp
    \brief FEMで水素原子に対するSchrödinger方程式を解く

    Copyright © 2019 @dc1394 All Rights Reserved.
    This software is released under the BSD 2-Clause License.
*/

#include "hydrogen_fem_cuda.h"
#include <chrono>               // for std::chrono
#include <iostream>             // for std::cout
#include <boost/format.hpp>     // for boost::format

int main()
{
    using namespace std::chrono;

    auto const start = system_clock::now();
    hydrogen_fem::Hydrogen_FEM hl;

    std::cout << boost::format("計算が終わりました: 基底状態のエネルギー固有値E = %.14f (Hartree)\n") % hl.do_run();
    auto const end = system_clock::now();
    std::cout << boost::format("計算時間 = %.14f（秒）\n") % duration_cast< duration<double> >(end - start).count();

    hl.save_result();

    std::cout << "計算結果を" << hydrogen_fem::Hydrogen_FEM::EIGENFUNC_FILENAME << "と" << hydrogen_fem::Hydrogen_FEM::EIGENVAL_FILENAME << "に書き込みました" << std::endl;

    return 0;
}