#include "linalg.h"

#include <iostream>
#include <stim/envi/envi.h>
#include "stim/envi/envi_header.h"
#include "stim/parser/arguments.h"
#include <stim/image/image.h>
#include <stim/math/matrix.h>
#include <time.h>
#include <fstream>
#include <thread>
#include "stim/envi/bil.h"

//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"




void progress_thread_envi(stim::envi* e);		//progress bar threaded function


extern unsigned char* MASK;					//pointer to an input file mask if one is provided
extern stim::envi ENVI;								//input ENVI file to be processed
extern unsigned long long X, Y, B;						//registers to quickly access the image size



														///comparison function to compare two values corresponding to two memory addresses a and b
														///is used in qsort function to sort elements of an array (ex.eigenvalues) and getting the corresponding indices.
int cmpfunc(const void * a, const void * b)
{

	double *x = *(double **)a; double *y = *(double **)b;

	return *x > *y ? 1 : -1;

}


/// multiply matrix A (nxm) by matrix B (mxp)
/// @param MatA is a pointer to the matrix A
/// @param MatB is a pointer to the matrix B
/// @param MatC is a pointer to pre-allocated memory of size [n * p * sizeof(double)] that stores the result of multiplication.
/// @param rowsX and @param columnsX are integers refering to the number of rows and columns of the matrix X respectively.
/// Note that columnA = rowsB
void multiply_2matrices(const double* MatA, int rowsA, int columnsA, const double* MatB, int rowsB, int columnsB, double* MatC, int rowsC, int columnsC) {
	for (int i = 0; i < rowsA; i++) {
		for (int j = 0; j < columnsB; j++) {
			double sum = 0.0;
			for (int k = 0; k < rowsB; k++)
				sum = sum + MatA[i * columnsA + k] * MatB[k * columnsB + j];
			MatC[i * columnsC + j] = sum;                                                                //resulting matrix C is [nxp]
		}
	}
}



