/**
 All of this code is written by Aman Agrawal
 (Indian Institute of Technology, Delhi)
*/

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

#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

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;

clock_t total_begin = clock();

genotype g;
MatrixXdr geno_matrix; //(p,n)

int MAX_ITER;
int k,p,n;
int k_orig;

MatrixXdr c; //(p,k)
MatrixXdr x; //(k,n)
MatrixXdr v; //(p,k)
MatrixXdr means; //(p,1)
MatrixXdr stds; //(p,1)

options command_line_opts;

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

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();
	}

	#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;

	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);
	}
}

pair<double,double> get_error_norm(MatrixXdr &c){
	HouseholderQR<MatrixXdr> qr(c);
	MatrixXdr Q;
	Q = qr.householderQ() * MatrixXdr::Identity(p,k);
	MatrixXdr q_t(k,p);
	q_t = Q.transpose();
	MatrixXdr b(k,n);
	multiply_y_post(q_t,k,b,true);
	JacobiSVD<MatrixXdr> b_svd(b, ComputeThinU | ComputeThinV);
	MatrixXdr u_l,d_l,v_l;
	if(fast_mode)
        u_l = b_svd.matrixU();
    else
        u_l = Q * b_svd.matrixU();
	v_l = b_svd.matrixV();
	d_l = MatrixXdr::Zero(k,k);
	for(int kk=0;kk<k; kk++)
		d_l(kk,kk) = (b_svd.singularValues())(kk);

	MatrixXdr u_k,v_k,d_k;
	u_k = u_l.leftCols(k_orig);
	v_k = v_l.leftCols(k_orig);
	d_k = MatrixXdr::Zero(k_orig,k_orig);
	for(int kk =0 ; kk < k_orig ; kk++)
		d_k(kk,kk)  =(b_svd.singularValues())(kk);

	MatrixXdr b_l,b_k;
    b_l = u_l * d_l * (v_l.transpose());
    b_k = u_k * d_k * (v_k.transpose());

    if(fast_mode){
        double temp_k = b_k.cwiseProduct(b).sum();
        double temp_l = b_l.cwiseProduct(b).sum();
        double b_knorm = b_k.norm();
        double b_lnorm = b_l.norm();
        double norm_k = (b_knorm*b_knorm) - (2*temp_k);
        double norm_l = (b_lnorm*b_lnorm) - (2*temp_l);
        return make_pair(norm_k,norm_l);
    }
    else{
        MatrixXdr e_l(p,n);
        MatrixXdr e_k(p,n);
        for(int p_iter=0;p_iter<p;p_iter++){
            for(int n_iter=0;n_iter<n;n_iter++){
                e_l(p_iter,n_iter) = g.get_geno(p_iter,n_iter,var_normalize) - b_l(p_iter,n_iter);
                e_k(p_iter,n_iter) = g.get_geno(p_iter,n_iter,var_normalize) - b_k(p_iter,n_iter);
            }
        }

        double ek_norm = e_k.norm();
        double el_norm = e_l.norm();
        return make_pair(ek_norm,el_norm);
    }
}

MatrixXdr run_EM_not_missing(MatrixXdr &c_orig){

	#if DEBUG==1
		if(debug){
			print_time ();
			cout << "Enter: run_EM_not_missing" << endl;
		}
	#endif

	MatrixXdr c_temp(k,p);
	MatrixXdr c_new(p,k);
	c_temp = ( (c_orig.transpose()*c_orig).inverse() ) * (c_orig.transpose());

	#if DEBUG==1
		if(debug){
			print_timenl ();
		}
	#endif

	MatrixXdr x_fn(k,n);
	multiply_y_post(c_temp,k,x_fn,true);

	#if DEBUG==1
		if(debug){
			print_timenl ();
		}
	#endif

	MatrixXdr x_temp(n,k);
	x_temp = (x_fn.transpose()) * ((x_fn*(x_fn.transpose())).inverse());
	multiply_y_pre(x_temp,k,c_new,true);

	#if DEBUG==1
		if(debug){
			print_time ();
			cout << "Exiting: run_EM_not_missing" << endl;
		}
	#endif

	return c_new;
}

MatrixXdr run_EM_missing(MatrixXdr &c_orig){

	MatrixXdr c_new(p,k);

	MatrixXdr mu(k,n);

	// E step
	MatrixXdr c_temp(k,k);
	c_temp = c_orig.transpose() * c_orig;

	MatrixXdr T(k,n);
	MatrixXdr c_fn;
	c_fn = c_orig.transpose();
	multiply_y_post(c_fn,k,T,false);

	MatrixXdr M_temp(k,1);
	M_temp = c_orig.transpose() *  means;

	for(int j=0;j<n;j++){
		MatrixXdr D(k,k);
		MatrixXdr M_to_remove(k,1);
		D = MatrixXdr::Zero(k,k);
		M_to_remove = MatrixXdr::Zero(k,1);
		for(int i=0;i<g.not_O_j[j].size();i++){
			int idx = g.not_O_j[j][i];
			D = D + (c_orig.row(idx).transpose() * c_orig.row(idx));
			M_to_remove = M_to_remove + (c_orig.row(idx).transpose()*g.get_col_mean(idx));
		}
		mu.col(j) = (c_temp-D).inverse() * ( T.col(j) - M_temp + M_to_remove);
	}

	// M step

	MatrixXdr mu_temp(k,k);
	mu_temp = mu * mu.transpose();
	MatrixXdr T1(p,k);
	MatrixXdr mu_fn;
	mu_fn = mu.transpose();
	multiply_y_pre(mu_fn,k,T1,false);
	MatrixXdr mu_sum(k,1);
	mu_sum = MatrixXdr::Zero(k,1);
	mu_sum = mu.rowwise().sum();

	for(int i=0;i<p;i++){
		MatrixXdr D(k,k);
		MatrixXdr mu_to_remove(k,1);
		D = MatrixXdr::Zero(k,k);
		mu_to_remove = MatrixXdr::Zero(k,1);
		for(int j=0;j<g.not_O_i[i].size();j++){
			int idx = g.not_O_i[i][j];
			D = D + (mu.col(idx) * mu.col(idx).transpose());
			mu_to_remove = mu_to_remove + (mu.col(idx));
		}
		c_new.row(i) = (((mu_temp-D).inverse()) * (T1.row(i).transpose() -  ( g.get_col_mean(i) * (mu_sum-mu_to_remove)))).transpose();
		double mean;
		mean = g.get_col_sum(i);
		mean = mean -  (c_orig.row(i)*(mu_sum-mu_to_remove))(0,0);
		mean = mean * 1.0 / (n-g.not_O_i[i].size());
		g.update_col_mean(i,mean);
	}
	return c_new;
}

MatrixXdr run_EM(MatrixXdr &c_orig){

	if(missing)
		return run_EM_missing(c_orig);
	else
		return run_EM_not_missing(c_orig);
}

void print_vals(){

	HouseholderQR<MatrixXdr> qr(c);
	MatrixXdr Q;
	Q = qr.householderQ() * MatrixXdr::Identity(p,k);
	MatrixXdr q_t(k,p);
	q_t = Q.transpose();
	MatrixXdr b(k,n);
	multiply_y_post(q_t,k,b,true);
	JacobiSVD<MatrixXdr> b_svd(b, ComputeThinU | ComputeThinV);
	MatrixXdr u_l;
	u_l = b_svd.matrixU();
	MatrixXdr v_l;
	v_l = b_svd.matrixV();
	MatrixXdr u_k;
	MatrixXdr v_k,d_k;
	u_k = u_l.leftCols(k_orig);
	v_k = v_l.leftCols(k_orig);

	ofstream evec_file;
	evec_file.open((string(command_line_opts.OUTPUT_PATH)+string("evecs.txt")).c_str());
	evec_file<< std::setprecision(15) << Q*u_k << endl;
	evec_file.close();
	ofstream eval_file;
	eval_file.open((string(command_line_opts.OUTPUT_PATH)+string("evals.txt")).c_str());
	for(int kk =0 ; kk < k_orig ; kk++)
		eval_file << std::setprecision(15)<< (b_svd.singularValues())(kk)<<endl;
	eval_file.close();

	ofstream proj_file;
	proj_file.open((string(command_line_opts.OUTPUT_PATH) + string("projections.txt")).c_str());
	proj_file << std::setprecision(15)<< v_k<<endl;
	proj_file.close();
	if(debug){
		ofstream c_file;
		c_file.open((string(command_line_opts.OUTPUT_PATH)+string("cvals.txt")).c_str());
		c_file<<c<<endl;
		c_file.close();

		d_k = MatrixXdr::Zero(k_orig,k_orig);
		for(int kk =0 ; kk < k_orig ; kk++)
			d_k(kk,kk)  =(b_svd.singularValues())(kk);
		MatrixXdr x_k;
		x_k = d_k * (v_k.transpose());
		ofstream x_file;
		x_file.open((string(command_line_opts.OUTPUT_PATH) + string("xvals.txt")).c_str());
		x_file<<x_k.transpose()<<endl;
		x_file.close();
	}
}

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

	clock_t io_begin = clock();
    clock_gettime (CLOCK_REALTIME, &t0);

	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;

	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 =  command_line_opts.max_iterations ;
	k_orig = command_line_opts.num_of_evec ;
	debug = command_line_opts.debugmode ;
	check_accuracy = command_line_opts.getaccuracy;
	var_normalize = command_line_opts.var_normalize;
	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;
	convergence_limit = command_line_opts.convergence_limit;
	bool toStop=false;
	if(convergence_limit!=-1)
		toStop=true;
	srand((unsigned int) time(0));
	c.resize(p,k);
	x.resize(k,n);
	v.resize(p,k);
	means.resize(p,1);
	stds.resize(p,1);

	if(!fast_mode && !memory_efficient){
		geno_matrix.resize(p,n);
		g.generate_eigen_geno(geno_matrix,var_normalize);
	}

	clock_t io_end = clock();

	//TODO: Initialization of c with gaussian distribution
	c = MatrixXdr::Random(p,k);

	// Initial intermediate data structures
	blocksize = k;
	int hsegsize = g.segment_size_hori; 	// = log_3(n)
	int hsize = pow(3,hsegsize);
	int vsegsize = g.segment_size_ver; 		// = log_3(p)
	int vsize = pow(3,vsegsize);

	partialsums = new double [blocksize];
	sum_op = new double[blocksize];
	yint_e = new double [hsize*blocksize];
	yint_m = new double [hsize*blocksize];
	memset (yint_m, 0, hsize*blocksize * sizeof(double));
	memset (yint_e, 0, hsize*blocksize * sizeof(double));

	y_e  = new double*[g.Nindv];
	for (int i = 0 ; i < g.Nindv ; i++) {
		y_e[i] = new double[blocksize];
		memset (y_e[i], 0, blocksize * sizeof(double));
	}

	y_m = new double*[hsegsize];
	for (int i = 0 ; i < hsegsize ; i++)
		y_m[i] = new double[blocksize];

	for(int i=0;i<p;i++){
		means(i,0) = g.get_col_mean(i);
		stds(i,0) = g.get_col_std(i);
	}

	ofstream c_file;
	if(debug){
		c_file.open((string(command_line_opts.OUTPUT_PATH)+string("cvals_orig.txt")).c_str());
		c_file<<c<<endl;
		c_file.close();
		printf("Read Matrix\n");
	}

	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

	clock_t it_begin = clock();
	for(int i=0;i<MAX_ITER;i++){

		MatrixXdr c1,c2,cint,r,v;
		double a,nll;
		if(debug){
			print_time ();
			cout << "*********** Begin epoch " << i << "***********" << endl;
		}
		if(accelerated_em!=0){
			#if DEBUG==1
				if(debug){
					print_time();
					cout << "Before EM" << endl;
				}
			#endif
			c1 = run_EM(c);
			c2 = run_EM(c1);
			#if DEBUG==1
				if(debug){
					print_time();
					cout << "After EM but before acceleration" << endl;
				}
			#endif
			r = c1-c;
			v = (c2-c1) - r;
			a = -1.0 * r.norm() / (v.norm()) ;
			if(accelerated_em==1){
				if(a>-1){
					a=-1;
					cint=c2;
				}
				else {
					cint = c - 2*a*r + a*a*v;
					nll = get_error_norm(cint).second;
					if(i>0){
						while(nll>prevnll && a<-1){
							a = 0.5 * (a-1);
							cint = c - 2*a*r +(a*a*v);
							nll = get_error_norm(cint).second;
						}
					}
				}
				c = cint;
			}
			else if(accelerated_em==2){
				cint = c - 2*a*r + a*a*v;
				c = cint;
				// c = run_EM(cint);
			}
		}
		else{
			c = run_EM(c);
		}

		if ( accelerated_em == 1 || check_accuracy || toStop) {
			pair<double,double> e = get_error_norm(c);
			prevnll = e.second;
			if(check_accuracy)
				cout<<"Iteration "<<i+1<<"  "<<std::setprecision(15)<<e.first<<"  "<<e.second<<endl;
			if(abs(e.first-prev_error.first)<=convergence_limit){
				cout<<"Breaking after "<<i+1<<" iterations"<<endl;
				break;
			}
			prev_error = e;
		}
		if(debug){
			print_time ();
			cout << "*********** End epoch " << i << "***********" << endl;
		}

	}
	clock_t it_end = clock();

    print_vals();

	clock_t total_end = clock();
	double io_time = double(io_end - io_begin) / CLOCKS_PER_SEC;
	double avg_it_time = double(it_end - it_begin) / (MAX_ITER * 1.0 * CLOCKS_PER_SEC);
	double total_time = double(total_end - total_begin) / CLOCKS_PER_SEC;
	cout<<"IO Time:  "<< io_time << "\nAVG Iteration Time:  "<<avg_it_time<<"\nTotal runtime:   "<<total_time<<endl;

	delete[] sum_op;
	delete[] partialsums;
	delete[] yint_e;
	delete[] yint_m;

	for (int i  = 0 ; i < hsegsize; i++)
		delete[] y_m [i];
	delete[] y_m;

	for (int i  = 0 ; i < g.Nindv; i++)
		delete[] y_e[i];
	delete[] y_e;

	return 0;
}


