#include <stdio.h>
#include <stdlib.h>
#include <math.h>
typedef double _Complex DCOMPLEX;
#define MKL_Complex16 DCOMPLEX
#include <mkl.h>
#include <mpi.h>
#include <omp.h>
#include <assert.h>
#include <time.h>
#include <complex.h>
#include <mkl_blacs.h>
#include <mkl_scalapack.h>

#ifndef M_PI
    #define M_PI 3.14159265358979323846
#endif


/* constants */
const double e_Charge = 1.60217662e-19; //electron charge
const double e_Mass = 9.10938356e-31; // electron mass
const double h = 6.62607015e-34; // Planck's constant
const double hbar = 1.05457182e-34; // reduced Planck's constant
const double kb = 1.38064852e-23;

void Test_QW_PWEM(); // Tests PWEM using complex matrix eigensolver
int INDXL2G(int iloc, int nb, int iproc, int isrcproc, int nprocs);
double complex Hmm(int m0, int m1, int elec, double field); // Returns matrix elements (complex)
double Chi(int m); // Returns chi function calculation
double G(int m); // Returns wavenumber
double Calculate_Ternary_Parameter(double x, double E1, double E2, double C); //Calculating material parameter
double Calculate_Band_Gap(double T, double Eg0, double alpha, double beta); // Calculating band gap of material

 /* mpi */
int myid, nprocs;
MKL_INT na_cols, na_rows;
MKL_INT np_cols, np_rows;
MKL_INT my_prow, my_pcol;
MKL_INT mpi_comm;

/* blacs & ScaLapack*/
MKL_INT my_blacs_ctxt, sc_desc[9], sc_descz[9], info;
MKL_INT nblk = 32;


/* parameters */
int order;
double L = 7e-9;
double SL = 21e-9;

double E_Mass_GaAs;
double E_Mass_AlAs;
double E_Mass0;
double E_Mass1;
double GaAs_Potential;
double AlAs_Potential;
double AlGaAs_Potential;
double Zone_Potential_Electrons;
double E_field = 10e6;


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

	order = atoi(argv[1]);

	MPI_Init(&argc, &argv);
	MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
	MPI_Comm_rank(MPI_COMM_WORLD, &myid);

	for (np_cols = (int) sqrt((double) nprocs); np_cols > 1; np_cols--) {
     	if (nprocs % np_cols == 0) {
       		break;
     	}
   	}

	np_rows = nprocs/np_cols;

	/* set up blacs */
	MKL_INT globalComm = MPI_COMM_WORLD;
	MKL_INT CommWhat = 0;
	blacs_get(&globalComm, &CommWhat, &my_blacs_ctxt);

	void* tag_ub;
	int flag;
	MPI_Comm_get_attr(MPI_COMM_WORLD, MPI_TAG_UB, &tag_ub, &flag);

	MKL_INT ID_range_set[2];
	ID_range_set[0] = 0;
	ID_range_set[1] = 262143;
	CommWhat = 1;
	//blacs_set(&globalComm, &CommWhat, ID_range_set);

	blacs_gridinit(&my_blacs_ctxt, "C", &np_rows, &np_cols);

	MKL_INT ID_range[2];
	CommWhat = 1;
	blacs_get(&my_blacs_ctxt, &CommWhat, ID_range);

	blacs_gridinfo(&my_blacs_ctxt, &np_rows, &np_cols, &my_prow, &my_pcol);

	#ifdef MKL_ILP64
	{

		printf("Process %d, ID_range min = %lld\n", myid, ID_range[0]);
		printf("Process %d, ID_range max = %lld\n", myid, ID_range[1]);

	}
	#else
	{

		printf("Process %d, ID_range min = %d\n", myid, ID_range[0]);
		printf("Process %d, ID_range max = %d\n", myid, ID_range[1]);

	}
	#endif

	if (flag) printf("Process %d, tag ub = %d\n", myid, *(int*)tag_ub);

	
	/* setting parameters */
	E_Mass_GaAs = 0.065 * e_Mass;
	E_Mass_AlAs = 0.15 * e_Mass;
	E_Mass0 = Calculate_Ternary_Parameter(0.3, E_Mass_GaAs, E_Mass_AlAs, 0);
	E_Mass1 = E_Mass_GaAs;
	GaAs_Potential = ((-6.92 + (0.34 / 3)) * e_Charge) + Calculate_Band_Gap(300, (1.519) * e_Charge, 0.5405e-3, 204);
	AlAs_Potential = Calculate_Band_Gap(300, (-4.27) * e_Charge, 0.885e-3, 530);
	AlGaAs_Potential = GaAs_Potential + (0.85 * 0.3 * e_Charge);
	Zone_Potential_Electrons = AlGaAs_Potential - GaAs_Potential;


	Test_QW_PWEM();


	blacs_gridexit(&my_blacs_ctxt);
	MPI_Finalize();

	return 0;
}