void mnf(std::string outfile, int keptComponents, std::string NoiseFractions, int cuda_device){
	///throw an Error in case of invalid values for the number of keptComponents
	if (keptComponents <= 0 || keptComponents > B) {
		std::cout << "ERROR: keptComponents must be in the range [0  " <<B<<"] (number of bands)"<< std::endl;
		exit(1);
	}
	//throw an error if the ENVI file isn't BIP
	if (ENVI.header.interleave != stim::envi_header::BIP) {
		std::cout << "ERROR: MNF basis calculation requires a BIP file to make covariance calculation practical" << std::endl;
		exit(1);
	}

	unsigned long long XY = X * Y;									//number of pixels in the image
	double* mu = (double*)malloc(sizeof(double) * B);               //allocate space for the mean

	std::thread t1(progress_thread_envi, &ENVI);					//start the progress bar thread
	std::cout << "Calculating mean spectrum..." << std::endl;
	ENVI.mean_spectrum(mu, NULL, MASK, true);									//calculate average spectrum

	for (size_t b = 0; b < B; b++) {						//for each band, test for non-finite values
		if (!std::isfinite(mu[b]))							//C++11 implementation
			std::cout << "WARNING: the mean spectrum contains non-finite values in band " << b << ". Use --mask-finite to create a mask of finite values." << std::endl;
	}
	t1.join();																//wait for the progress bar thread to finish (it probably already is)

	stim::matrix<double> cov(B, B);											//allocate a double-precision covariance matrix
	t1 = std::thread(progress_thread_envi, &ENVI);							//start the progress bar thread
	ENVI.co_matrix(cov.data(), mu, MASK, cuda_device, true);				//calculate covariance matrix
	t1.join();																//wait for the progress bar thread to finish (it probably already is)

	stim::matrix<double> ncov(B, B);										//allocate a double-precision noise covariance matrix
	t1 = std::thread(progress_thread_envi, &ENVI);							//start the progress bar thread
	ENVI.coNoise_matrix(ncov.data(), mu, MASK, cuda_device, true);			//calculate covariance of noise matrix
	t1.join();																//wait for the progress bar thread to finish (it probably already is)


	std::cout << std::endl << "Calculating the inverse covariance matrix S^(-1)...";
	int *IPIV = (int*)malloc(sizeof(int) * B);
	
	//LAPACKE_dgetrf(LAPACK_COL_MAJOR, (int)B, (int)B, cov.data(), (int)B, IPIV);		//perform LU factorization
	//LAPACKE_dgetri(LAPACK_COL_MAJOR, (int)B, cov.data(), (int)B, IPIV);			//calculate matrix inverse
	
	// REPLACED THE ABOVE LAPACKE functions with new CLAPACK functions in linalg.cpp
	LINALG_dgetrf((int)B, (int)B, cov.data(), (int)B, IPIV);		//perform LU factorization
	LINALG_dgetri((int)B, cov.data(), (int)B, IPIV);			//calculate matrix inverse

	free(IPIV);
	std::cout << "done." << std::endl;

	std::cout << std::endl << "Calculating Q = S_n * S^(-1)...";							//calculate the Q matrix
	stim::matrix<double> Q = ncov * cov;
	std::cout << "done." << std::endl;

	free(mu);          //free memory for the mean spectrum

	///computing the eigenvalues and left eigenvectors of matrix Qmat.
	double *EigenvaluesReal = (double*)malloc(B * sizeof(double));       //allocate space for the real part of eigenvalues
	double *EigenvaluesIm = (double*)malloc(B * sizeof(double));         //allocate space for the imaginary part of eigenvalues

	std::cout << std::endl << "Calculating left eigenvectors X * A = v * X...";
	stim::matrix<double> ev_left(B, B);
	//LAPACKE_dgeev(LAPACK_COL_MAJOR, 'V', 'N', (int)B, Q.data(), (int)B, EigenvaluesReal, EigenvaluesIm, ev_left.data(), (int)B, 0, (int)B);       //perform eigenvalue decomposition

	// REPLACED THE ABOVE LAPACKE functions with new CLAPACK functions in linalg.cpp
	LINALG_dgeev('V', 'N', (int)B, Q.data(), (int)B, EigenvaluesReal, EigenvaluesIm, ev_left.data(), (int)B, 0, (int)B);       //perform eigenvalue decomposition

	std::cout << "done." << std::endl;

	std::cout << std::endl << "Sorting eigenvectors...";
	///SORTING EIGENVECTORS CORRESPONDING TO THEIR EIGENVALUES
	//EigenvaluesReal is the original array and "pointers" is the array of corresponding pointers, whose size is the size of the original array.
	//Let cmp be the comparison function for qsort().
	double **pointers = (double**)malloc(B * sizeof(double*));
	for (int i = 0; i < B; i++) {
		pointers[i] = &EigenvaluesReal[i];         // assign the address of integer.
	}

	///SORTING EIGENVALUES AND GETTING THE INDICES
	qsort(pointers, B, sizeof(pointers[0]), cmpfunc);
	size_t *EigenSortedIndices = (size_t*)malloc(B * sizeof(size_t));               //allocate space for the Indices of sorted eigenvalues
	for (size_t i = 0; i < B; ++i)
		EigenSortedIndices[i] = pointers[i] - EigenvaluesReal;                      //Calculate indices of the sorted eigenvalues
	
	stim::matrix<double> As = ev_left.sort_cols(EigenSortedIndices);
	stim::matrix<double> Ast = As.transpose();
	std::cout << "done." << std::endl;

	free(EigenvaluesReal);
	free(EigenvaluesIm);
	free(EigenSortedIndices);

	std::cout << std::endl << "Removing noise components...";
	///Noise removal and then backward transformation (builing one general/final basis matrix to be multiplied wih the raw data)
	//General Transformation/Final Basis Matrix = gtm = (inv(As).') * (Keye) * (Ast) ;
	//Note that the forward transform is done using transpose of mnf basis (Ast).However, for inverse transform, we use transpose of mnf basis inversed.

	//creating a BxB matrix with the number of 'keptComponents' 1 on the diagonal and 0 elsewhere.
	stim::matrix<double> K = stim::matrix<double>::I(B);							//create an identity matrix
	for (size_t i = keptComponents; i < B; i++)	K(i, i) = 0.0;						//assign zeros to components that will not be used

	stim::matrix<double> K_Ast = K * Ast;

	//calculate inverse of As matrix
	int *IPIV2 = (int*)malloc(sizeof(int) * B);
	//LAPACKE_dgetrf(LAPACK_COL_MAJOR, (int)B, (int)B, As.data(), (int)B, IPIV2);
	//LAPACKE_dgetri(LAPACK_COL_MAJOR, (int)B, As.data(), (int)B, IPIV2);

	// REPLACED THE ABOVE LAPACKE functions with new CLAPACK functions in linalg.cpp
	LINALG_dgetrf((int)B, (int)B, As.data(), (int)B, IPIV2);
	LINALG_dgetri((int)B, As.data(), (int)B, IPIV2);

	free(IPIV2);

	//calculate transpose of As inversed matrix (transpose of transformation matrix inveresed  ) .
	stim::matrix<double> Asit = As.transpose();

	//multiplying As (inversed transposed matrix) by KeyeAst : (inv(As)).' * (KeyeAst) .
	stim::matrix<double> G = Asit * K_Ast;
	stim::matrix<double> G_t = G.transpose();
	std::cout << "done." << std::endl;

	std::ofstream basisfile(outfile);
	basisfile << 0;
	for (size_t i = 1; i < B; i++) basisfile << ", " << 0;
	basisfile << std::endl;
	basisfile << G.csv();
	basisfile.close();

	//output noise fractions
	if (NoiseFractions.size() > 0) {													//if a file name is specified for the noise fractions, save them
		std::cout << std::endl << "Outputing noise fractions in a CSV file (infile-NoiseFractions)..." << std::endl;
		//outputing the sorted Eigenvalues in a CSV file (Noise Fractions)
		std::ofstream csv(NoiseFractions.c_str());								//open a CSV file to write the sorted Eigenvalues
		csv << EigenSortedIndices[0] << ", " << EigenvaluesReal[EigenSortedIndices[0]];														//output the first variable
		for (unsigned long long b = 1; b < B; b++)
			csv << "\n" << EigenSortedIndices[b] << ", " << EigenvaluesReal[EigenSortedIndices[b]];												//output the next variable
		csv.close();
	}
}	//end MNF