/*======================================================================
Copyright 2019 Google LLC

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

    https://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 <cmath>

#include <Eigen/Dense>
#include "PCV_Types.h"

#include "Caster.h"

using namespace Eigen;

Caster::Caster(
            double _Kx, double _Ky, double _ang,
            double  _b, double  _r,
            double  _f, double _Mf, double _If,
            double _Ih, double _Ii, double _Is, double _It, double _Ij,
            double _Ns, double _Nt, double _Nw,
            double _px, double _py, double _Mp, double _Ip)
    :Lambda(M3Z), Mu(V3Z),
     Ce(M3I), Jdot(M3Z), A(M3Z), CC(V3Z), L_Pk(M3Z), Mu_Pk(V3Z)
{
 Kx = _Kx;   Ky = _Ky;
  b = _b;    bi= 1.0/b;   r = _r;    ri= 1.0/r;
  f = _f;    e = f-b;    Mf = _Mf;  If = _If;
 Ih = _Ih;  Ii = _Ii;    Is = _Is;  It = _It;    Ij = _Ij;
 Ns = _Ns;  Nt = _Nt;    Nw = _Nw;
 px = _px;  py = _py;    Mp = _Mp;  Ip = _Ip;

 Px = Kx + px;
 Py = Ky + py;
 P2  = pow(Px,2) + pow(Py,2);

 enc_offset = (long)(((M_PI_2 - _ang) * Ns * ENC_RAD2CNT)+0.5);

 Fill_A();
 Fill_L_Pk();
}

void
Caster::Fill_A()
{
  double mfe2Ifih= Mf*e*e + If + Ii + Ih;
  double IsNs    = Is * Ns;
  double ItNt    = It * Nt;
  double ItNt2   = ItNt*Nt;

  A(0,0) = mfe2Ifih + IsNs*Ns + ItNt;
  A(0,1) = (Ii - Ih - ItNt2) * Nw;
  A(0,2) = mfe2Ifih - IsNs - ItNt;

  A(1,0) = A(0,1);
  A(1,1) = Mf*r*r + (Ii + Ih + ItNt2)*Nw*Nw + Ij;
  A(1,2) = (Ii - Ih + ItNt ) * Nw;

  A(2,0) = A(0,2);
  A(2,1) = A(1,2);
  A(2,2) = mfe2Ifih + Is + It;
}

inline void
Caster::Fill_CC()
{
  register double  u_w = u + w;
  register double Mfre = Mf*r*e;

  CC(0) =  Mfre * v * u_w;
  CC(1) = -Mfre * pow(u_w,2);
  CC(2) =  CC[0];
}

void
Caster::Fill_L_Pk()
{
  L_Pk(0,0) =  Mp;
  L_Pk(0,1) =  0.0;
  L_Pk(0,2) = -Mp * Py;

  L_Pk(1,0) =  0.0;
  L_Pk(1,1) =  Mp;
  L_Pk(1,2) =  Mp * Px;

  L_Pk(2,0) = L_Pk(0,2);
  L_Pk(2,1) = L_Pk(1,2);
  L_Pk(2,2) = Ip + Mp*P2;
}

inline void
Caster::Fill_Mu_Pk()
{
  register double w2  = pow(w,2);

  Mu_Pk(0) = -Mp * Px * w2;
  Mu_Pk(1) = -Mp * Py * w2;
  Mu_Pk(2) =  0.0;
}

inline void
Caster::Fill_Ce()  // Ce (C_theta) (a.k.a. Ji )
{
  Ce(0,0) =  bi*s;
  Ce(0,1) = -bi*c;
  Ce(0,2) = -bi*(Kx*c + Ky*s) - 1.0;

  Ce(1,0) =  ri*c;
  Ce(1,1) =  ri*s;
  Ce(1,2) =  ri*(Kx*s - Ky*c);

  Ce(2,0) =  0.0;
  Ce(2,1) =  0.0;
  Ce(2,2) =  1.0;
}

inline void
Caster::Fill_Jdot()
{
  register double u_w = u + w;

  Jdot(0,0) =  b*c*u_w;
  Jdot(0,1) = -r*s*u_w;
  Jdot(0,2) =  Kx*w + b*c*u_w;

  Jdot(1,0) =  b*s*u_w;
  Jdot(1,1) =  r*c*u_w;
  Jdot(1,2) =  Ky*w + b*s*u_w;

  Jdot(2,0) =  0.0;
  Jdot(2,1) =  0.0;
  Jdot(2,2) =  0.0;
}

void
Caster::Fill_LM( double const _u,   // sigma dot
                 double const _v,   //  rho  dot
                 double const _w)   // theta dot
{

  static Vector3d   Qdot(3);

  // SET STATE
  Qdot[0] = u = _u;
  Qdot[1] = v = _v;
  Qdot[2] = w = _w;

  // FILL JACOBIANS
  Fill_Ce();
  Fill_Jdot();

  // FILL JT-SPACE CC VECTOR (NOTE: 'A' IS CONSTANT)
  Fill_CC();
  // FILL OP-SPACE CC VECTOR (NOTE: 'L_Pk' IS CONSTANT)
  Fill_Mu_Pk();

  //Lambda = L_Pk + (Cet * A * Ce)
  Lambda = L_Pk + Ce.transpose() * A * Ce;

  //Mu = Mu_Pk + [Cet * ( CC - (A*(Ce*(Jdot*Qdot))) )]
  Mu = Mu_Pk + Ce.transpose()*(CC-(A*(Ce*(Jdot*Qdot))));
}


#ifndef STAN_MATH_PRIM_FUN_LAMBERT_W_HPP
#define STAN_MATH_PRIM_FUN_LAMBERT_W_HPP

#include <stan/math/prim/meta.hpp>
#include <stan/math/prim/fun/boost_policy.hpp>
#include <boost/math/special_functions/lambert_w.hpp>

namespace stan {
namespace math {

/**
 * Compute the Lambert W function on W0 branch for a value x.
 *
 * @tparam T type of value
 * @param x value
 * @return value of the W0 branch of the Lambert W function for x
 * @throw std::domain_error if x is less than or equal to `-e^(-1)`
 */
template <typename T, require_arithmetic_t<T>* = nullptr>
inline double lambert_w0(const T& x) {
  return boost::math::lambert_w0(x, boost_policy_t<>());
}

/**
 * Compute the Lambert W function on W-1 branch for a value x.
 *
 * @tparam T type of value
 * @param x value
 * @return value of the W-1 branch of the Lambert W function for x
 * @throw std::domain_error if x is less than or equal to `-e^(-1)` or greater
 * than or equal to 0
 */
template <typename T, require_arithmetic_t<T>* = nullptr>
inline double lambert_wm1(const T& x) {
  return boost::math::lambert_wm1(x, boost_policy_t<>());
}

