//OpenCV
#include <opencv2/opencv.hpp>
#include <stim/math/matrix.h>
#include <stim/math/constants.h>
#include <sstream>
#include "progress_thread.h"
#include <limits>
#include <chrono>

//LAPACKE support for Visual Studio
#include <complex>
#ifndef LAPACK_COMPLEX_CUSTOM
#define LAPACK_COMPLEX_CUSTOM
#define lapack_complex_float std::complex<float>
#define lapack_complex_double std::complex<double>
#endif
#include "lapacke.h"

class stim_EM : public cv::EM{
public:
	cv::vector<cv::Mat> getCovs() {
		return covs;
	}

	stim_EM(int nclusters = EM::DEFAULT_NCLUSTERS, int covMatType = EM::COV_MAT_DIAGONAL,
		const cv::TermCriteria& termCrit = cv::TermCriteria(cv::TermCriteria::COUNT + cv::TermCriteria::EPS,
			EM::DEFAULT_MAX_ITERS, FLT_EPSILON)) : cv::EM(nclusters, covMatType, termCrit) {

	}
};

//define a structure for a multi-class GMM
class GMM{
public:
	size_t K;						//number of Gaussians per class
	size_t F;						//number of features
	double t_gauss;

	std::vector< double > w;		//array of K weights for each Gaussian
	std::vector< stim::matrix< double > > mu;		//a vector storing K mean vectors of length F
	//std::vector< std::vector< gmm_mat > > sigma;	//(C x K) array of covariance matrices (F x F)
	std::vector< stim::matrix<double> > sigma;		//array of K (F x F) covariance matrices for each Gaussian
	std::vector< stim::matrix<double> > sigma_i;	//stores the inverse covariance matrices
	std::vector< double > sqrt_tau_sigma_det;		//stores sqrt(2*pi*|sigma|)

	void init(){
		w.resize(K);											//allocate space for weights
		mu.resize(K);						//allocate space for means
		sigma.resize(K);										//allocate space for each covariance matrix
		for (size_t k = 0; k < K; k++) {
			mu[k] = stim::matrix<double>(F, 1);
			sigma[k] = stim::matrix<double>(F, F);
		}
		t_gauss = 0;
	}

	//calculate the inverse sigma matrices
	void invert_sigmas() {
		sigma_i.resize(K);										//allocate space for K inverse matrices
		int *IPIV = (int*)malloc(sizeof(int) * F);				//allocate space for the row indices
		for (size_t k = 0; k < K; k++) {						//for each sigma matrix			
			sigma_i[k] = sigma[k];								//copy the covariance matrix
			LAPACKE_dgetrf(LAPACK_COL_MAJOR, (int)F, (int)F, sigma_i[k].data(), (int)F, IPIV);		//perform LU factorization
			LAPACKE_dgetri(LAPACK_COL_MAJOR, (int)F, sigma_i[k].data(), (int)F, IPIV);			//calculate matrix inverse
		}
		free(IPIV);
	}

	void calc_sqrt_tau_sigma_det() {
		sqrt_tau_sigma_det.resize(K);
		for (size_t k = 0; k < K; k++) {
			sqrt_tau_sigma_det[k] = sqrt(sigma[k].det() * stim::TAU);
		}
	}

	//initialize predictors for improving calculation of responses
	void init_predictors() {
		invert_sigmas();
		calc_sqrt_tau_sigma_det();
	}

	//calculate the value of a multi-variate gaussian distribution given a vector of means and a covariance matrix
	double mvgauss(stim::matrix<double> x, size_t k) {
		std::chrono::high_resolution_clock::time_point t0 = std::chrono::high_resolution_clock::now();
		stim::matrix<double> xmu = x - mu[k];
		stim::matrix<double> xmu_t = xmu.transpose();
		stim::matrix<double> xmu_t_sigma_i = xmu_t * sigma_i[k];
		stim::matrix<double> xmu_t_sigma_i_xmu = xmu_t_sigma_i * xmu;
		double a = -0.5 * xmu_t_sigma_i_xmu(0, 0);
		double numer = exp(a);
		stim::matrix<double> tau_sigma = sigma[k] * stim::TAU;
		double determinant = tau_sigma.det();
		double denom = sqrt(determinant);

		std::chrono::high_resolution_clock::time_point t1 = std::chrono::high_resolution_clock::now();
		t_gauss += std::chrono::duration_cast< std::chrono::duration<double> >(t1 - t0).count();
		return numer / denom;
	}

	double mvgauss(double* x, size_t k, double* scratch) {
		std::chrono::high_resolution_clock::time_point t0 = std::chrono::high_resolution_clock::now();
		for (size_t f = 0; f < F; f++)
			scratch[f] = x[f] - mu[k](f, 0);
		stim::matrix<double> xmu(F, 1, scratch);
		stim::matrix<double> xmu_t(1, F, scratch);
		stim::matrix<double> xmu_t_sigma_i = xmu_t * sigma_i[k];
		stim::matrix<double> xmu_t_sigma_i_xmu = xmu_t_sigma_i * xmu;
		double a = -0.5 * xmu_t_sigma_i_xmu(0, 0);
		double numer = exp(a);

		std::chrono::high_resolution_clock::time_point t1 = std::chrono::high_resolution_clock::now();
		t_gauss += std::chrono::duration_cast< std::chrono::duration<double> >(t1 - t0).count();

		return numer / sqrt_tau_sigma_det[k];
	}

