#include <fstream>
#include <boost/program_options.hpp>
#include <iostream>
#include <algorithm>
#include <iterator>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <cmath>
#include <omp.h>

#include <fftw3.h>
#include <math.h>
#include "cmodules/timer.h"

#include "param.h"
#include "bioem.h"
#include "model.h"
#include "map.h"

#include "bioem_algorithm.h"

using namespace boost;
namespace po = boost::program_options;

using namespace std;

// A helper function of Boost
template<class T>
ostream& operator<<(ostream& os, const vector<T>& v)
{
	copy(v.begin(), v.end(), ostream_iterator<T>(os, " "));
	return os;
}

bioem::bioem()
{
	FFTAlgo = getenv("FFTALGO") == NULL ? 0 : atoi(getenv("FFTALGO"));
}

bioem::~bioem()
{

}

int bioem::configure(int ac, char* av[])
{
	// **************************************************************************************
	// **** Configuration Routine using boost for extracting parameters, models and maps ****
	// **************************************************************************************
	// ****** And Precalculating necessary grids, map crosscorrelations and kernels  ********
	// *************************************************************************************

	// *** Inizialzing default variables ***
	std::string infile, modelfile, mapfile;
	Model.readPDB = false;
	param.writeAngles = false;
	param.dumpMap = false;
	param.loadMap = false;
	RefMap.readMRC = false;
	RefMap.readMultMRC = false;

	// *************************************************************************************
	cout << " ++++++++++++ FROM COMMAND LINE +++++++++++\n\n";
	// *************************************************************************************

	// ********************* Command line reading input with BOOST ************************

	try {
		po::options_description desc("Command line inputs");
		desc.add_options()
		("Inputfile", po::value<std::string>(), "(Mandatory) Name of input parameter file")
		("Modelfile", po::value< std::string>() , "(Mandatory) Name of model file")
		("Particlesfile", po::value< std::string>(), "(Mandatory) Name of paricles file")
		("ReadPDB", "(Optional) If reading model file in PDB format")
		("ReadMRC", "(Optional) If reading particle file in MRC format")
		("ReadMultipleMRC", "(Optional) If reading Multiple MRCs")
		("DumpMaps", "(Optional) Dump maps after they were red from maps file")
		("LoadMapDump", "(Optional) Read Maps from dump instead of maps file")
		("help", "(Optional) Produce help message")
		;


		po::positional_options_description p;
		p.add("Inputfile", -1);
		p.add("Modelfile", -1);
		p.add("Particlesfile", -1);
		p.add("ReadPDB", -1);
		p.add("ReadMRC", -1);
		p.add("ReadMultipleMRC", -1);
		p.add("DumpMaps", -1);
		p.add("LoadMapDump", -1);

		po::variables_map vm;
		po::store(po::command_line_parser(ac, av).
				  options(desc).positional(p).run(), vm);
		po::notify(vm);

		if((ac < 6)) {
			std::cout << desc << std::endl;
			return 1;
		}
		if (vm.count("help")) {
			cout << "Usage: options_description [options]\n";
			cout << desc;
			return 1;
		}

		if (vm.count("Inputfile"))
		{
			cout << "Input file is: ";
			cout << vm["Inputfile"].as< std::string >() << "\n";
			infile = vm["Inputfile"].as< std::string >();
		}
		if (vm.count("Modelfile"))
		{
			cout << "Model file is: "
				 << vm["Modelfile"].as<  std::string  >() << "\n";
			modelfile = vm["Modelfile"].as<  std::string  >();
		}

		if (vm.count("ReadPDB"))
		{
			cout << "Reading model file in PDB format.\n";
			Model.readPDB = true;
		}

		if (vm.count("ReadMRC"))
		{
			cout << "Reading particle file in MRC format.\n";
			RefMap.readMRC=true;
		}

		if (vm.count("ReadMultipleMRC"))
		{
			cout << "Reading Multiple MRCs.\n";
			RefMap.readMultMRC=true;
		}

		if (vm.count("DumpMaps"))
		{
			cout << "Dumping Maps after reading from file.\n";
			param.dumpMap = true;
		}

		if (vm.count("LoadMapDump"))
		{
			cout << "Loading Map dump.\n";
			param.loadMap = true;
		}

		if (vm.count("Particlesfile"))
		{
			cout << "Paricle file is: "
				 << vm["Particlesfile"].as< std::string >() << "\n";
			mapfile = vm["Particlesfile"].as< std::string >();
		}
	}
	catch(std::exception& e)
	{
		cout << e.what() << "\n";
		return 1;
	}

	// ********************* Reading Parameter Input ***************************
	// copying inputfile to param class
	param.fileinput = infile.c_str();
	param.readParameters();

	// ********************* Reading Model Input ******************************
	// copying modelfile to model class
	Model.filemodel = modelfile.c_str();
	Model.readModel();

	// ********************* Reading Particle Maps Input **********************
	// ********* HERE: PROBLEM if maps dont fit on the memory!! ***************
	// copying mapfile to ref map class
	param.filemap = mapfile.c_str();
	RefMap.readRefMaps(param);

	// ****************** Precalculating Necessary Stuff *********************
	precalculate();

	if (getenv("BIOEM_DEBUG_BREAK"))
	{
		param.nTotGridAngles = atoi(getenv("BIOEM_DEBUG_BREAK"));
		param.nTotCTFs = atoi(getenv("BIOEM_DEBUG_BREAK"));
	}

	pProb.init(RefMap.ntotRefMap, param.nTotGridAngles, *this);

	deviceInit();

	return(0);
}