namespace internal {

/**
 * Structure to wrap lambert_w0() so it can be vectorized.
 *
 * @tparam T type of variable
 * @param x variable
 * @return value of the W0 branch of the Lambert W function at x.
 * @throw std::domain_error if x is less than or equal to `-e^(-1)`
 */
struct lambert_w0_fun {
  template <typename T>
  static inline T fun(const T& x) {
    return lambert_w0(x);
  }
};

/**
 * Structure to wrap lambert_wm1() so it can be vectorized.
 *
 * @tparam T type of variable
 * @param x variable
 * @return value of the W-1 branch of the Lambert W function at x.
 * @throw std::domain_error if x is less than or equal to `-e^(-1)` or greater
 * than or equal to 0
 */
struct lambert_wm1_fun {
  template <typename T>
  static inline T fun(const T& x) {
    return lambert_wm1(x);
  }
};
}  // namespace internal

/**
 * Vectorized version of lambert_w0().
 *
 * @tparam T type of container
 * @param x container
 * @return value of the W0 branch of the Lambert W function for each value in x
 * @throw std::domain_error if x is less than or equal to `-e^(-1)`
 */
template <typename T, require_not_stan_scalar_t<T>* = nullptr>
inline auto lambert_w0(const T& x) {
  return apply_scalar_unary<internal::lambert_w0_fun, T>::apply(x);
}

/**
 * Vectorized version of lambert_wm1().
 *
 * @tparam T type of container
 * @param x container
 * @return value of the W0 branch of the Lambert W function for each value in x
 * @throw std::domain_error if x is less than or equal to `-e^(-1)` or greater
 * than or equal to 0
 */
template <typename T, require_not_stan_scalar_t<T>* = nullptr>
inline auto lambert_wm1(const T& x) {
  return apply_scalar_unary<internal::lambert_wm1_fun, T>::apply(x);
}

}  // namespace math
}  // namespace stan

#endif


/*=========================================================================*/
/*                                                                         */
/* @file            Optimizer.cpp                                          */
/* @author          Chirantan Ekbote (ekbote@seas.harvard.edu)             */
/* @date            2012/11/14                                             */
/* @version         0.3                                                    */
/* @brief           Optimizer for generating bright and dark images        */
/*                                                                         */
/*=========================================================================*/

#ifdef USE_OPENMP
#include <omp.h>
#endif
#include <Eigen/SVD>

#include "../inc/Optimizer.h"

#define ROW(n)	1*n,	2*n,	3*n,	4*n,	3*n,	2*n,	1*n

namespace pvrtex {
// Optimization window information
const int Optimizer::window_width_ = 11;
const int Optimizer::window_height_ = 11;
const int Optimizer::matrix_rows_ = 121;
const int Optimizer::matrix_cols_ = 8;
const int Optimizer::offset_x_ = 4;
const int Optimizer::offset_y_ = 4;

// Weight matrices
const float Optimizer::kTopLeft[] =
{
  ROW(0.0625f), 0, 0, 0, 0,	// 1/16
  ROW(0.125f), 0, 0, 0, 0,	// 2/16
  ROW(0.1875f), 0, 0, 0, 0,	// 3/16
  ROW(0.25f), 0, 0, 0, 0,		// 4/16
  ROW(0.1875f), 0, 0, 0, 0,	// 3/16
  ROW(0.125f), 0, 0, 0, 0,	// 2/16
  ROW(0.0625f), 0, 0, 0, 0,	// 1/16
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0
};
const float Optimizer::kTopRight[] =
{
  0, 0, 0, 0, ROW(0.0625f),	// 1/16
  0, 0, 0, 0, ROW(0.125f),	// 2/16
  0, 0, 0, 0, ROW(0.1875f),	// 3/16
  0, 0, 0, 0, ROW(0.25f),		// 4/16
  0, 0, 0, 0, ROW(0.1875f),	// 3/16
  0, 0, 0, 0, ROW(0.125f),	// 2/16
  0, 0, 0, 0, ROW(0.0625f),	// 1/16
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0
};
const float Optimizer::kBottomLeft[] =
{
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0,
  ROW(0.0625f), 0, 0, 0, 0,	// 1/16
  ROW(0.125f), 0, 0, 0, 0,	// 2/16
  ROW(0.1875f), 0, 0, 0, 0,	// 3/16
  ROW(0.25f), 0, 0, 0, 0,		// 4/16
  ROW(0.1875f), 0, 0, 0, 0,	// 3/16
  ROW(0.125f), 0, 0, 0, 0,	// 2/16
  ROW(0.0625f), 0, 0, 0, 0,	// 1/16
};
const float Optimizer::kBottomRight[] =
{
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0,
  ROW(0), 0, 0, 0, 0,
  0, 0, 0, 0, ROW(0.0625f),	// 1/16
  0, 0, 0, 0, ROW(0.125f),	// 2/16
  0, 0, 0, 0, ROW(0.1875f),	// 3/16
  0, 0, 0, 0, ROW(0.25f),		// 4/16
  0, 0, 0, 0, ROW(0.1875f),	// 3/16
  0, 0, 0, 0, ROW(0.125f),	// 2/16
  0, 0, 0, 0, ROW(0.0625f)	// 1/16
};

Optimizer::Optimizer(const Eigen::MatrixXi &o, Eigen::MatrixXi &d,
                     Eigen::MatrixXi &b, SOLVER s, util::DATA_FORMAT f) :

  dark_(d),
  bright_(b),
  orig_(o),
  solv_(s),
  format_(f)
{
  red_ = Eigen::MatrixXi(orig_.rows(), orig_.cols());
  green_ = Eigen::MatrixXi(orig_.rows(), orig_.cols());
  blue_ = Eigen::MatrixXi(orig_.rows(), orig_.cols());
}

Optimizer::~Optimizer() {
}

void Optimizer::ComputeUpdateVector() {
  Eigen::MatrixXi comp = util::ModulateImage(util::Upscale4x4(dark_, format_),
                                             util::Upscale4x4(bright_, format_),
                                             mod_);
#ifdef USE_OPENMP
#pragma omp parallel for
#endif
  for (int y = 0; y < orig_.rows(); ++y) {
    for (int x = 0; x < orig_.cols(); ++x) {
      Eigen::Vector3i diff;
      diff = (util::MakeColorVector(orig_(y,x), util::PVR888) -
              util::MakeColorVector(comp(y,x), util::PVR888));
      //        red_(y,x) = util::MakeRed(orig_(y,x));
      //        green_(y,x) = util::MakeGreen(orig_(y,x));
      //        blue_(y,x) = util::MakeBlue(orig_(y,x));
      red_(y,x) = diff(0);
      green_(y,x) = diff(1);
      blue_(y,x) = diff(2);
    }
  }
}

void Optimizer::OptimizeWindow(int j, int i) {
  Eigen::MatrixXf a(matrix_rows_, matrix_cols_);
  Eigen::MatrixXf w(matrix_rows_, matrix_cols_ / 2);
  Eigen::VectorXf red(matrix_rows_);
  Eigen::VectorXf green(matrix_rows_);
  Eigen::VectorXf blue(matrix_rows_);
  int idx, pixel_x, pixel_y;
  float m, distance;
  float r, g, b;

  /* Construct the optimization window */
  for (int y = 0; y < window_height_; ++y) {
    for (int x = 0; x < window_width_; ++x) {
      /* Get the position of the pixel we want to fetch */
      idx = y*window_width_ + x;
      pixel_x = util::Clamp(offset_x_*i-1 + x, 0, mod_.cols()-1);
      pixel_y = util::Clamp(offset_y_*j-1 + y, 0, mod_.rows()-1);

      /* Fetch the modulation value and the original color*/
      m = mod_(pixel_y, pixel_x);
      r = static_cast<float>(red_(pixel_y, pixel_x));
      g = static_cast<float>(green_(pixel_y, pixel_x));
      b = static_cast<float>(blue_(pixel_y, pixel_x));

      /* Fetch the distance weights and construct the matrix */
      distance = kTopLeft[idx];
      a(idx, 0) = distance * (1.0f - m);
      a(idx, 1) = distance * m;
      w(idx, 0) = distance;
      red(idx) = distance * r;
      green(idx) = distance * g;
      blue(idx) = distance * b;

      /* Top right pxel */
      distance = kTopRight[idx];
      a(idx, 2) = distance * (1.0f - m);
      a(idx, 3) = distance * m;
      w(idx, 1) = distance;
      red(idx) += distance * r;
      green(idx) += distance * g;
      blue(idx) += distance * b;

      /* Bottom left pixel */
      distance = kBottomLeft[idx];
      a(idx, 4) = distance * (1.0f - m);
      a(idx, 5) = distance * m;
      w(idx, 2) = distance;
      red(idx) += distance * r;
      green(idx) += distance * g;
      blue(idx) += distance * b;

      /* Bottom right pixel */
      distance = kBottomRight[idx];
      a(idx, 6) = distance * (1.0f - m);
      a(idx, 7) = distance * m;
      w(idx, 3) = distance;
      red(idx) += distance * r;
      green(idx) += distance * g;
      blue(idx) += distance * b;

    }
  }

  /* Solve for the best colors */
  Eigen::JacobiSVD<Eigen::MatrixXf> svd(a, Eigen::ComputeThinU |
                                        Eigen::ComputeThinV);
  Eigen::VectorXi optimal_red = svd.solve(red).cast<int>();
  Eigen::VectorXi optimal_green;
  Eigen::VectorXi optimal_blue;
  if (format_ == util::PVR444) {
    Eigen::JacobiSVD<Eigen::MatrixXf> svd_w(w, Eigen::ComputeThinU |
                                            Eigen::ComputeThinV);
    Eigen::Vector3i update;
    optimal_green = svd_w.solve(green).cast<int>();
    optimal_blue = svd_w.solve(blue).cast<int>();

    /* Update the dark and bright images */
    for (int x = 0; x < 2; ++x) {
      for (int y = 0; y < 2; ++y) {
        idx = 4*y + 2*x;
        update = Eigen::Vector3i(util::Clamp(optimal_red(idx), -32, 32),
                                 util::Clamp(optimal_green(idx/2), -32, 32),
                                 util::Clamp(optimal_blue(idx/2), -32, 32));
        dark_(j+y, i+x) = util::MakeRGB(
                              util::MakeColorVector(dark_(j+y, i+x), format_) +
                              update, format_);
        update(0) = util::Clamp(optimal_red(idx+1), -32, 32);
        bright_(j+y, i+x) = util::MakeRGB(
                                util::MakeColorVector(bright_(j+y, i+x), format_) +
                                    update, format_);

      }
    }
  } else {
    optimal_green = svd.solve(green).cast<int>();
    optimal_blue = svd.solve(blue).cast<int>();

    /* Update the dark and bright images */
    for (int x = 0; x < 2; ++x) {
      for (int y = 0; y < 2; ++y) {
        idx = 4*y + 2*x;
        dark_(j+y, i+x) = util::MakeRGB(
                              util::MakeColorVector(dark_(j+y, i+x), format_) +
                              Eigen::Vector3i(
                                  util::Clamp(optimal_red(idx), -32, 32),
                                  util::Clamp(optimal_green(idx), -32, 32),
                                  util::Clamp(optimal_blue(idx), -32, 32)),
                                        format_);
        bright_(j+y, i+x) = util::MakeRGB(
                                util::MakeColorVector(bright_(j+y, i+x), format_) +
                                Eigen::Vector3i(
                                    util::Clamp(optimal_red(idx+1), -32, 32),
                                    util::Clamp(optimal_green(idx+1), -32, 32),
                                    util::Clamp(optimal_blue(idx+1), -32, 32)),
                                          format_);
        //      dark_(j+y, i+x) = util::Make565RGB(Eigen::Vector3i(optimal_red(idx),
        //                                                         optimal_green(idx),
        //                                                         optimal_blue(idx)));
        //      bright_(j+y, i+x) = util::Make565RGB(Eigen::Vector3i(optimal_red(idx+1),
        //                                                           optimal_green(idx+1),
        //                                                           optimal_blue(idx+1)));

      }
    }

  }

}

void Optimizer::Optimize(const Eigen::MatrixXf &m) {
  mod_ = m;
  ComputeUpdateVector();
#ifdef USE_OPENMP
#pragma omp parallel for
#endif
  for (int j = 0; j < dark_.rows(); j+=2) {
    for (int i = 0; i < dark_.cols(); i+=2) {
      OptimizeWindow(j, i);
    }
  }
}

void Optimizer::WriteToFile(const char *filename) {

}

} /* namespace pvrtex */


