#include <petscksp.h>

static char help[] = "";

const int send_tag = 100;

#undef __FUNCT__
#define __FUNCT__ "solveLinSystem"
double solveLinSystem(KSP ksp, Vec rhs, PetscInt xSize, int retValInd){
	double retVal;
	PetscErrorCode ierr;	
	ierr = KSPSolve(ksp, rhs, x);
	VecGetValues(x, 1, &retValInd, &retVal);
	return retVal;
}

#undef __FUNCT__
#define __FUNCT__ "loadBinMatrix"
void loadMatrixBin(char filein[], Mat *A, const MatType mattype, MPI_Comm comm){
	PetscErrorCode ierr;
	PetscViewer    view;
	PetscViewerBinaryOpen(comm, filein, FILE_MODE_READ, &view);
	MatCreate(comm, A);
	ierr = MatSetType(*A, mattype);		
	MatLoad(*A,view);  
	PetscViewerDestroy(&view);
}

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc,char **args){
	Mat A, B;
	KSP ksp;
	char *fileLocA, *fileLocB;
	int rhs_size, x_size;
	double trace = 0.0;
	int myrank, noProcesses;
	PetscLogDouble start, end, dif, tmp_start, tmp_end, loadTime;
	
	fileLocA = "./calcTraceExp/A10000";
	fileLocB = "./calcTraceExp/B10000";
	
	PetscInitialize(&argc,&args,(char *)0,help);
	
	PetscGetTime(&tmp_start);	
	
	MPI_Comm_size(MPI_COMM_WORLD, &noProcesses);
	MPI_Comm_rank(MPI_COMM_WORLD, &myrank);
	
	loadMatrixBin(fileLocA, &A, MATSEQAIJ, PETSC_COMM_SELF);
	loadMatrixBin(fileLocB, &B, MATMPIDENSE, PETSC_COMM_WORLD);
	
	PetscGetTime(&tmp_end);	
	loadTime = tmp_end - tmp_start;
	PetscPrintf(PETSC_COMM_WORLD, "Loading took = %G\n", loadTime);
	
	PetscGetTime(&start);	
	
	KSPCreate(PETSC_COMM_SELF, &ksp); //Create linear solver context
	KSPSetOperators(ksp, A, A, SAME_PRECONDITIONER);
	#ifdef PETSC_HAVE_MUMPS
		PC pc;
		PetscInt ival,icntl;
		
		PetscPrintf(PETSC_COMM_WORLD, "Using Cholesky factorization\n");
		MatSetOption(A, MAT_SPD, PETSC_TRUE);
		KSPSetType(ksp, KSPPREONLY);
		KSPGetPC(ksp,&pc);
		PCSetType(pc,PCCHOLESKY);
		PCFactorSetMatSolverPackage(pc, MATSOLVERMUMPS);		
		PCFactorSetUpMatSolverPackage(pc);
		
		KSPSetFromOptions(ksp);
		KSPSetUp(ksp);
	#endif  
	
	// matrix B is transposed because its faster to access rows then columns
	MatTranspose(B, MAT_REUSE_MATRIX, &B);
	MatGetSize(B, PETSC_NULL, &rhs_size);	
	MatGetSize(A, PETSC_NULL, &x_size);	
	
	// Each process calculates A\b_i, where b_i is i-th row from B. Each process does this just for the rows from 
	// B it owns
	
	// Eventualno posto se faktorizacija radi samo pri prvom pozivu mozda jednom pozvati A\b pa onda slati ksp
	
	int startRow, endRow, rowInd;
	MatGetOwnershipRange(B, &startRow, &endRow);
	for(rowInd = startRow; rowInd < endRow; rowInd++){
		//PetscGetTime(&tmp_start);
	
		Vec rhs, x;
		MatGetVecs(A, &x, &rhs);
		
		PetscInt ncols;
		const PetscInt *cols;
		const PetscScalar *vals;
		MatGetRow(B, rowInd, &ncols, &cols, &vals);
		VecSetValues(rhs, ncols, cols, vals, INSERT_VALUES); 
		MatRestoreRow(B, rowInd, &ncols, &cols, &vals);
		
		trace += solveLinSystem(ksp, rhs, x, rowInd);		
		VecDestroy(&rhs); 
		VecDestroy(&x); 
	}
	
	// send calculated trace from each process to root
	if(myrank == 0){
		// receive values from all processes and add them to trace
		int processInd;
		for(processInd=1; processInd<noProcesses; processInd++){
			double result;
			MPI_Status status;
			MPI_Recv(&result, 1, MPI_DOUBLE, MPI_ANY_SOURCE, send_tag, MPI_COMM_WORLD, &status);
			trace += result;
		}
	}else{
		// send calculated trace to zero process
		MPI_Send(&trace, 1, MPI_DOUBLE, 0, send_tag, MPI_COMM_WORLD);
	}
	
	PetscPrintf(PETSC_COMM_WORLD, "Trace = %G\n", trace);
	PetscGetTime(&end);
	dif = end - start;
	PetscPrintf(PETSC_COMM_WORLD, "Trace calculations took = %G\n", dif);
	PetscPrintf(PETSC_COMM_WORLD, "Total time = %G\n", dif + loadTime);
	
	KSPDestroy(&ksp);
	MatDestroy(&A);
	MatDestroy(&B);
	
	PetscFinalize();
	return 0;
}