void bioem::cleanup()
{
	//Deleting allocated pointers
	free_device_host(pProb.ptr);
	RefMap.freePointers();
}

int bioem::precalculate()
{
	// **************************************************************************************
	// **Precalculating Routine of Orientation grids, Map crosscorrelations and CTF Kernels**
	// **************************************************************************************

	// Generating Grids of orientations
	param.CalculateGridsParam();

	// Precalculating CTF Kernels stored in class Param
	param.CalculateRefCTF();

	//Precalculate Maps
	RefMap.precalculate(param, *this);

	return(0);
}

int bioem::run()
{
	// **************************************************************************************
	// **** Main BioEM routine, projects, convolutes and compares with Map using OpenMP ****
	// **************************************************************************************

	// **** If we want to control the number of threads -> omp_set_num_threads(XX); ******
	// ****************** Declarying class of Probability Pointer  *************************

	printf("\tInitializing Probabilities\n");
	// Inizialzing Probabilites to zero and constant to -Infinity
	for (int iRefMap = 0; iRefMap < RefMap.ntotRefMap; iRefMap ++)
	{
		bioem_Probability_map& pProbMap = pProb.getProbMap(iRefMap);

		pProbMap.Total = 0.0;
		pProbMap.Constoadd = -9999999;
		pProbMap.max_prob = -9999999;
		for (int iOrient = 0; iOrient < param.nTotGridAngles; iOrient ++)
		{
			bioem_Probability_angle& pProbAngle = pProb.getProbAngle(iRefMap, iOrient);

			pProbAngle.forAngles = 0.0;
			pProbAngle.ConstAngle = -99999999;
		}
	}
	// **************************************************************************************
	deviceStartRun();
	{
		const int count = omp_get_max_threads();
		localCCT = new mycomplex_t*[count];
		lCC = new myfloat_t*[count];
		for (int i = 0;i < count;i++)
		{
			localCCT[i] = (mycomplex_t *) myfftw_malloc(sizeof(mycomplex_t) * param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D);
			lCC[i] = (myfloat_t *) myfftw_malloc(sizeof(myfloat_t) * param.param_device.NumberPixels * param.param_device.NumberPixels);
		}
	}

	// ******************************** MAIN CYCLE ******************************************

	// *** Declaring Private variables for each thread *****
	mycomplex_t* proj_mapFFT;
	myfloat_t* conv_map = new myfloat_t[param.param_device.NumberPixels * param.param_device.NumberPixels];
	mycomplex_t* conv_mapFFT;
	myfloat_t sumCONV, sumsquareCONV;

	//allocating fftw_complex vector
	proj_mapFFT = (mycomplex_t *) myfftw_malloc(sizeof(mycomplex_t) * param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D);
	conv_mapFFT = (mycomplex_t *) myfftw_malloc(sizeof(mycomplex_t) * param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D);

	HighResTimer timer;

	printf("\tMain Loop (GridAngles %d, CTFs %d, RefMaps %d, Shifts (%d/%d)²), Pixels %d²\n", param.nTotGridAngles, param.nTotCTFs, RefMap.ntotRefMap, 2 * param.param_device.maxDisplaceCenter + param.param_device.GridSpaceCenter, param.param_device.GridSpaceCenter, param.param_device.NumberPixels);
	printf("\tInner Loop Count (%d %d %d) %lld\n", param.param_device.maxDisplaceCenter, param.param_device.GridSpaceCenter, param.param_device.NumberPixels, (long long int) (param.param_device.NumberPixels * param.param_device.NumberPixels * (2 * param.param_device.maxDisplaceCenter / param.param_device.GridSpaceCenter + 1) * (2 * param.param_device.maxDisplaceCenter / param.param_device.GridSpaceCenter + 1)));
	for (int iOrient = 0; iOrient < param.nTotGridAngles; iOrient++)
	{
		// ***************************************************************************************
		// ***** Creating Projection for given orientation and transforming to Fourier space *****
		timer.ResetStart();
		createProjection(iOrient, proj_mapFFT);
		printf("Time Projection %d: %f\n", iOrient, timer.GetCurrentElapsedTime());

		// ***************************************************************************************
		// ***** **** Internal Loop over convolutions **** *****
		for (int iConv = 0; iConv < param.nTotCTFs; iConv++)
		{
			printf("\t\tConvolution %d %d\n", iOrient, iConv);
			// *** Calculating convolutions of projection map and crosscorrelations ***

			timer.ResetStart();
			createConvolutedProjectionMap(iOrient, iConv, proj_mapFFT, conv_map, conv_mapFFT, sumCONV, sumsquareCONV);
			printf("Time Convolution %d %d: %f\n", iOrient, iConv, timer.GetCurrentElapsedTime());

			// ***************************************************************************************
			// *** Comparing each calculated convoluted map with all experimental maps ***
			timer.ResetStart();
			compareRefMaps(iOrient, iConv, conv_map, conv_mapFFT, sumCONV, sumsquareCONV);

			const double compTime = timer.GetCurrentElapsedTime();
			const int nShifts = 2 * param.param_device.maxDisplaceCenter / param.param_device.GridSpaceCenter + 1;
			const double nFlops = (double) RefMap.ntotRefMap * (double) nShifts * (double) nShifts *
								  (((double) param.param_device.NumberPixels - (double) param.param_device.maxDisplaceCenter / 2.) * ((double) param.param_device.NumberPixels - (double) param.param_device.maxDisplaceCenter / 2.) * 5. + 25.) / compTime;
			const double nGBs = (double) RefMap.ntotRefMap * (double) nShifts * (double) nShifts *
								(((double) param.param_device.NumberPixels - (double) param.param_device.maxDisplaceCenter / 2.) * ((double) param.param_device.NumberPixels - (double) param.param_device.maxDisplaceCenter / 2.) * 2. + 8.) * (double) sizeof(myfloat_t) / compTime;
			const double nGBs2 = (double) RefMap.ntotRefMap * ((double) param.param_device.NumberPixels * (double) param.param_device.NumberPixels + 8.) * (double) sizeof(myfloat_t) / compTime;

			printf("Time Comparison %d %d: %f sec (%f GFlops, %f GB/s (cached), %f GB/s)\n", iOrient, iConv, compTime, nFlops / 1000000000., nGBs / 1000000000., nGBs2 / 1000000000.);
		}
	}
	//deallocating fftw_complex vector
	myfftw_free(proj_mapFFT);
	myfftw_free(conv_mapFFT);
	delete[] conv_map;

	deviceFinishRun();
	{
		const int count = omp_get_max_threads();
		for (int i = 0;i < count;i++)
		{
			myfftw_free(localCCT[i]);
			myfftw_free(lCC[i]);
		}
		delete[] localCCT;
		delete[] lCC;
	}

	// ************* Writing Out Probabilities ***************

	// *** Angular Probability ***

	// if(param.writeAngles){
	ofstream angProbfile;
	angProbfile.open ("ANG_PROB_iRefMap");
	// }

	ofstream outputProbFile;
	outputProbFile.open ("Output_Probabilities");
	for (int iRefMap = 0; iRefMap < RefMap.ntotRefMap; iRefMap ++)
	{
		// **** Total Probability ***
		bioem_Probability_map& pProbMap = pProb.getProbMap(iRefMap);

		outputProbFile << "RefMap " << iRefMap << " Probability  "  << log(pProbMap.Total) + pProbMap.Constoadd + 0.5 * log(M_PI) + (1 - param.param_device.Ntotpi * 0.5)*(log(2 * M_PI) + 1) + log(param.param_device.volu) << " Constant " << pProbMap.Constoadd  << "\n";

		outputProbFile << "RefMap " << iRefMap << " Maximizing Param: ";

		// *** Param that maximize probability****
		outputProbFile << (pProbMap.max_prob + 0.5 * log(M_PI) + (1 - param.param_device.Ntotpi * 0.5) * (log(2 * M_PI) + 1) + log(param.param_device.volu)) << " ";
		outputProbFile << param.angles[pProbMap.max_prob_orient].pos[0] << " ";
		outputProbFile << param.angles[pProbMap.max_prob_orient].pos[1] << " ";
		outputProbFile << param.angles[pProbMap.max_prob_orient].pos[2] << " ";
		outputProbFile << param.CtfParam[pProbMap.max_prob_conv].pos[0] << " ";
		outputProbFile << param.CtfParam[pProbMap.max_prob_conv].pos[1] << " ";
		outputProbFile << param.CtfParam[pProbMap.max_prob_conv].pos[2] << " ";
		outputProbFile << pProbMap.max_prob_cent_x << " ";
		outputProbFile << pProbMap.max_prob_cent_y;
		outputProbFile << "\n";

		// *** For individual files*** //angProbfile.open ("ANG_PROB_"iRefMap);

		if(param.writeAngles)
		{
			for (int iOrient = 0; iOrient < param.nTotGridAngles; iOrient++)
			{
				bioem_Probability_angle& pProbAngle = pProb.getProbAngle(iRefMap, iOrient);

				angProbfile << " " << iRefMap << " " << param.angles[iOrient].pos[0] << " " << param.angles[iOrient].pos[1] << " " << param.angles[iOrient].pos[2] << " " << log(pProbAngle.forAngles) + pProbAngle.ConstAngle + 0.5 * log(M_PI) + (1 - param.param_device.Ntotpi * 0.5)*(log(2 * M_PI) + 1) + log(param.param_device.volu) << " " << log(param.param_device.volu) << "\n";
			}
		}
	}

	angProbfile.close();
	outputProbFile.close();

	return(0);
}