//
//  KDTree.cpp
//
//
//  Created by Sivaram Ambikasaran on 12/3/13.
//
//

#include "KDTree.hpp"
#include <Eigen/Dense>

KDTree::KDTree(const unsigned n_Locations, const unsigned n_Dimension, double* locations, const unsigned n_Properties, double* properties, const unsigned MinParticlesInLeaf, const unsigned nLevels) {
        this->n_Locations       =       n_Locations;
        this->n_Dimension       =       n_Dimension;
        this->n_Properties      =       n_Properties;
        this->MinParticlesInLeaf=       MinParticlesInLeaf;
        this->nLevels           =       nLevels;

        sorted_Contents         =       Eigen::MatrixXd(n_Locations, n_Dimension+n_Properties);

        unsigned count_Location =       0;
        unsigned count_Property =       0;

        for (unsigned j=0; j<n_Locations; ++j) {
                for (unsigned k=0; k<n_Dimension; ++k) {
                        sorted_Contents(j,k)    =       locations[count_Location];
                        ++count_Location;
                }
                for (unsigned k=n_Dimension; k<n_Dimension+n_Properties; ++k) {
                        sorted_Contents(j,k)    =       properties[count_Property];
                        ++count_Property;
                }
        }
}

void KDTree::merge_Sorted_Lists(unsigned n_Left_Start, unsigned n_Left_Size, unsigned n_Right_Start, unsigned n_Right_Size, unsigned n_Index) {

        Eigen::MatrixXd temp_List(n_Left_Size+n_Right_Size, n_Dimension+n_Properties);

        unsigned j_Left =       n_Left_Start;
        unsigned j_Right=       n_Right_Start;

        unsigned j      =       0;

        while (j_Left < n_Left_Start + n_Left_Size && j_Right < n_Right_Start + n_Right_Size) {
                if (sorted_Contents(j_Left, n_Index) < sorted_Contents(j_Right, n_Index)) {
                        temp_List.row(j)=       sorted_Contents.row(j_Left);
                        ++j_Left;
                }
                else {
                        temp_List.row(j)=       sorted_Contents.row(j_Right);
                        ++j_Right;
                }
                ++j;
        }

        while (j_Left < n_Left_Start + n_Left_Size) {
                temp_List.row(j)        =       sorted_Contents.row(j_Left);
                ++j_Left;
                ++j;
        }

        while (j_Right < n_Right_Start + n_Right_Size) {
                temp_List.row(j)        =       sorted_Contents.row(j_Right);
                ++j_Right;
                ++j;
        }

        sorted_Contents.block(n_Left_Start, 0, n_Left_Size + n_Right_Size, n_Dimension+n_Properties)        =       temp_List;
}

void KDTree::merge_Sort(unsigned n_Start, unsigned n_Size, unsigned n_Index) {
        if (n_Size<=1) {
                //      Do nothing
                return;
        }
        else {
                unsigned n_Left_Start   =       n_Start;
                unsigned n_Left_Size    =       n_Size/2;
                unsigned n_Right_Start  =       n_Start+n_Left_Size;
                unsigned n_Right_Size   =       n_Size-n_Left_Size;

                merge_Sort(n_Left_Start, n_Left_Size, n_Index);
                merge_Sort(n_Right_Start, n_Right_Size, n_Index);

                merge_Sorted_Lists(n_Left_Start, n_Left_Size, n_Right_Start, n_Right_Size, n_Index);
        }
}

int KDTree::sort_KDTree(unsigned n_Start, unsigned n_Size, unsigned n_Index, unsigned local_level, std::vector<int>& NumberOfParticlesInLeaves) {
    // static int nLevels = 0;
    n_Index                 =       n_Index%n_Dimension;
    merge_Sort(n_Start, n_Size, n_Index);

    unsigned n_Left_Start   =       n_Start;
    unsigned n_Left_Size    =       n_Size/2;
    unsigned n_Right_Start  =       n_Start+n_Left_Size;
    unsigned n_Right_Size   =       n_Size-n_Left_Size;

    ++n_Index;
    n_Index                 =       n_Index%n_Dimension;

    merge_Sort(n_Left_Start, n_Left_Size, n_Index);
    unsigned n_Left_Bottom_Start   =       n_Left_Start;
    unsigned n_Left_Bottom_Size    =       n_Left_Size/2;
    unsigned n_Left_Top_Start  =       n_Left_Start+n_Left_Bottom_Size;
    unsigned n_Left_Top_Size   =       n_Left_Size-n_Left_Bottom_Size;

    merge_Sort(n_Right_Start, n_Right_Size, n_Index);
    unsigned n_Right_Bottom_Start   =       n_Right_Start;
    unsigned n_Right_Bottom_Size    =       n_Right_Size/2;
    unsigned n_Right_Top_Start  =       n_Right_Start+n_Right_Bottom_Size;
    unsigned n_Right_Top_Size   =       n_Right_Size-n_Right_Bottom_Size;

    ++n_Index;
    n_Index                 =       n_Index%n_Dimension;

    if (local_level < nLevels) {
    // if (n_Left_Bottom_Size >= MinParticlesInLeaf &&  n_Left_Top_Size >= MinParticlesInLeaf && n_Right_Bottom_Size >= MinParticlesInLeaf && n_Right_Top_Size >= MinParticlesInLeaf && local_level < nLevels) {
      local_level++;
      if(local_level==nLevels) {
        // std::cout << n_Left_Bottom_Size << std::endl;
        // std::cout << n_Left_Top_Size << std::endl;
        // std::cout << n_Right_Bottom_Size << std::endl;
        // std::cout << n_Right_Top_Size << std::endl;
        // std::cout << "--------------------------------" << std::endl;
        NumberOfParticlesInLeaves.push_back(n_Left_Bottom_Size);
        NumberOfParticlesInLeaves.push_back(n_Left_Top_Size);
        NumberOfParticlesInLeaves.push_back(n_Right_Bottom_Size);
        NumberOfParticlesInLeaves.push_back(n_Right_Top_Size);
      }
      // std::cout << "local_level: " << local_level << std::endl;
      sort_KDTree(n_Left_Bottom_Start, n_Left_Bottom_Size, n_Index, local_level, NumberOfParticlesInLeaves);
      sort_KDTree(n_Left_Top_Start, n_Left_Top_Size, n_Index, local_level, NumberOfParticlesInLeaves);
      sort_KDTree(n_Right_Bottom_Start, n_Right_Bottom_Size, n_Index, local_level, NumberOfParticlesInLeaves);
      sort_KDTree(n_Right_Top_Start, n_Right_Top_Size, n_Index, local_level, NumberOfParticlesInLeaves);
      // nLevels += 1;
    }
    return 0;
    // return nLevels;
}

int KDTree::sort_KDTree(std::vector<int>& NumberOfParticlesInLeaves) {
        if (n_Locations<=1) {
                //      Do nothing
                return 0;
        }
        else {
                //      Number of point on the left cluster
                return sort_KDTree(0, n_Locations, 0, 0, NumberOfParticlesInLeaves);
        }
}

void KDTree::get_Location_Properties(double* locations, double* properties) {
        unsigned count_Location         =       0;
        unsigned count_Property         =       0;
        for (unsigned j=0; j<n_Locations; ++j) {
                for (unsigned k=0; k<n_Dimension; ++k) {
                        locations[count_Location]       =       sorted_Contents(j,k);
                        ++count_Location;
                }
                for (unsigned k=n_Dimension; k<n_Dimension+n_Properties; ++k) {
                        properties[count_Property]      =       sorted_Contents(j,k);
                        ++count_Property;
                }
        }
}

void KDTree::get_Location_Properties(const unsigned n_Index, double* location, double* properties) {
        for (unsigned k=0; k<n_Dimension; ++k) {
                location[k]     =       sorted_Contents(n_Index,k);
        }
        for (unsigned k=0; k<n_Properties; ++k) {
                properties[k]   =       sorted_Contents(n_Index,k+n_Dimension);
        }
}

/////////////////GLOBAL FUNCTIONS TO BE CALLED FROM MAIN ///////////////////////////
void display(std::string display_String, unsigned n_Locations, unsigned n_Dimension, double* locations, unsigned n_Properties, double* properties) {
        std::cout << display_String << std::endl;

        unsigned count_Location  =       0;
        unsigned count_Property  =       0;

        for (unsigned j=0; j<n_Locations; ++j) {
                for (unsigned k=0; k<n_Dimension; ++k) {
                        std::cout << locations[count_Location] << "\t";
                        ++count_Location;
                }
                for (unsigned k=0; k<n_Properties; ++k) {
                        std::cout << properties[count_Property] << "\t";
                        ++count_Property;
                }
                std::cout << std::endl;
        }
        std::cout << std::endl;
}

