#include <petscksp.h>
#include <iostream>
#include "mpi.h"

#define RUNKSP

#ifdef RUNKSP

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

	/*-----------------------------------------------------------------------------
	 *  Declare vars and objects
	 *-----------------------------------------------------------------------------*/
	Vec rhs,sol,ex_sol;
	Mat MM;
	KSP ksp;
	PC pc;
	PetscReal norm;
	PetscInt i,j,nrows,ncols,its,tb=1000,tt;
	tt = tb;
	PetscErrorCode ierr;
	const PetscInt *rows;
	const PetscInt *cols;
	IS isrows,iscols;
	PetscBool flg = PETSC_TRUE;
	PetscScalar* v;
	PetscRandom rand;
	PetscScalar rval;
	int procid,size;

	/*-----------------------------------------------------------------------------
	 *  Init petsc
	 *-----------------------------------------------------------------------------*/
	PetscInitialize(&argc,&argv,NULL,NULL);
	MPI_Comm_rank(MPI_COMM_WORLD,&procid);
	MPI_Comm_size(MPI_COMM_WORLD,&size);

	ierr = PetscOptionsGetInt(NULL,"-tb",&tb, NULL); CHKERRQ(ierr);
	ierr = PetscOptionsGetInt(NULL,"-tt",&tt, NULL); CHKERRQ(ierr);

	ierr = PetscRandomCreate(PETSC_COMM_WORLD,&rand);CHKERRQ(ierr);
	ierr = PetscRandomSetFromOptions(rand);CHKERRQ(ierr);

	/*-----------------------------------------------------------------------------
	 *  Set up Matrix MM
	 *-----------------------------------------------------------------------------*/
	ierr = MatCreate(PETSC_COMM_WORLD,&MM);CHKERRQ(ierr);
	ierr = MatSetSizes(MM,PETSC_DECIDE,PETSC_DECIDE,tt,tb);CHKERRQ(ierr);
	ierr = MatSetType(MM,MATELEMENTAL);CHKERRQ(ierr);
	ierr = MatSetFromOptions(MM);CHKERRQ(ierr);
	ierr = MatSetUp(MM);CHKERRQ(ierr);


	/*-----------------------------------------------------------------------------
	 *  Set Local Values to Matrix MM
	 *-----------------------------------------------------------------------------*/
	ierr=MatGetOwnershipIS(MM,&isrows,&iscols);
	ierr=ISGetLocalSize(isrows,&nrows);
	ierr=ISGetIndices(isrows,&rows);
	ierr=ISGetLocalSize(iscols,&ncols);
	ierr=ISGetIndices(iscols,&cols);
	ierr=PetscMalloc1(nrows*ncols,&v);
	for (i=0; i<nrows; i++) {
		for (j=0; j<ncols; j++) {
			PetscRandomGetValue(rand,&rval);
			v[i*ncols+j] = rval; // PetscReal((i==j)?1:0);
		}
	}
	ierr =  MatSetValues(MM,nrows,rows,ncols,cols,v,INSERT_VALUES);CHKERRQ(ierr);
	ierr = MatAssemblyBegin(MM,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
	ierr = MatAssemblyEnd(MM,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
	ierr = ISRestoreIndices(isrows,&rows);CHKERRQ(ierr);
	ierr = ISRestoreIndices(iscols,&cols);CHKERRQ(ierr);
	ierr = ISDestroy(&isrows);CHKERRQ(ierr);
	ierr = ISDestroy(&iscols);CHKERRQ(ierr);
	if (flg) {
//		ierr=MatView(MM,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
	}

	/*-----------------------------------------------------------------------------
	 *  Set up sol and rhs vectors s.t rhs = MM*ex_sol
	 *-----------------------------------------------------------------------------*/
	ierr = VecCreate(PETSC_COMM_WORLD,&sol);CHKERRQ(ierr);
	ierr = VecSetSizes(sol,PETSC_DECIDE,tb);CHKERRQ(ierr);
	ierr = VecSetFromOptions(sol);CHKERRQ(ierr);
	ierr = VecDuplicate(sol,&rhs);CHKERRQ(ierr);
	ierr = VecDuplicate(sol,&ex_sol);CHKERRQ(ierr);
	ierr = VecSet(sol,0.0);CHKERRQ(ierr);
	ierr = VecSet(rhs,0.0);CHKERRQ(ierr);
	ierr = VecSet(ex_sol,1.0);CHKERRQ(ierr);
	ierr = MatMult(MM,ex_sol,rhs);CHKERRQ(ierr);
	if(flg){
		ierr = VecView(rhs,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
		ierr = VecView(sol,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
		ierr = VecView(ex_sol,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
	}

	 /*-----------------------------------------------------------------------------
	  *  Solve system MM*sol = rhs
	  *-----------------------------------------------------------------------------*/
	 ierr = KSPCreate(PETSC_COMM_WORLD,&ksp);CHKERRQ(ierr);
	 ierr = KSPSetType(ksp,KSPPREONLY);CHKERRQ(ierr);// optional
	 ierr = KSPGetPC(ksp,&pc);CHKERRQ(ierr);
	 ierr = PCSetType(pc,PCLU);CHKERRQ(ierr);
	 ierr = KSPSetOperators(ksp,MM,MM);CHKERRQ(ierr);
	 ierr = KSPSetTolerances(ksp,1.e-5,1.e-16,PETSC_DEFAULT,PETSC_DEFAULT);CHKERRQ(ierr);
	 ierr = PCFactorSetMatSolverPackage(pc,MATSOLVERELEMENTAL);
	 ierr = KSPSetFromOptions(ksp);CHKERRQ(ierr);
	 ierr = KSPSolve(ksp,rhs,sol);CHKERRQ(ierr);

	 
	 /*-----------------------------------------------------------------------------
	  *  Check Error
	  *-----------------------------------------------------------------------------*/
	if(flg){
		if(procid==0)
			std::cout<<"sol is "<<"\n";
		MPI_Barrier(MPI_COMM_WORLD);
		ierr = VecView(sol,PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);
	}
	ierr = VecAXPY(sol,-1.0,ex_sol);CHKERRQ(ierr);
	ierr = VecNorm(sol,NORM_2,&norm);CHKERRQ(ierr);
	ierr = KSPGetIterationNumber(ksp,&its);CHKERRQ(ierr);
	ierr = PetscPrintf(PETSC_COMM_WORLD,"Norm of error %g iterations %D\n",(double)norm,its);CHKERRQ(ierr);

	 /*-----------------------------------------------------------------------------
	  *  Destroy objects
	  *-----------------------------------------------------------------------------*/
	ierr = PetscFree(v);CHKERRQ(ierr);
	ierr = MatDestroy(&MM);CHKERRQ(ierr);
	ierr = VecDestroy(&rhs);CHKERRQ(ierr);
	ierr = VecDestroy(&sol);CHKERRQ(ierr);
	ierr = VecDestroy(&ex_sol);CHKERRQ(ierr);
	ierr = PetscRandomDestroy(&rand);CHKERRQ(ierr);
	ierr = KSPDestroy(&ksp);CHKERRQ(ierr);
	MPI_Barrier(MPI_COMM_WORLD);
	PetscFinalize();

	return 0;
}


#endif