int bioem::compareRefMaps(int iOrient, int iConv, const myfloat_t* conv_map, mycomplex_t* localmultFFT, myfloat_t sumC, myfloat_t sumsquareC, const int startMap)
{
	//***************************************************************************************
	//***** BioEM routine for comparing reference maps to convoluted maps *****
	if (FFTAlgo)
	{
		//With FFT Algorithm
		#pragma omp parallel for
		for (int iRefMap = startMap; iRefMap < RefMap.ntotRefMap; iRefMap ++)
		{
			const int num = omp_get_thread_num();
			calculateCCFFT(iRefMap, iOrient, iConv, sumC, sumsquareC, localmultFFT, localCCT[num], lCC[num]);
		}
	}
	else
	{
		//Without FFT Algorithm
		#pragma omp parallel for
		for (int iRefMap = startMap; iRefMap < RefMap.ntotRefMap; iRefMap ++)
		{
			compareRefMapShifted < -1 > (iRefMap, iOrient, iConv, conv_map, pProb, param.param_device, RefMap);
		}
	}
	return(0);
}

inline void bioem::calculateCCFFT(int iRefMap, int iOrient, int iConv, myfloat_t sumC, myfloat_t sumsquareC, mycomplex_t* localConvFFT, mycomplex_t* localCCT, myfloat_t* lCC)
{
	//***************************************************************************************
	//***** Calculating cross correlation in FFTALGOrithm *****

	const mycomplex_t* RefMapFFT = &RefMap.RefMapsFFT[iRefMap * param.FFTMapSize];
	for(int i = 0; i < param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D; i++)
	{
		localCCT[i][0] = localConvFFT[i][0] * RefMapFFT[i][0] + localConvFFT[i][1] * RefMapFFT[i][1];
		localCCT[i][1] = localConvFFT[i][1] * RefMapFFT[i][0] - localConvFFT[i][0] * RefMapFFT[i][1];
	}

	myfftw_execute_dft_c2r(param.fft_plan_c2r_backward, localCCT, lCC);

	doRefMapFFT(iRefMap, iOrient, iConv, lCC, sumC, sumsquareC, pProb, param.param_device, RefMap);
}