void sort_KDTree(unsigned N, unsigned n_Dimension, double* locations, unsigned n_Properties, double* properties, unsigned MinParticlesInLeaf, unsigned nLevels, double* sorted_Locations, double* sorted_Properties, std::vector<std::vector<int> >& boxNumbers, std::vector<int>& NumberOfParticlesInLeaves) {
    KDTree* B                       =       new KDTree(N, n_Dimension, locations, n_Properties, properties, MinParticlesInLeaf, nLevels);
    delete locations;
    delete properties;
    NumberOfParticlesInLeaves.reserve(pow(4,nLevels));
    //      Sorts the locations based on KDTree.
    B->sort_KDTree(NumberOfParticlesInLeaves);
    // converting N ordering to the mirrored C ordering.
    // The sorted locations that the KD Tree outputs, correspond to boxes in N ordering.
    // So here we are generating N ordering sequence of box numbers in terms of the mirrored C ordering.
    // only leaf level ordering is what we want: boxNumbers[nLevels]
    std::vector<int> boxNumbersInLevel0;
    boxNumbersInLevel0.push_back(0);//level 0
    boxNumbers.push_back(boxNumbersInLevel0);
    for (size_t j = 1; j <= nLevels; j++) {
      std::vector<int> boxNumbersInALevel;
      for (size_t k = 0; k < boxNumbers[j-1].size(); k++) {
        boxNumbersInALevel.push_back(boxNumbers[j-1][k]*4 + 0);
        boxNumbersInALevel.push_back(boxNumbers[j-1][k]*4 + 3);
        boxNumbersInALevel.push_back(boxNumbers[j-1][k]*4 + 1);
        boxNumbersInALevel.push_back(boxNumbers[j-1][k]*4 + 2);
      }
      boxNumbers.push_back(boxNumbersInALevel);
    }
    // Obtains the sorted location.
    // sorted_Locations contains the locations sorted as per the KD Tree.(In N ordering of boxes)
    B->get_Location_Properties(sorted_Locations, sorted_Properties);

    // Display the sorted contents.
   // display("Sorted contents: ", N, n_Dimension, sorted_Locations, n_Properties, sorted_Properties);
   // exit(0);
   // Take away from kD tree class: sorted_Locations, sorted_Properties, boxNumbers[nLevels], NumberOfParticlesInLeaves. Give these as inputs to FMM2DTree constructor.
    /////////////////////////////////////////////////////////////////////////
    /////////////////////////////////////////////////////////////////////////

    /* example:
    boxNumbers NumberOfParticlesInLeaves
    0           4
    3           4
    1           4
    2           4
    */
    delete B;
}


//==================================================================================================
/*!
  @file

    @copyright 2016 NumScale SAS

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_MINMOD_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_MINMOD_HPP_INCLUDED

#if defined(DOXYGEN_ONLY)
namespace boost { namespace simd
{
  /*!
    @ingroup group-arithmetic

    This function object computes the smallest of its parameter if they share the same sign,
    zero instead.

    @par Header <boost/simd/function/minmod.hpp>

    @par Notes

    Using `minmod(x, y)` is similar to `x*y > 0 ? min(x, y) : 0`

    @see min, minnum, minnummag, minmag

    @par Example:

      @snippet minmod.cpp minmod

    @par Possible output:

      @snippet minmod.txt minmod

  **/
  Value minmod(Value const& x, Value const& y);
} }
#endif

#include <boost/simd/function/scalar/minmod.hpp>
#include <boost/simd/function/simd/minmod.hpp>

#endif


#include "pre.h"

#include "label.h"

#ifdef CMAKE_USE_DLIB
#include <dlib/statistics.h>
#include <dlib/svm.h>
#endif

#include <cassert>
#include <cmath>

namespace pre
{

dat::Dataset logarithm( const dat::Dataset & d )
{
    double min {};
    dat::apply( [ & ] ( label::Num, const dat::Spectrum & s )
        {
            for( const auto & point : s._y )
            {
                if( point < min )
                {
                    min = point;
                }
            }
        }     , d );

    dat::Dataset ret{ {}, d.second };
    dat::apply( [ & ] ( label::Num l, const dat::Spectrum & s )
        {
            dat::Spectrum transformed{};
            unsigned i {};
            for( const auto & intensity : s._y )
            {
                const auto positive = intensity - min + 1;
                const auto point = std::log( positive );
                transformed._y[ i++ ] = point;
            }
            ret.first[ l ].push_back( transformed );
        }      , ret );

    return ret;
}

Log::Log( const dat::Dataset & )
{
}

dat::Dataset Log::operator()( const dat::Dataset & d ) const
{
    return logarithm( d );
}

void normalize( dat::Dataset & d )
{
    dat::Spectrum mean {};
    dat::apply( [ & ] ( label::Num, const dat::Spectrum & s )
    {
        auto src = s._y.cbegin();
        const auto end = s._y.cend();
        auto dest = mean._y.begin();
        while( src < end )
        {
            * dest++ += * src++;
        }
    }      , d );
    const auto c = dat::count( d );
    for( auto & point : mean._y )
    {
        point /= c;
    }

    dat::Spectrum variance;
    dat::apply( [ & ] ( label::Num, const dat::Spectrum & s )
    {
        auto sp = s._y.cbegin();
        const auto end = s._y.cend();
        auto dest = variance._y.begin();
        while( sp < end )
        {
            const auto a = mean._y[ sp - s._y.cbegin()] ;
            const auto diff = * sp - a;
            * dest += diff * diff;
            assert( diff * diff >= 0 );
            ++ sp;
            ++ dest;
        }
    }      , d );
    for( auto & point : variance._y )
    {
        point = sqrt( point / ( c - 1 ) );
        assert( point >= 0 );
    }

    dat::mutate( [ & ] ( label::Num, dat::Spectrum & s )
    {
        auto src = s._y.begin();
        const auto end = s._y.cend();
        auto m_it = mean._y.cbegin();
        auto v_it = variance._y.cbegin();
        while( src < end )
        {
            if( * v_it < 1e-12  )
            {
                * src = * m_it;
            }
            else
            {
                * src = ( * src - * m_it ) / * v_it;
            }

            ++src;
            ++m_it;
            ++v_it;
        }
    }      , d );
}

Norm::Norm( const dat::Dataset & )
{
}

dat::Dataset Norm::operator()( const dat::Dataset & d ) const
{
    dat::Dataset ret{ d };
    normalize( ret );
    return ret;
}

std::vector< size_t > rank_features( const dat::Dataset & )
{
    return {};
}

#ifdef CMAKE_USE_DLIB
dat::Dataset lda( const dat::Dataset & d )
{
    dlib::matrix< double > samples;
    samples.set_size( static_cast< long >( dat::count( d ) )
                    , dat::Spectrum::_num_points );
    std::vector< unsigned long > labels;
    auto i = 0u;
    dat::apply( [ & ] ( const label::Num l, const dat::Spectrum & s )
        {
            labels.push_back( static_cast< unsigned long >( l ) );
            std::copy( s._y.cbegin()
                     , s._y.cend()
                     , samples.begin() + dat::Spectrum::_num_points * i++ );
        }     , d );
    const auto samples2 = samples;

    dlib::matrix< double, 0, 1 > means;
    dlib::compute_lda_transform( samples, means, labels );

    dlib::matrix< double > ret = samples2 * samples - means; //crash
    std::cout << samples.nc() << ", " << samples.nr() << std::endl;
    std::cout << samples2.nc() << ", " << samples2.nr() << std::endl;
    std::cout << ret.nc() << ", " << ret.nr() << std::endl;

    return {};
}
#endif // CMAKE_USE_DLIB

shark::LinearModel<>
train_encoder( std::vector< shark::RealVector > & inputs
             , unsigned N
             )
{
    const auto dataset{ shark::createUnlabeledDataFromRange( inputs ) };
    shark::PCA pca{ dataset };
    shark::LinearModel<> enc;
    pca.encoder( enc, N );
    return enc;
}

shark::LinearModel<> train_encoder( const dat::Dataset & train
                                  , unsigned N
                                  )
{
    std::vector< shark::RealVector > inputs;
    dat::apply( [ & ] ( auto, const dat::Spectrum & s )
    {
        inputs.push_back( dat::to_shark_vector( s ) );
    }
    , train
    );

    return train_encoder( inputs, N );
}

shark::LinearModel<>
train_encoder( const shark::ClassificationDataset & train
             , unsigned N
             )
{
    std::vector< shark::RealVector > vec;
    for( auto e{ train.elements().begin()}; e < train.elements().end(); ++e )
    {
        vec.push_back( e->input );
    }

    return train_encoder( vec, N );
}

PCA::PCA( const dat::Dataset & train, unsigned dim )
    : _enc{ train_encoder( train, dim ) }
{
}

PCA::PCA( const shark::ClassificationDataset & train, unsigned dim )
    : _enc{ train_encoder( train, dim ) }
{
}

dat::Dataset PCA::operator()( const dat::Dataset & d ) const
{
    if( d.first.empty() )
    {
        return {};
    }

    std::vector< shark::RealVector > inputs;
    std::vector< label::Num > labels;
    dat::apply( [&] ( label::Num l, const dat::Spectrum & s )
    {
        shark::RealVector enc;
        const auto vec{ dat::to_shark_vector( s ) };
        _enc.eval( vec, enc );
        inputs.push_back( std::move( enc ) );
        labels.push_back( l );
    }
              , d );

    const auto sharked{ shark::createLabeledDataFromRange( inputs, labels ) };
    return dat::from_shark_dataset( sharked, d.second );
}

const std::vector< std::string > ALL_PRE{ "log"
                                        , "norm"

#if defined(CMAKE_USE_OPENCV) || defined(CMAKE_USE_SHARK)
                                         , "pca"
#endif
#if defined (CMAKE_USE_OPENCV) || defined (CMAKE_USE_DLIB)
                                        , "lda"
#endif
                                        };

}  // namespace pre


/*
 * DeltaMeasure.cpp
 *
 *  Created on: 31.08.2017
 *      Author: thies
 */

#include <deal.II/base/logstream.h>
#include <deal.II/base/timer.h>
#include <deal.II/grid/grid_tools.h>

#include <base/ConstantMesh.h>
#include <base/Norm.h>
#include <measurements/DeltaMeasure.h>
#include <measurements/SensorValues.h>

#include <stddef.h>
#include <cmath>
#include <iostream>
#include <string>
#include <vector>

