#include <petscmat.h>
#include <petscviewer.h>

int main(int argc,char **args)
{
  PetscErrorCode ierr;
  PetscMPIInt    size,rank;
  Mat            A,F,X,spRHST;
  PetscInt       m,n,nrhs,M,N,i,test;
  PetscScalar    v;
  PetscReal      norm,tol=PETSC_SQRT_MACHINE_EPSILON;
  PetscRandom    rand;
  PetscBool      displ=PETSC_FALSE;
 
  ierr = PetscInitialize(&argc, &args, NULL, NULL);if (ierr) return ierr;
  ierr = MPI_Comm_size(PETSC_COMM_WORLD,&size);CHKERRQ(ierr);
  ierr = MPI_Comm_rank(PETSC_COMM_WORLD,&rank);CHKERRQ(ierr);
  ierr = PetscOptionsGetBool(NULL,NULL,"-displ",&displ,NULL);CHKERRQ(ierr);

  /* Load matrix A */
  PetscViewer viewerA;
  ierr = PetscViewerBinaryOpen(PETSC_COMM_WORLD, "A1.petsc", FILE_MODE_READ, &viewerA);CHKERRQ(ierr);
  ierr = MatCreate(PETSC_COMM_WORLD, &A);CHKERRQ(ierr);
  ierr = MatLoad(A, viewerA);CHKERRQ(ierr);
  ierr = PetscViewerDestroy(&viewerA);CHKERRQ(ierr);
  ierr = MatGetLocalSize(A,&m,&n);CHKERRQ(ierr);
  ierr = MatGetSize(A,&M,&N);CHKERRQ(ierr);
  if (m != n) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_SIZ, "The matrix is not square (%d, %d)", m, n);

  /* Create dense matrix X */
  nrhs = N;
  ierr = PetscOptionsGetInt(NULL,NULL,"-nrhs",&nrhs,NULL);CHKERRQ(ierr);
  ierr = MatCreate(PETSC_COMM_WORLD, &X);CHKERRQ(ierr);
  ierr = MatSetSizes(X, m, PETSC_DECIDE, PETSC_DECIDE, nrhs);CHKERRQ(ierr);
  ierr = MatSetType(X, MATDENSE);CHKERRQ(ierr);
  ierr = MatSetFromOptions(X);CHKERRQ(ierr);
  ierr = MatSetUp(X);CHKERRQ(ierr);
  ierr = PetscRandomCreate(PETSC_COMM_WORLD,&rand);CHKERRQ(ierr);
  ierr = PetscRandomSetFromOptions(rand);CHKERRQ(ierr);
  ierr = MatSetRandom(X,rand);CHKERRQ(ierr);

  // factorise 'A' using LU Factorization 
  ierr = PetscPrintf(PETSC_COMM_WORLD,"using LU factorization\n");CHKERRQ(ierr);
  ierr = MatGetFactor(A,MATSOLVERMUMPS,MAT_FACTOR_LU,&F);CHKERRQ(ierr);
  ierr = MatLUFactorSymbolic(F,A,NULL,NULL,NULL);CHKERRQ(ierr);
  ierr = MatLUFactorNumeric(F,A,NULL);CHKERRQ(ierr);

  // Create spRHST: PETSc does not support compressed column format which is required by MUMPS for sparse RHS matrix,
  // thus user must create spRHST=spRHS^T 
  ierr = MatCreate(PETSC_COMM_WORLD,&spRHST);CHKERRQ(ierr);
  if (!rank) {
      /* MUMPS requires RHS be centralized on the host! */
      ierr = MatSetSizes(spRHST,nrhs,M,PETSC_DECIDE,PETSC_DECIDE);CHKERRQ(ierr);
  } else {
      ierr = MatSetSizes(spRHST,0,0,PETSC_DECIDE,PETSC_DECIDE);CHKERRQ(ierr);
  }
  ierr = MatSetType(spRHST,MATAIJ);CHKERRQ(ierr);
  ierr = MatSetFromOptions(spRHST);CHKERRQ(ierr);
  ierr = MatSetUp(spRHST);CHKERRQ(ierr);
  if (!rank) {
      v = 1.0;
      for (i=0; i<nrhs; i++) {
        ierr = MatSetValues(spRHST,1,&i,1,&i,&v,INSERT_VALUES);CHKERRQ(ierr);
      }
  }
  ierr = MatAssemblyBegin(spRHST,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(spRHST,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  
  /* Create spRHS = spRHST^T. Two matrices share internal matrix data structure */
  Mat spRHS = NULL;
  ierr = MatCreateTranspose(spRHST,&spRHS);CHKERRQ(ierr);
  
  /* get inv(A) using MatMumpsGetInverse() & sparse RHS: */
  ierr = MatMumpsGetInverse(F,spRHS);CHKERRQ(ierr);
  ierr = MatDestroy(&spRHS);CHKERRQ(ierr);

  /* Free data structures */
  ierr = MatDestroy(&A);CHKERRQ(ierr);
  ierr = MatDestroy(&spRHS);CHKERRQ(ierr);
  ierr = MatDestroy(&spRHST);CHKERRQ(ierr);
  ierr = MatDestroy(&X);CHKERRQ(ierr);
  ierr = PetscRandomDestroy(&rand);CHKERRQ(ierr);
  ierr = PetscFinalize();
  return ierr;
 
}