	/// returns the probability density of the membership of v in all K clusters
	std::vector<double> G(stim::matrix<double> x) {
		std::vector<double> result(K);					//allocate space for all K probabilities
		for (size_t k = 0; k < K; k++) {				//for each gaussian
			result[k] = mvgauss(x, k);
		}
		return result;
	}

	/// Calculate the response to x among all K clusters given pointers to pre-allocated arrays
	void G(double* x, double* r) {
		double* scratch = (double*)malloc(F * sizeof(double));
		for (size_t k = 0; k < K; k++) {				//for each gaussian
			r[k] = mvgauss(x, k, scratch);
		}
		free(scratch);
	}

	/// Return the cluster most closely corresponding to the input vector x
	size_t get_cluster(stim::matrix<double> x) {
		size_t cluster;									//stores the cluster ID
		std::vector<double> posteriors = G(x);
		double largest = posteriors[0];
		for (size_t k = 0; k < K; k++) {
			if (posteriors[k] >= largest) {
				largest = posteriors[k];
				cluster = k;
			}
		}
		return cluster;
	}

	///Return the posterior probability of the vector x based on the current Gaussian mixture model
	double P(stim::matrix<double> x) {
		std::vector<double> posteriors = G(x);
		double p = 0;
		for (size_t k = 0; k < K; k++) {
			p += w[k] * posteriors[k];					//calculate the weighted sum of all Gaussian functions
		}
		return p;
	}

	double P(double* x) {
		double* posteriors = (double*)malloc(K * sizeof(double));
		G(x, posteriors);
		double p = 0;
		for (size_t k = 0; k < K; k++) {
			p += w[k] * posteriors[k];					//calculate the weighted sum of all Gaussian functions
		}
		return p;
	}

public:

	GMM() {
		K = 0;
		F = 0;
	}

	GMM(size_t clusters, size_t features){
		K = clusters;
		F = features;
		init();
	}

	void set(const cv::Mat weights, const cv::Mat means, const std::vector<cv::Mat> cov) {
		for (size_t k = 0; k < K; k++)
			w[k] = weights.at<double>(0, (int)k);

		for (size_t k = 0; k < K; k++)
			for (size_t f = 0; f < F; f++)
				mu[k](f, 0) = means.at<double>((int)k, (int)f);

		for (size_t k = 0; k < K; k++) {
			for (size_t fi = 0; fi < F; fi++) {
				for (size_t fj = 0; fj < F; fj++) {
					sigma[k](fi, fj) = cov[k].at<double>((int)fi, (int)fj);
				}
			}
		}
		init_predictors();											//calculate the inverse covariance matrices				
	}	

	std::string str() {
		std::stringstream ss;
		ss << "weights:" << std::endl;
		for (size_t k = 0; k < K; k++)
			ss << "     " << w[k] << std::endl;

		ss << std::endl << "centers:" << std::endl;
		for (size_t k = 0; k < K; k++)
			ss << mu[k].toStr() << std::endl;

		ss << std::endl << "covariances:" << std::endl;
		for (size_t k = 0; k < K; k++)
			ss << sigma[k].toStr() << std::endl;
		return ss.str();
	}

	void save(std::ostream& out) {
		out << K << std::endl;							//save the number of clusters
		out << F << std::endl;							//save the number of features
		for (size_t k = 0; k < K; k++)
			out << std::fixed << w[k] << std::endl;
		for (size_t k = 0; k < K; k++)
			out << mu[k].csv() << std::endl;

		for (size_t k = 0; k < K; k++)
			out << sigma[k].csv() << std::endl;
	}

	void save(std::string filename) {
		std::ofstream outfile(filename);
		int digits = std::numeric_limits<double>::max_digits10;
		outfile.precision(digits);
		save(outfile);
		outfile.close();
	}

	//load a GMM
	void load(std::istream& in) {
		in >> K;										//load the number of clusters
		in >> F;										//load the number of features
		init();
		for (size_t k = 0; k < K; k++)
			in >> w[k];
		for (size_t k = 0; k < K; k++)
			mu[k].csv(in);

		for (size_t k = 0; k < K; k++)
			sigma[k].csv(in);
		init_predictors();								//calculate the inverse covariance matrices
	}

	void load(std::string filename) {
		std::ifstream infile(filename);
		load(infile);
		infile.close();
	}

};

/// Multi-class supervised GMM
class multiGMM {
public:
	size_t C;										//number of classes

	std::vector<GMM> gmms;							//vector of Gaussian Mixture models

	/// Generate an empty GMM for each class
	void init() {
		for (size_t c = 0; c < C; c++) {
			gmms.resize(C);
		}
	}

	multiGMM(size_t classes) {
		C = classes;								//store the number of classes
		init();
	}