namespace wavepi {
namespace measurements {

using namespace dealii;
using namespace wavepi::base;

template <int dim>
DeltaMeasure<dim>::DeltaMeasure(std::shared_ptr<SpaceTimeMesh<dim>> mesh,
                                std::shared_ptr<SensorDistribution<dim>> points,
                                std::shared_ptr<Norm<DiscretizedFunction<dim>>> norm)
    : mesh(mesh), sensor_distribution(points), norm(norm) {
  AssertThrow(mesh && norm, ExcNotInitialized());
}

template <int dim>
SensorValues<dim> DeltaMeasure<dim>::evaluate(const DiscretizedFunction<dim>& field) {
  AssertThrow(sensor_distribution && sensor_distribution->size(), ExcNotInitialized());
  AssertThrow(mesh == field.get_mesh(), ExcMessage("DeltaMeasure called with different meshes"));
  AssertThrow(*norm == *field.get_norm(), ExcMessage("DeltaMeasure called with different norms"));

  SensorValues<dim> res(sensor_distribution);
  auto mapping = StaticMappingQ1<dim>::mapping;

  if (std::dynamic_pointer_cast<ConstantMesh<dim>, SpaceTimeMesh<dim>>(mesh) &&
      sensor_distribution->times_per_point_available()) {
    // specialized implementation that is ordered by points, not time
    // (most expensive operation is find_active_cell_around_point(...))
    auto dof_handler = mesh->get_dof_handler(0);

    for (size_t mpi = 0; mpi < sensor_distribution->get_points().size(); mpi++) {
      auto pos = sensor_distribution->get_points()[mpi];

      std::pair<typename DoFHandler<dim>::active_cell_iterator, Point<dim>> cell_point =
          GridTools::find_active_cell_around_point(mapping, *dof_handler, pos);

      Quadrature<dim> q(GeometryInfo<dim>::project_to_unit_cell(cell_point.second));
      FEValues<dim> fe_values(mapping, dof_handler->get_fe(), q, UpdateFlags(update_values));
      fe_values.reinit(cell_point.first);

      const unsigned int dofs_per_cell = dof_handler->get_fe().dofs_per_cell;

      std::vector<types::global_dof_index> local_dof_indices(dofs_per_cell);
      cell_point.first->get_dof_indices(local_dof_indices);

      for (size_t mti = 0; mti < sensor_distribution->get_times_per_point(mpi).size(); mti++) {
        double mtime = sensor_distribution->get_times_per_point(mpi)[mti];
        size_t ti    = mesh->nearest_time(mtime);

        // (φ_i, δ_p) = fe_values.shape_value(i, 0)
        // this loop calculates the dot product between field[ti] and δ_p directly
        for (unsigned int i = 0; i < dofs_per_cell; i++)
          res[sensor_distribution->index_times_per_point(mpi, mti)] +=
              field[ti][local_dof_indices[i]] * fe_values.shape_value(i, 0);
      }
    }
  } else {
    size_t offset = 0;

    for (size_t mti = 0; mti < sensor_distribution->get_times().size(); mti++) {
      double mtime = sensor_distribution->get_times()[mti];

      size_t ti        = mesh->nearest_time(mtime);
      auto dof_handler = mesh->get_dof_handler(ti);

      for (size_t msi = 0; msi < sensor_distribution->get_points_per_time(mti).size(); msi++) {
        auto pos = sensor_distribution->get_points_per_time(mti)[msi];

        std::pair<typename DoFHandler<dim>::active_cell_iterator, Point<dim>> cell_point =
            GridTools::find_active_cell_around_point(mapping, *dof_handler, pos);

        Quadrature<dim> q(GeometryInfo<dim>::project_to_unit_cell(cell_point.second));
        FEValues<dim> fe_values(mapping, dof_handler->get_fe(), q, UpdateFlags(update_values));
        fe_values.reinit(cell_point.first);

        const unsigned int dofs_per_cell = dof_handler->get_fe().dofs_per_cell;

        std::vector<types::global_dof_index> local_dof_indices(dofs_per_cell);
        cell_point.first->get_dof_indices(local_dof_indices);

        // (φ_i, δ_p) = fe_values.shape_value(i, 0)
        // this loop calculates the dot product between field[ti] and δ_p directly
        for (unsigned int i = 0; i < dofs_per_cell; i++)
          res[msi + offset] += field[ti][local_dof_indices[i]] * fe_values.shape_value(i, 0);
      }

      offset += sensor_distribution->get_points_per_time(mti).size();
    }
  }

  return res;
}

template <int dim>
SensorValues<dim> DeltaMeasure<dim>::zero() {
  return SensorValues<dim>(sensor_distribution);
}

template <int dim>
DiscretizedFunction<dim> DeltaMeasure<dim>::adjoint(const SensorValues<dim>& measurements) {
  AssertThrow(mesh && sensor_distribution && sensor_distribution->size(), ExcNotInitialized());

  DiscretizedFunction<dim> res(mesh);
  auto mapping  = StaticMappingQ1<dim>::mapping;
  size_t offset = 0;

  if (std::dynamic_pointer_cast<ConstantMesh<dim>, SpaceTimeMesh<dim>>(mesh) &&
      sensor_distribution->times_per_point_available()) {
    // specialized implementation that is ordered by points, not time
    // (most expensive operation is find_active_cell_around_point(...))
    auto dof_handler = mesh->get_dof_handler(0);

    for (size_t mpi = 0; mpi < sensor_distribution->get_points().size(); mpi++) {
      auto pos = sensor_distribution->get_points()[mpi];

      std::pair<typename DoFHandler<dim>::active_cell_iterator, Point<dim>> cell_point =
          GridTools::find_active_cell_around_point(mapping, *dof_handler, pos);

      Quadrature<dim> q(GeometryInfo<dim>::project_to_unit_cell(cell_point.second));
      FEValues<dim> fe_values(mapping, dof_handler->get_fe(), q, UpdateFlags(update_values));
      fe_values.reinit(cell_point.first);

      const unsigned int dofs_per_cell = dof_handler->get_fe().dofs_per_cell;

      std::vector<types::global_dof_index> local_dof_indices(dofs_per_cell);
      cell_point.first->get_dof_indices(local_dof_indices);

      for (size_t mti = 0; mti < sensor_distribution->get_times_per_point(mpi).size(); mti++) {
        double mtime = sensor_distribution->get_times_per_point(mpi)[mti];
        size_t ti    = mesh->nearest_time(mtime);

        // (φ_i, δ_p) = fe_values.shape_value(i, 0)
        // this loop calculates the dot product between field[ti] and δ_p directly
        for (unsigned int i = 0; i < dofs_per_cell; i++)
          res[ti][local_dof_indices[i]] +=
              measurements[sensor_distribution->index_times_per_point(mpi, mti)] * fe_values.shape_value(i, 0);
      }
    }
  } else {
    for (size_t mti = 0; mti < sensor_distribution->get_times().size(); mti++) {
      double mtime = sensor_distribution->get_times()[mti];

      size_t ti        = mesh->nearest_time(mtime);
      auto dof_handler = mesh->get_dof_handler(ti);

      for (size_t msi = 0; msi < sensor_distribution->get_points_per_time(mti).size(); msi++) {
        auto pos = sensor_distribution->get_points_per_time(mti)[msi];

        std::pair<typename DoFHandler<dim>::active_cell_iterator, Point<dim>> cell_point =
            GridTools::find_active_cell_around_point(mapping, *dof_handler, pos);

        Quadrature<dim> q(GeometryInfo<dim>::project_to_unit_cell(cell_point.second));
        FEValues<dim> fe_values(mapping, dof_handler->get_fe(), q, UpdateFlags(update_values));
        fe_values.reinit(cell_point.first);

        const unsigned int dofs_per_cell = dof_handler->get_fe().dofs_per_cell;

        std::vector<types::global_dof_index> local_dof_indices(dofs_per_cell);
        cell_point.first->get_dof_indices(local_dof_indices);

        // (φ_i, δ_p) = fe_values.shape_value(i, 0)
        // this loop calculates dot products, not coefficients.
        for (unsigned int i = 0; i < dofs_per_cell; i++)
          res[ti][local_dof_indices[i]] += measurements[msi + offset] * fe_values.shape_value(i, 0);
      }

      offset += sensor_distribution->get_points_per_time(mti).size();
    }
  }

  // indicate which norm we used for the adjoint
  res.set_norm(norm);
  res.dot_transform_inverse();

  return res;
}

template class DeltaMeasure<1>;
template class DeltaMeasure<2>;
template class DeltaMeasure<3>;

}  // namespace measurements
}  // namespace wavepi


// This file is part of LatNet Builder.
//
// Copyright (C) 2012-2021  The LatNet Builder author's, supervised by Pierre L'Ecuyer, Universite de Montreal.
//
// 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 "netbuilder/NetConstructionTraits.h"

#include "latbuilder/GenSeq/GeneratingValues-PLR.h"
#include "latbuilder/SeqCombiner.h"
#include "latbuilder/Util.h"

#include <NTL/GF2X.h>
#include <sstream>
#include <boost/algorithm/string/erase.hpp>

#include "latticetester/ntlwrap.h"

namespace NetBuilder {

    const std::string NetConstructionTraits<NetConstruction::POLYNOMIAL>::name = "Polynomial";

    typedef typename NetConstructionTraits<NetConstruction::POLYNOMIAL>::GenValue GenValue;

    typedef typename NetConstructionTraits<NetConstruction::POLYNOMIAL>::SizeParameter SizeParameter;

    bool NetConstructionTraits<NetConstruction::POLYNOMIAL>::checkGenValue(const GenValue& genValue, const SizeParameter& sizeParameter)
    {
        return IsOne(GCD(genValue,sizeParameter));
    }

    unsigned int NetConstructionTraits<NetConstruction::POLYNOMIAL>::nRows(const SizeParameter& sizeParameter) {return (unsigned int) deg(sizeParameter); }

    unsigned int NetConstructionTraits<NetConstruction::POLYNOMIAL>::nCols(const SizeParameter& sizeParameter) {return (unsigned int) deg(sizeParameter); }

    void expandSeries(const GenValue& genValue, const SizeParameter& sizeParameter, std::vector<unsigned int>& expansion, unsigned int expansion_limit){
        int m = (int) deg(sizeParameter);
        for(int l = 1; l<= (int) expansion_limit ; l++){
            int res =  (m-l >=0 && IsOne(coeff(genValue, m-l)))? 1 : 0;
            int start = (l-m > 1) ? (l-m) : 1;
            for( int p = start; p < l; p++){
                res = ( res + expansion[p-1] * NTL::conv<int>(coeff(sizeParameter, m-(l-p)))) %2;
            }
            expansion[l-1] = res;
        }
    }

