#include <petscksp.h>

static char help[] = "";

#undef __FUNCT__
#define __FUNCT__ "loadBinMatrix"
void loadBinMatrix(char filein[], Mat *A, const MatType mattype){
	PetscErrorCode ierr;
	PetscViewer    view;

	// Open binary file.  Note that we use FILE_MODE_READ to indicate reading from this file.
	PetscViewerBinaryOpen(PETSC_COMM_WORLD,filein,FILE_MODE_READ,&view);
	//Load the matrix; then destroy the viewer.    
	MatCreate(PETSC_COMM_WORLD,A);
	ierr = MatSetType(*A, mattype);CHKERRQ(ierr);		
	ierr = MatSetFromOptions(*A);CHKERRQ(ierr);
	//ierr = MatSetUp(*A);CHKERRQ(ierr);
	MatLoad(*A,view);  
	PetscViewerDestroy(&view);
}

// prams: lin solver context, right hand side, column index of large matrix B (index of rhs vector)
void solveLinSystem(KSP *ksp, Vec *rhs, PetscInt i, PetscInt rowsA, PetscInt colsA, Vec *diagonal){
	PetscScalar val;
	PetscErrorCode ierr;
	PetscInt start,end;
	ierr = KSPSolve(*ksp,*rhs,*rhs);CHKERRQ(ierr);
	ierr = VecGetOwnershipRange(*rhs,&start,&end);CHKERRQ(ierr);
	if((i>=start) && (i<end)){
		VecGetValues(*rhs, 1, &i, &val);	
		VecSetValues(*diagonal, 1, &i, &val, INSERT_VALUES);
	}	
}

#undef __FUNCT__
#define __FUNCT__ "main"
int main(int argc,char **args){
	char fileLocA[PETSC_MAX_PATH_LEN], fileLocb[PETSC_MAX_PATH_LEN], fileLocx[PETSC_MAX_PATH_LEN];
	Mat A, B; /* linear system matrix, matrix containing right hand sides */
	KSP ksp;     /* linear solver context */
	PetscInt rowsA, colsA, rowsB, colsB, i, noLinSystems;
	PetscErrorCode ierr;
	PetscLogDouble start, end, elapsed_time;

	PetscInitialize(&argc,&args,(char *)0,help);
  
	ierr = PetscOptionsGetString(PETSC_NULL,"-finA",fileLocA,PETSC_MAX_PATH_LEN,PETSC_NULL);CHKERRQ(ierr);
	ierr = PetscOptionsGetString(PETSC_NULL,"-finB",fileLocb,PETSC_MAX_PATH_LEN,PETSC_NULL);CHKERRQ(ierr);
	// read sparse matrix A from binary file
	loadBinMatrix(fileLocA, &A, MATMPIAIJ); //MATSEQAIJ
	// read matrix B from binary file
	loadBinMatrix(fileLocb, &B, MATDENSE);
	
	ierr = PetscGetTime(&start); CHKERRQ(ierr); // returns the current time of day in seconds
	
	MatGetSize(B, 0, &noLinSystems);
	MatTranspose(B, MAT_REUSE_MATRIX, &B);
	MatGetLocalSize(A, &rowsA, &colsA);
	MatGetLocalSize(B, &rowsB, &colsB);

	ierr = KSPCreate(PETSC_COMM_WORLD,&ksp);CHKERRQ(ierr); //Create linear solver context
	ierr = KSPSetOperators(ksp,A,A,SAME_PRECONDITIONER);CHKERRQ(ierr);
	ierr = KSPSetFromOptions(ksp);CHKERRQ(ierr);
	PetscReal trace = 0;
	Vec diagonal;
	VecCreateMPI(PETSC_COMM_WORLD, PETSC_DECIDE,  noLinSystems, &diagonal);
	
	PetscInt startRow, endRow, globalRow;
	ierr = MatGetOwnershipRange(B, &startRow, &endRow);CHKERRQ(ierr);	
	
	for (i=0; i<rowsB; i++) {
		globalRow = startRow + i;
		PetscScalar *vals, *cols;
		ierr = MatGetRow(B, globalRow, NULL, &cols, &vals); CHKERRQ(ierr);
		Vec b; // Vector in which we will put the column from B	
		VecCreateMPI(PETSC_COMM_WORLD, PETSC_DECIDE, noLinSystems, &b);
		VecSetValues(b, noLinSystems, cols, vals, INSERT_VALUES);
		MatRestoreRow(B, globalRow, NULL, &cols, &vals);				
		VecAssemblyBegin(b);
		VecAssemblyEnd(b); 								
		solveLinSystem(&ksp, &b, i, rowsA, colsA, &diagonal); 
		ierr = VecDestroy(&b);CHKERRQ(ierr); 
	}
	
	VecAssemblyBegin(diagonal);
	VecAssemblyEnd(diagonal);
	VecSum(diagonal, &trace);
	PetscPrintf(PETSC_COMM_WORLD,"Trace = %3G\n",trace);
	
	// Clean up
	ierr = KSPDestroy(&ksp);CHKERRQ(ierr);
	ierr = MatDestroy(&A);CHKERRQ(ierr);
	ierr = MatDestroy(&B);CHKERRQ(ierr);
	ierr = VecDestroy(&diagonal);CHKERRQ(ierr);
	
	ierr = PetscGetTime(&end); CHKERRQ(ierr);
	elapsed_time = end - start;
	
	PetscPrintf(PETSC_COMM_WORLD,"Execution time: %5G seconds\n", elapsed_time);
	
	ierr = PetscFinalize();
	return 0;
}