	//get the class that most likely corresponds to x
	size_t get_class(stim::matrix<double> x) {
		
		double p0;
		size_t c_p = 0;								//stores the most likely class label
		double p = gmms[0].P(x);					//get the posterior probability of class 0
		for (size_t c = 1; c < C; c++) {			//for each class
			p0 = gmms[c].P(x);						//get the posterior probability of membership given x
			if (p0 > p) {							//if the new class is most likely
				p = p0;								//update the maximum probability
				c_p = c;							//update the class ID
			}
		}
		return c_p;
	}

	void save(std::string filename) {
		std::ofstream outfile(filename);			//open an output file stream
		if (outfile) {
			int digits = std::numeric_limits<double>::max_digits10;
			outfile.precision(digits);
			outfile << C << std::endl;					//save the number of classes
			for (size_t c = 0; c < C; c++) {
				gmms[c].save(outfile);					//save each individual GMM
			}
			outfile.close();
		}
		else {
			std::cout << "ERROR creating GMM file " << filename << std::endl;
			exit(1);
		}
	}

	bool load(std::string filename) {
		std::ifstream infile(filename);				//open the input file
		if (!infile) return false;
		infile >> C;								//load the number of classes
		gmms.resize(C);								//resize the GMM array to match the number of classes

		for (size_t c = 0; c < C; c++)				//load each GMM (one per class)
			gmms[c].load(infile);
		return true;
	}
};

/// trains a single Gaussian Mixture model using expectation maximization in OpenCV
GMM train_gmm(cv::Mat &F, int k, int attempts, int iters, double epsilon){
	
	GMM new_gmm(k, F.cols);							//create a new GMM classifier
	stim_EM em(k, cv::EM::COV_MAT_DIAGONAL, cv::TermCriteria( CV_TERMCRIT_EPS+CV_TERMCRIT_ITER, iters, epsilon));
	if(!em.train(F)) {
		std::cout << "ERROR training GMM" << std::endl;
		exit(1);
	}
	size_t nc = em.get<int>("nclusters");
	
	cv::Mat output;
	cv::Mat means = em.get<cv::Mat>("means");
	cv::Mat weights = em.get<cv::Mat>("weights");
	cv::vector<cv::Mat> covs = em.getCovs();
	new_gmm.set(weights, means, covs);
	return new_gmm;
}

//Predict a set of classes based on given centroid vectors
std::vector< stim::image<unsigned char> > predict_gmm(stim::envi* E, multiGMM* gmm, std::vector< stim::image<float> >& responses, unsigned char* MASK = NULL){
	size_t nC = gmm->C;											//get the number of classes
	if (nC == 1)
		nC = gmm->gmms[0].K;									//if there is only one GMM, classify based on clusters
	size_t X = E->header.samples;								//store ENVI file size parameters
	size_t Y = E->header.lines;
	size_t B = E->header.bands;
	size_t XY = E->header.samples * E->header.lines;

	size_t tP = 0;												//calculate the total number of pixels
	if(MASK){
		for(size_t xy = 0; xy < XY; xy++){
			if(MASK[xy]) tP++;
		}
	}
	else
		tP = X * Y;

	std::vector< stim::image<unsigned char> > C;				//create an array of mask images
	C.resize(nC);
	responses.resize(nC);										//allocate space for the response images
	
	for(size_t c = 0; c < nC; c++){								//for each class mask
		C[c] = stim::image<unsigned char>(X, Y, 1);				//allocate space for the mask
		memset(C[c].data(), 0, X * Y * sizeof(unsigned char));	//initialize all of the pixels to zero
		responses[c] = stim::image<float>(X, Y, 1);				//allocate space for the response image
		memset(responses[c].data(), 0, X * Y * sizeof(float));	//initialize the response image to zero
	}

	double progress = 0;										//initialize the progress bar variable
	std::thread t1(progressbar_thread, &progress);				//start the progress bar thread

	size_t t = 0;
	double* spectrum = (double*)malloc(sizeof(double) * B);					//allocate space to hold a spectrum
	double gm, maxgm;
	size_t maxc;
	for(size_t p = 0; p < XY; p++){										//for each pixel
		if(!MASK || MASK[p] > 0){
			E->spectrum<double>(spectrum, p);							//get the spectrum at pixel p
			maxc = 0;
			for (size_t c = 0; c < nC; c++) {
				gm = gmm->gmms[c].P(spectrum);								//evaluate the posterior for class c
				responses[c].data()[p] = (float)gm;

				if (c == 0) maxgm = gm;
				else if (gm > maxgm) {
					maxgm = gm;
					maxc = c;
				}
			}
			C[maxc].data()[p] = 255;
			t++;
			progress = (double)(t+1) / (double)(tP) * 100.0;		//update the progress bar variable
		}
	}
	t1.join();												//finish the progress bar thread

	for (size_t c = 0; c < gmm->gmms.size(); c++) {
		std::cout << "gauss-time (" << c << "): " << gmm->gmms[c].t_gauss << std::endl;
	}
	
	return C;
}