    GeneratingMatrix*  NetConstructionTraits<NetConstruction::POLYNOMIAL>::createGeneratingMatrix(const GenValue& genValue, const SizeParameter& sizeParameter, const Dimension& dimension_j, const unsigned int nRows)
    {
        unsigned int m = (unsigned int) (deg(sizeParameter));
        unsigned int finalnRows = (nRows == 0)? m : nRows;
        GeneratingMatrix* genMat = new GeneratingMatrix(finalnRows, m);
        std::vector<unsigned int> expansion(finalnRows + m);
        expandSeries(genValue, sizeParameter, expansion, finalnRows + m);
        for(unsigned int c = 0; c < m; c++)
        {
            for(unsigned int row = 0; row < finalnRows; row++)
            {
                (*genMat)(row,c) = expansion[c + row];
            }
        }
        return genMat;
    }

    typename NetConstructionTraits<NetConstruction::POLYNOMIAL>::GenValueSpaceCoordSeq NetConstructionTraits<NetConstruction::POLYNOMIAL>::genValueSpaceCoord(Dimension coord, const SizeParameter& sizeParameter)
    {
        if (coord==0)
        {
            Polynomial fakeModulus;
            SetX(fakeModulus);
            return GenValueSpaceCoordSeq(fakeModulus);
        }
        else
        {
           return GenValueSpaceCoordSeq(sizeParameter);
        }
    }

    typename NetConstructionTraits<NetConstruction::POLYNOMIAL>::GenValueSpaceSeq NetConstructionTraits<NetConstruction::POLYNOMIAL>::genValueSpace(Dimension dimension, const SizeParameter& sizeParameter)
    {
        std::vector<GenValueSpaceCoordSeq> seqs;
        seqs.reserve(dimension);
        for(Dimension coord = 0; coord < dimension; ++coord)
        {
            seqs.push_back(genValueSpaceCoord(coord, sizeParameter));
        }
        return GenValueSpaceSeq(seqs);
    }

    std::string NetConstructionTraits<NetConstruction::POLYNOMIAL>::format(const std::vector<std::shared_ptr<GeneratingMatrix>>& genMatrices, const std::vector<std::shared_ptr<GenValue>>& genVals, const SizeParameter& sizeParameter, OutputStyle outputStyle, unsigned int interlacingFactor)
    {
        std::string res;
        std::ostringstream stream;

        if (outputStyle == OutputStyle::TERMINAL){
            stream << "Polynomial Digital Net - Modulus = " << LatBuilder::IndexOfPolynomial(sizeParameter) << " - GeneratingVector =" << std::endl;
            for (unsigned int coord = 0; coord < genVals.size(); coord++){
                stream << "  " << LatBuilder::IndexOfPolynomial(*(genVals[coord])) << std::endl;
            }
        }

        else if (outputStyle == OutputStyle::LATTICE){
            stream << "# Parameters for a polynomial lattice rule in base 2" << std::endl;
            stream << genVals.size() / interlacingFactor << "      #  s =  " << genVals.size() / interlacingFactor << " dimensions" << std::endl;
            if (interlacingFactor > 1){
                stream << interlacingFactor << "    # Interlacing factor" << std::endl;
                stream << genVals.size() << "    # Number of components = interlacing factor x dimension" << std::endl;
            }
            stream << (int) deg(sizeParameter) << "      # n = 2^";
            stream <<  (int) deg(sizeParameter) << " = " << (int)pow(2,deg(sizeParameter) ) << " points"<< std::endl;

            stream << LatBuilder::IndexOfPolynomial(sizeParameter) << "   # polynomial modulus" << std::endl;
            stream << "# Coordinates of generating vector, starting at j=1" << std::endl;
            std::string res;
            for (unsigned int coord = 0; coord < genVals.size(); coord++){
                res += std::to_string(LatBuilder::IndexOfPolynomial(*(genVals[coord]))) + "\n";
            }
            res.pop_back();
            stream << res;
        }

        res += stream.str();
        return res;
    }
}


// Boost.Geometry

// Copyright (c) 2015-2016 Oracle and/or its affiliates.

// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

// 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_FORMULAS_ANDOYER_INVERSE_HPP
#define BOOST_GEOMETRY_FORMULAS_ANDOYER_INVERSE_HPP

#include <boost/math/constants/constants.hpp>

#include <boost/geometry/core/radius.hpp>
#include <boost/geometry/core/srs.hpp>

#include <boost/geometry/util/condition.hpp>
#include <boost/geometry/util/math.hpp>

#include <boost/geometry/algorithms/detail/flattening.hpp>

#include <boost/geometry/formulas/differential_quantities.hpp>
#include <boost/geometry/formulas/result_inverse.hpp>

namespace boost { namespace geometry { namespace formula
{

/*!
\brief The solution of the inverse problem of geodesics on latlong coordinates,
       Forsyth-Andoyer-Lambert type approximation with first order terms.
\author See
    - Technical Report: PAUL D. THOMAS, MATHEMATICAL MODELS FOR NAVIGATION SYSTEMS, 1965
      http://www.dtic.mil/docs/citations/AD0627893
    - Technical Report: PAUL D. THOMAS, SPHEROIDAL GEODESICS, REFERENCE SYSTEMS, AND LOCAL GEOMETRY, 1970
      http://www.dtic.mil/docs/citations/AD703541
*/
template <
    typename CT,
    bool EnableDistance,
    bool EnableAzimuth,
    bool EnableReverseAzimuth = false,
    bool EnableReducedLength = false,
    bool EnableGeodesicScale = false
>
class andoyer_inverse
{
    static const bool CalcQuantities = EnableReducedLength || EnableGeodesicScale;
    static const bool CalcAzimuths = EnableAzimuth || EnableReverseAzimuth || CalcQuantities;
    static const bool CalcFwdAzimuth = EnableAzimuth || CalcQuantities;
    static const bool CalcRevAzimuth = EnableReverseAzimuth || CalcQuantities;

public:
    typedef result_inverse<CT> result_type;

    template <typename T1, typename T2, typename Spheroid>
    static inline result_type apply(T1 const& lon1,
                                    T1 const& lat1,
                                    T2 const& lon2,
                                    T2 const& lat2,
                                    Spheroid const& spheroid)
    {
        result_type result;

        // coordinates in radians

        if ( math::equals(lon1, lon2) && math::equals(lat1, lat2) )
        {
            return result;
        }

        CT const c0 = CT(0);
        CT const c1 = CT(1);
        CT const pi = math::pi<CT>();
        CT const f = detail::flattening<CT>(spheroid);

        CT const dlon = lon2 - lon1;
        CT const sin_dlon = sin(dlon);
        CT const cos_dlon = cos(dlon);
        CT const sin_lat1 = sin(lat1);
        CT const cos_lat1 = cos(lat1);
        CT const sin_lat2 = sin(lat2);
        CT const cos_lat2 = cos(lat2);

        // H,G,T = infinity if cos_d = 1 or cos_d = -1
        // lat1 == +-90 && lat2 == +-90
        // lat1 == lat2 && lon1 == lon2
        CT cos_d = sin_lat1*sin_lat2 + cos_lat1*cos_lat2*cos_dlon;
        // on some platforms cos_d may be outside valid range
        if (cos_d < -c1)
            cos_d = -c1;
        else if (cos_d > c1)
            cos_d = c1;

        CT const d = acos(cos_d); // [0, pi]
        CT const sin_d = sin(d);  // [-1, 1]

        if ( BOOST_GEOMETRY_CONDITION(EnableDistance) )
        {
            CT const K = math::sqr(sin_lat1-sin_lat2);
            CT const L = math::sqr(sin_lat1+sin_lat2);
            CT const three_sin_d = CT(3) * sin_d;

            CT const one_minus_cos_d = c1 - cos_d;
            CT const one_plus_cos_d = c1 + cos_d;
            // cos_d = 1 or cos_d = -1 means that the points are antipodal

            CT const H = math::equals(one_minus_cos_d, c0) ?
                            c0 :
                            (d + three_sin_d) / one_minus_cos_d;
            CT const G = math::equals(one_plus_cos_d, c0) ?
                            c0 :
                            (d - three_sin_d) / one_plus_cos_d;

            CT const dd = -(f/CT(4))*(H*K+G*L);

            CT const a = get_radius<0>(spheroid);

            result.distance = a * (d + dd);
        }

        if ( BOOST_GEOMETRY_CONDITION(CalcAzimuths) )
        {
            // sin_d = 0 <=> antipodal points
            if (math::equals(sin_d, c0))
            {
                // T = inf
                // dA = inf
                // azimuth = -inf
                result.azimuth = lat1 <= lat2 ? c0 : pi;
            }
            else
            {
                CT const c2 = CT(2);

                CT A = c0;
                CT U = c0;
                if ( ! math::equals(cos_lat2, c0) )
                {
                    CT const tan_lat2 = sin_lat2/cos_lat2;
                    CT const M = cos_lat1*tan_lat2-sin_lat1*cos_dlon;
                    A = atan2(sin_dlon, M);
                    CT const sin_2A = sin(c2*A);
                    U = (f/ c2)*math::sqr(cos_lat1)*sin_2A;
                }

                CT B = c0;
                CT V = c0;
                if ( ! math::equals(cos_lat1, c0) )
                {
                    CT const tan_lat1 = sin_lat1/cos_lat1;
                    CT const N = cos_lat2*tan_lat1-sin_lat2*cos_dlon;
                    B = atan2(sin_dlon, N);
                    CT const sin_2B = sin(c2*B);
                    V = (f/ c2)*math::sqr(cos_lat2)*sin_2B;
                }

                CT const T = d / sin_d;

                // even with sin_d == 0 checked above if the second point
                // is somewhere in the antipodal area T may still be great
                // therefore dA and dB may be great and the resulting azimuths
                // may be some more or less arbitrary angles

                if (BOOST_GEOMETRY_CONDITION(CalcFwdAzimuth))
                {
                    CT const dA = V*T - U;
                    result.azimuth = A - dA;
                    normalize_azimuth(result.azimuth, A, dA);
                }

                if (BOOST_GEOMETRY_CONDITION(CalcRevAzimuth))
                {
                    CT const dB = -U*T + V;
                    result.reverse_azimuth = pi - B - dB;
                    if (result.reverse_azimuth > pi)
                    {
                        result.reverse_azimuth -= 2 * pi;
                    }
                    normalize_azimuth(result.reverse_azimuth, B, dB);
                }
            }
        }

        if (BOOST_GEOMETRY_CONDITION(CalcQuantities))
        {
            typedef differential_quantities<CT, EnableReducedLength, EnableGeodesicScale, 1> quantities;
            quantities::apply(dlon, sin_lat1, cos_lat1, sin_lat2, cos_lat2,
                              result.azimuth, result.reverse_azimuth,
                              get_radius<2>(spheroid), f,
                              result.reduced_length, result.geodesic_scale);
        }

        return result;
    }

private:
    static inline void normalize_azimuth(CT & azimuth, CT const& A, CT const& dA)
    {
        CT const c0 = 0;

        if (A >= c0) // A indicates Eastern hemisphere
        {
            if (dA >= c0) // A altered towards 0
            {
                if (azimuth < c0)
                {
                    azimuth = c0;
                }
            }
            else // dA < 0, A altered towards pi
            {
                CT const pi = math::pi<CT>();
                if (azimuth > pi)
                {
                    azimuth = pi;
                }
            }
        }
        else // A indicates Western hemisphere
        {
            if (dA <= c0) // A altered towards 0
            {
                if (azimuth > c0)
                {
                    azimuth = c0;
                }
            }
            else // dA > 0, A altered towards -pi
            {
                CT const minus_pi = -math::pi<CT>();
                if (azimuth < minus_pi)
                {
                    azimuth = minus_pi;
                }
            }
        }
    }
};

}}} // namespace boost::geometry::formula