void Test_QW_PWEM() {
	MKL_INT N = order; // matrix order
	MKL_INT m = 10; // eigenvalues to compute


	MKL_INT f = 0;
	na_rows = NUMROC(&N, &nblk, &my_prow, &f, &np_rows);
	na_cols = NUMROC(&N, &nblk, &my_pcol, &f, &np_cols);	
	
	MKL_INT na_rowsz = NUMROC(&N, &nblk, &my_prow, &f, &np_rows);
	MKL_INT na_colsz = NUMROC(&m, &nblk, &my_pcol, &f, &np_cols);
	
	DESCINIT(sc_desc, &N, &N, &nblk, &nblk, &f, &f, &my_blacs_ctxt, &na_rows, &info);
	DESCINIT(sc_descz, &N, &N, &nblk, &nblk, &f, &f, &my_blacs_ctxt, &na_rows, &info);


	double complex *values, *V;
	double *W;
	values = calloc(na_rows * na_cols, sizeof(double complex));
	W = calloc(N, sizeof(double));
	V = calloc(na_rowsz * na_colsz, sizeof(double complex));


	for (MKL_INT m0 = 0; m0 < na_rows; m0++) 
	{

		for (MKL_INT m1 = 0; m1 < na_cols; m1++) 
		{

			values[(m1 * na_rows) + m0] = Hmm((-N / 2) + INDXL2G(m0, nblk, my_prow, 0, np_rows), (-N / 2) + INDXL2G(m1, nblk, my_pcol, 0, np_cols), 0, E_field) / e_Charge; // matrix elements
			
		}
	
	}
	
	char jobz = 'V'; // Computes eigenvalues and eigenvectors
	char range = 'I'; // The routine computes eigenvalues with indices il to iu
	char uplo = 'L'; // Matrix stores lower triangular part
	
	MKL_INT ia = 1;
	MKL_INT ja = 1;
	
	double vl = 0; // If range = 'V', vl is the eigenvalue search lower bound
	double vu = 0; // If range = 'V', vr is the eigenvalue search upper bound
	
	MKL_INT il = 1; // Eigenvalue search indice lower bound
	MKL_INT iu = il + m - 1; // Eigenvalue search indice upper bound
	
	MKL_INT iz = 1;
	MKL_INT jz = 1;
	
	MKL_INT lwork = -1;
	double complex work_dum;
	double complex *work;
	
	MKL_INT liwork = -1;
	MKL_INT iwork_dum;
	MKL_INT *iwork;
	
	MKL_INT lrwork = -1;
	double rwork_dum;
	double *rwork;
	
	MKL_INT nz;
	
	MKL_INT info;
	
	double abstol = 2*pdlamch(&my_blacs_ctxt, "S"); // High accuracy setting for error tolerance
	
	double orfac = 0;
	
	MKL_INT *ifail = calloc(N, sizeof(MKL_INT));
	
	MKL_INT *iclustr = calloc(2*np_rows*np_cols, sizeof(MKL_INT));
	
	double *gap = calloc(np_rows*np_cols, sizeof(double));

	
	//pzheevx( &jobz, &range, &uplo, &N, values, &ia, &ja, sc_desc, &vl, &vu, &il, &iu, &abstol, &m, &nz, W, &orfac, V, &iz, &jz, sc_descz, &work_dum, &lwork, &rwork_dum, &lrwork, &iwork_dum, &liwork, ifail, iclustr, gap, &info );

	pzheevr( &jobz, &range, &uplo, &N, values, &ia, &ja, sc_desc, &vl, &vu, &il, &iu, &m, &nz, W, V, &iz, &jz, sc_descz, &work_dum, &lwork, &rwork_dum, &lrwork, &iwork_dum, &liwork, &info );
	
	
	lwork = (MKL_INT)(creal(work_dum) + 0.5);
	work = calloc(lwork, sizeof(double complex));
	
	lrwork = (MKL_INT)(rwork_dum + 0.5);
	rwork = calloc(lrwork, sizeof(double));
	
	liwork = iwork_dum;
	iwork = calloc(liwork, sizeof(MKL_INT));

	
	//pzheevx( &jobz, &range, &uplo, &N, values, &ia, &ja, sc_desc, &vl, &vu, &il, &iu, &abstol, &m, &nz, W, &orfac, V, &iz, &jz, sc_descz, work, &lwork, rwork, &lrwork, iwork, &liwork, ifail, iclustr, gap, &info );
	
	pzheevr( &jobz, &range, &uplo, &N, values, &ia, &ja, sc_desc, &vl, &vu, &il, &iu, &m, &nz, W, V, &iz, &jz, sc_descz, work, &lwork, rwork, &lrwork, iwork, &liwork, &info );
	
	if (myid == 0)
	{
		#ifdef MKL_ILP64
		{

			printf("\nInfo = %lld\n\n", info);

		}
		#else
		{

			printf("Eigensolver info = %d\n\n", info);

		}
		#endif
		
		for (int i = 0; i < m; i++)
		{

			printf("Energy %d = %.11lf eV.\n", i + 1, W[i]);
			
		}
	
	} 

	free(values);
	free(W);
	free(V);
	free(work);
	free(iwork);
	free(rwork);
	free(ifail);
	free(iclustr);
	free(gap);

}