int bioem::createProjection(int iMap, mycomplex_t* mapFFT)
{
	// **************************************************************************************
	// ****  BioEM Create Projection routine in Euler angle predefined grid******************
	// ********************* and turns projection into Fourier space ************************
	// **************************************************************************************

	myfloat3_t RotatedPointsModel[Model.nPointsModel];
	myfloat_t rotmat[3][3];
	myfloat_t alpha, gam, beta;
	myfloat_t* localproj;

	localproj = lCC[omp_get_thread_num()];
	memset(localproj, 0, param.param_device.NumberPixels * param.param_device.NumberPixels * sizeof(*localproj));

	alpha = param.angles[iMap].pos[0];
	beta = param.angles[iMap].pos[1];
	gam = param.angles[iMap].pos[2];

	// **** To see how things are going: cout << "Id " << omp_get_thread_num() <<  " Angs: " << alpha << " " << beta << " " << gam << "\n"; ***

	// ********** Creat Rotation with pre-defiend grid of orientations**********
	rotmat[0][0] = cos(gam) * cos(alpha) - cos(beta) * sin(alpha) * sin(gam);
	rotmat[0][1] = cos(gam) * sin(alpha) + cos(beta) * cos(alpha) * sin(gam);
	rotmat[0][2] = sin(gam) * sin(beta);
	rotmat[1][0] = -sin(gam) * cos(alpha) - cos(beta) * sin(alpha) * cos(gam);
	rotmat[1][1] = -sin(gam) * sin(alpha) + cos(beta) * cos(alpha) * cos(gam);
	rotmat[1][2] = cos(gam) * sin(beta);
	rotmat[2][0] = sin(beta) * sin(alpha);
	rotmat[2][1] = -sin(beta) * cos(alpha);
	rotmat[2][2] = cos(beta);

	for(int n = 0; n < Model.nPointsModel; n++)
	{
		RotatedPointsModel[n].pos[0] = 0.0;
		RotatedPointsModel[n].pos[1] = 0.0;
		RotatedPointsModel[n].pos[2] = 0.0;
	}
	for(int n = 0; n < Model.nPointsModel; n++)
	{
		for(int k = 0; k < 3; k++)
		{
			for(int j = 0; j < 3; j++)
			{
				RotatedPointsModel[n].pos[k] += rotmat[k][j] * Model.PointsModel[n].pos[j];
			}
		}
	}

	int i, j;

	// ************ Projection over the Z axis********************
	for(int n = 0; n < Model.nPointsModel; n++)
	{
		//Getting pixel that represents coordinates & shifting the start at to Numpix/2,Numpix/2 )
		i = floor(RotatedPointsModel[n].pos[0] / param.pixelSize + (myfloat_t) param.param_device.NumberPixels / 2.0f + 0.5f);
		j = floor(RotatedPointsModel[n].pos[1] / param.pixelSize + (myfloat_t) param.param_device.NumberPixels / 2.0f + 0.5f);

		localproj[i * param.param_device.NumberPixels + j] += Model.densityPointsModel[n] / Model.NormDen;
	}

	// **** Output Just to check****
	if(iMap == 10)
	{
		ofstream myexamplemap;
		ofstream myexampleRot;
		myexamplemap.open ("MAP_i10");
		myexampleRot.open ("Rot_i10");
		myexamplemap << "ANGLES " << alpha << " " << beta << " " << gam << "\n";
		for(int k = 0; k < param.param_device.NumberPixels; k++)
		{
			for(int j = 0; j < param.param_device.NumberPixels; j++) myexamplemap << "\nMAP " << k << " " << j << " " << localproj[k * param.param_device.NumberPixels + j];
		}
		myexamplemap << " \n";
		for(int n = 0; n < Model.nPointsModel; n++)myexampleRot << "\nCOOR " << RotatedPointsModel[n].pos[0] << " " << RotatedPointsModel[n].pos[1] << " " << RotatedPointsModel[n].pos[2];
		myexamplemap.close();
		myexampleRot.close();
	}

	// ***** Converting projection to Fourier Space for Convolution later with kernel****
	// ********** Omp Critical is necessary with FFTW*******
	myfftw_execute_dft_r2c(param.fft_plan_r2c_forward, localproj, mapFFT);

	return(0);
}