#endif // BOOST_GEOMETRY_FORMULAS_ANDOYER_INVERSE_HPP


#include <vector>
#include <type_traits>

#include <boost/range/adaptor/transformed.hpp>
namespace ba = boost::adaptors;

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/functional.h>
namespace py = pybind11;

#include <dionysus/row-reduction.h>
#include <dionysus/ordinary-persistence.h>
#include <dionysus/standard-reduction.h>

#include "filtration.h"
#include "persistence.h"                // to get access to PyReducedMatrix::Chain
#include "zigzag-persistence.h"
#include "diagram.h"
#include "progress.h"

PYBIND11_MAKE_OPAQUE(PyReducedMatrix::Chain);      // we want to provide our own binding for Chain

struct Time
{
            Time(float t_, size_t i_, unsigned dim_, bool dir_):
                t(t_), i(i_), dim(dim_), dir(dir_)             {}

    bool    operator<(const Time& other) const
    {
        if (t == other.t)
        {
            if (dir && !other.dir)      // add comes before remove
                return true;
            else if (!dir && other.dir)
                return false;
            else if (dim == other.dim)
                return i < other.i;
            else if (dir)
                return dim < other.dim;
            else // if (!dir)
                return other.dim < dim;
        }
        else
            return t < other.t;
    }

    float       t;
    size_t      i;
    unsigned    dim;
    bool        dir;
};

class PyTimeIndexMap
{
    public:
        using ZZIndex           = PyZigzagPersistence::Index;
        using FIndex            = size_t;
        using Map               = std::unordered_map<ZZIndex, FIndex>;
        using const_iterator    = Map::const_iterator;

        void                    set(const ZZIndex& x, const FIndex& y)  { m_[x] = y; }
        void                    remove(const ZZIndex& x)                { m_.erase(x); }

        FIndex                  operator[](const ZZIndex& x) const      { return m_.find(x)->second; }
        size_t                  size() const                            { return m_.size(); }

        const_iterator          begin() const                           { return m_.begin(); }
        const_iterator          end() const                             { return m_.end(); }

    private:
        Map     m_;
};

using Times     = std::vector<std::vector<float>>;
using Callback  = std::function<void(size_t, float, bool, const PyZigzagPersistence*, const PyTimeIndexMap*)>;

std::tuple<PyZigzagPersistence, std::vector<PyDiagram>, PyTimeIndexMap>
zigzag_homology_persistence(const PyFiltration&     f,
                            const Times&            times_,
                            PyZpField::Element      prime,
                            const Callback&         callback,
                            bool                    show_progress)
{
    using Index          = PyZigzagPersistence::Index;
    using CellChainEntry = dionysus::ChainEntry<PyZpField, PySimplex>;
    using ChainEntry     = dionysus::ChainEntry<PyZpField, Index>;

    std::vector<Time> times;
    for (size_t i = 0; i < times_.size(); ++i)
    {
        int dim = f[i].dimension();
        bool dir = true;
        for (float t : times_[i])
        {
            times.emplace_back(t, i, dim, dir);
            dir = !dir;
        }
    }
    std::sort(times.begin(), times.end());

    std::unique_ptr<Progress> progress(new NoProgress);
    if (show_progress)
        progress = std::unique_ptr<Progress>(new ShowProgress(times.size()));

    std::vector<PyDiagram> diagrams;
    PyZpField field(prime);
    PyZigzagPersistence persistence(field);
    unsigned op = 0;
    unsigned cell = 0;
    std::vector<unsigned>   cells(f.size(), -1);
    PyTimeIndexMap          cells_inv_;
    for (auto& tt : times)
    {
        (*progress)();

        size_t i = tt.i; float t = tt.t; bool dir = tt.dir;

        auto& c = f[i];
        if (dir)
        {
            cells_inv_.set(cell, i);
            cells[i] = cell++;

            Index pair = persistence.add(c.boundary(persistence.field()) |
                                                    ba::transformed([&](const CellChainEntry& e)
                                                    {
                                                        auto idx = f.index(e.index());
                                                        return ChainEntry(e.element(), cells[idx]);
                                                    }));

            if (pair != persistence.unpaired())
            {
                auto t_birth = times[pair].t;
                if (t_birth != t)
                {
                    int dim = c.dimension()-1;
                    while (dim+1 > diagrams.size())
                        diagrams.emplace_back();
                    diagrams[dim].emplace_back(t_birth, t, pair);
                }
            }

            ++op;
        } else
        {
            Index pair = persistence.remove(cells[i]);
            cells_inv_.remove(cells[i]);
            cells[i] = -1;
            if (pair != persistence.unpaired())
            {
                auto t_birth = times[pair].t;
                if (t_birth != t)
                {
                    int dim = c.dimension();
                    while (dim+1 > diagrams.size())
                        diagrams.emplace_back();
                    diagrams[dim].emplace_back(t_birth, t, pair);
                }
            }
            ++op;
        }

        callback(i,t,dir,&persistence,&cells_inv_);
    }

    // add infinite points
    constexpr float inf = std::numeric_limits<float>::infinity();
    for (auto& birth_idx : persistence.alive_ops())
    {
        auto i_birth   = times[birth_idx].i;
        auto t_birth   = times[birth_idx].t;
        auto dir_birth = times[birth_idx].dir;

        int dim = f[i_birth].dimension();
        if (!dir_birth)     // born on removal
            dim -= 1;
        while (dim+1 > diagrams.size())
            diagrams.emplace_back();
        diagrams[dim].emplace_back(t_birth, inf, birth_idx);
    }

    return std::make_tuple(std::move(persistence), std::move(diagrams), std::move(cells_inv_));
}

// custom iterator that preserves right filtration order
struct PyZZAliveCycleIterator
{
            PyZZAliveCycleIterator(const PyZigzagPersistence& zz_, py::object ref_):
                zz(zz_), ref(ref_)
    {
        for (PyZigzagPersistence::Index x : zz.alive_cycles())
            indices.push_back(x);
        std::sort(indices.begin(), indices.end());
    }

    PyReducedMatrix::Chain  next()
    {
        if (idx == indices.size())
            throw py::stop_iteration();

        PyReducedMatrix::Chain c;
        for (auto& x : zz.cycle(indices[idx]))
            c.emplace_back(x.element(), std::get<0>(x.index()));
        ++idx;
        return c;
    }

    const PyZigzagPersistence&                  zz;
    py::object                                  ref;
    size_t                                      idx = 0;
    std::vector<PyZigzagPersistence::Index>     indices;
};

#include "chain.h"

void init_zigzag_persistence(py::module& m)
{
    using namespace pybind11::literals;
    m.def("zigzag_homology_persistence",   &zigzag_homology_persistence, "filtration"_a, "times"_a, "prime"_a = 2,
                                                                         "callback"_a = Callback([](size_t, float, bool, const PyZigzagPersistence*, const PyTimeIndexMap*){}),
                                                                         "progress"_a = false,
          R"(
          compute zigzag homology persistence of the filtration with respect to the given times

          Args:
              filtration: an instance of :class:`~dionysus._dionysus.Filtration` with the set of simplices used in the zigzag construction
              times:      a list of lists; the outer list runs parallel with the filtration;
                          the inner list specifies for each simplex when it enters and leaves the zigzag
                          (even entries, starting the indexing from 0, are interpreted as appearance times, odd entires as disappearance)
              prime:      prime modulo which to perform computation
              callback:   function to call after every step in the zigzag; it gets arguments `(i,t,d,zz,cells)`,
                          where `i` is the index of the simplex being added or removed, `t` is the time,
                          `d` is the "direction" (`True` if the simplex is being added, `False` if it`s being removed),
                          `zz` is the current state of the :class:`~dionysus._dionysus.ZigzagPersistence`,
                          `cells` is the map from the internal indices of the zigzag representation to the filtration indices,
              progress:   show a progress bar.

          Returns:
              A triple. The first element is an instance of
              :class:`~dionysus._dionysus.ZigzagPersistence`, which offers access to the cycles
              alive at the end of the zigzag; the second is a list of
              persistence diagrams; the third is an instance of
              :class:`~dionysus._dionysus.TimeIndexMap` for translating cycles
              from the internal representation to filtration indices.
          )");

   py::class_<PyZZAliveCycleIterator>(m, "ZZAliveCycleIterator")
        .def("__iter__", [](PyZZAliveCycleIterator& it) -> PyZZAliveCycleIterator& { return it; })
        .def("__next__", &PyZZAliveCycleIterator::next);

    py::class_<PyZigzagPersistence>(m, "ZigzagPersistence", "representation of the current homology basis")
        .def(py::init<PyZpField>())
        .def("__len__",     &PyZigzagPersistence::alive_size,       "number of alive cycles")
        .def("__iter__",    [](py::object zz)   { return PyZZAliveCycleIterator(zz.cast<const PyZigzagPersistence&>(), zz); }, "iterator over the alive cycles")
        .def("__repr__",    [](const PyZigzagPersistence& zz)
                            { std::ostringstream oss; oss << "Zigzag persistence with " << zz.alive_size() << " alive cycles"; return oss.str(); })
    ;

    py::class_<PyTimeIndexMap>(m, "TimeIndexMap", "map from the internal representation of zigzag persistence to filtration indices")
        .def("__len__",     &PyTimeIndexMap::size,                      "size of the map")
        .def("__getitem__", [](const PyTimeIndexMap& m, size_t i) { return m[i]; }, "access the filtration index of the given internal index")
        .def("__iter__",    [](const PyTimeIndexMap& m) { return py::make_iterator(m.begin(), m.end()); },
                                py::keep_alive<0, 1>() /* Essential: keep object alive while iterator exists */,
                                "iterate over the entries of the map")
    ;
}