int INDXL2G(int iloc, int nb, int iproc, int isrcproc, int nprocs)
{
	
	return (nprocs * nb * ((iloc) / nb) + ((iloc) % nb) + ((nprocs + iproc - isrcproc) % nprocs) * nb);

}

double complex Hmm(int m0, int m1, int elec, double field ) 
{

	double complex val = 0;	

	if (!elec) 
	{

		if (m0 == m1)
		{

			val += (hbar * hbar * G(m0) * G(m1) / (2 * E_Mass0)) + Zone_Potential_Electrons;

		}
		else
		{

			val -= I * e_Charge * field * (1 - (2 * (abs((m0 - m1)) % 2))) / (G(m0) - G(m1));
			
		}

		val += Chi(m1 - m0) * ((hbar * hbar * G(m0) * G(m1) * ((1 / E_Mass1) - (1 / E_Mass0)) / (2)) - Zone_Potential_Electrons);

	}
	return val;

}

double Chi(int m) 
{

	if (m == 0)
	{

		return L / SL;

	}
	else
	{

		return L * sin( M_PI * m * L / SL ) / (M_PI * m * L);

	}

}

double Calculate_Ternary_Parameter(double x, double E1, double E2, double C) 
{
	
	return ((1 - x) * E1) + (x * E2) - (x * (1 - x) * C);

}

double Calculate_Band_Gap(double T, double Eg0, double alpha, double beta) 
{

	return Eg0 - ((alpha * T * T) * (e_Charge) / (T + beta));
	
}

double G(int m) 
{

	return 2 * M_PI * m / SL;

}