int bioem::createConvolutedProjectionMap(int iMap, int iConv, mycomplex_t* lproj, myfloat_t* Mapconv, mycomplex_t* localmultFFT, myfloat_t& sumC, myfloat_t& sumsquareC)
{
	// **************************************************************************************
	// ****  BioEM Create Convoluted Projection Map routine, multiplies in Fourier **********
	// **************** calculated Projection with convoluted precalculated Kernel***********
	// *************** and Backtransforming it to real Space ********************************
	// **************************************************************************************

	mycomplex_t* tmp = localCCT[omp_get_thread_num()];

	// **** Multiplying FFTmap with corresponding kernel ****
	const mycomplex_t* refCTF = &param.refCTF[iConv * param.FFTMapSize];
	for(int i = 0; i < param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D; i++)
	{
		localmultFFT[i][0] = lproj[i][0] * refCTF[i][0] + lproj[i][1] * refCTF[i][1];
		localmultFFT[i][1] = lproj[i][1] * refCTF[i][0] - lproj[i][0] * refCTF[i][1];
		// cout << "GG " << i << " " << j << " " << refCTF[i][0] << " " << refCTF[i][1] <<" " <<lproj[i][0] <<" " <<lproj[i][1] << "\n";
	}

	//FFTW_C2R will destroy the input array, so we have to work on a copy here
	memcpy(tmp, localmultFFT, sizeof(mycomplex_t) * param.param_device.NumberPixels * param.param_device.NumberFFTPixels1D);

	// **** Bringing convoluted Map to real Space ****
	myfftw_execute_dft_c2r(param.fft_plan_c2r_backward, tmp, Mapconv);

	// *** Calculating Cross-correlations of cal-convoluted map with its self *****
	sumC = 0;
	sumsquareC = 0;
	for(int i = 0; i < param.param_device.NumberPixels * param.param_device.NumberPixels; i++)
	{
		sumC += Mapconv[i];
		sumsquareC += Mapconv[i] * Mapconv[i];
	}
	// *** The DTF gives an unnormalized value so have to divded by the total number of pixels in Fourier ***
	// Normalizing
	myfloat_t norm2 = (myfloat_t) (param.param_device.NumberPixels * param.param_device.NumberPixels);
	myfloat_t norm4 = norm2 * norm2;
	sumC = sumC / norm2;
	sumsquareC = sumsquareC / norm4;

	return(0);
}

int bioem::calcross_cor(myfloat_t* localmap, myfloat_t& sum, myfloat_t& sumsquare)
{
	// *********************** Routine to calculate Cross correlations***********************

	sum = 0.0;
	sumsquare = 0.0;
	for (int i = 0; i < param.param_device.NumberPixels; i++)
	{
		for (int j = 0; j < param.param_device.NumberPixels; j++)
		{
			// Calculate Sum of pixels
			sum += localmap[i * param.param_device.NumberPixels + j];
			// Calculate Sum of pixels squared
			sumsquare += localmap[i * param.param_device.NumberPixels + j] * localmap[i * param.param_device.NumberPixels + j];
		}
	}
	return(0);
}

int bioem::deviceInit()
{
	return(0);
}

int bioem::deviceStartRun()
{
	return(0);
}

int bioem::deviceFinishRun()
{
	return(0);
}

void* bioem::malloc_device_host(size_t size)
{
	return(mallocchk(size));
}

void bioem::free_device_host(void* ptr)
{
	free(ptr);
}