#include <petsc.h>
#include <mpi.h>

static char help[] = "superlu_dist seg fault test.\n\n";

int main(int argc,char **args)
{
  Mat  A, B, X, F;
  KSP  ksp;
  PC  pc;
  PetscInt  inode, nnodes;
  PetscViewer  v_file;
  PetscScalar  pp;
  PetscErrorCode  ierr;

  ierr = PetscInitialize(&argc, &args, NULL, help);  CHKERRQ(ierr);

  MPI_Comm_size(PETSC_COMM_WORLD, &nnodes);
  MPI_Comm_rank(PETSC_COMM_WORLD, &inode);

  MatCreate(PETSC_COMM_WORLD, &A);
  MatSetType(A, MATAIJ);

  PetscViewerBinaryOpen(PETSC_COMM_WORLD,"mat_a.dat",FILE_MODE_READ,&v_file);
  ierr = MatLoad(A, v_file);
  printf("MatLoad %i\n", ierr);  CHKERRQ(ierr);
  PetscViewerDestroy(&v_file);

  PetscInt    M,N,m,n,rstart,rend,cols[2],i,nrhs=4000;
  PetscScalar vals[2];
  ierr = MatGetSize(A,&M,&N);CHKERRQ(ierr);
  ierr = MatGetLocalSize(A,&m,&n);CHKERRQ(ierr);
  printf("A: %d, %d, %d %d\n",m,n,M,N);

  ierr = MatCreateDense(PETSC_COMM_WORLD,m,PETSC_DECIDE,PETSC_DECIDE,nrhs,NULL,&B);CHKERRQ(ierr);
  ierr = MatCreateDense(PETSC_COMM_WORLD,m,PETSC_DECIDE,PETSC_DECIDE,nrhs,NULL,&X);CHKERRQ(ierr);
  //MatZeroEntries(B);
  //MatZeroEntries(X);

  ierr = MatGetOwnershipRange(A,&rstart,&rend);CHKERRQ(ierr);
  for (i=rstart; i<rend; i++) {
    cols[0] = 0; cols[1] = 1; vals[0] = 1.0; vals[1] = 1.0;
    ierr = MatSetValues(B,1,&i,2,cols,vals,INSERT_VALUES);CHKERRQ(ierr);
  }
  ierr = MatAssemblyBegin(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(B,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyBegin(X,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(X,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

  //pp = 1.0;
  //MatShift(B, pp);
  //MatView(B,PETSC_VIEWER_STDOUT_WORLD);

  KSPCreate(PETSC_COMM_WORLD, &ksp);
  KSPSetOperators(ksp, A, A);

  KSPSetType(ksp, KSPPREONLY);
  KSPGetPC(ksp, &pc);

  PCSetType(pc, PCLU);
  ierr = PCFactorSetMatSolverType(pc,MATSOLVERSUPERLU_DIST);CHKERRQ(ierr);

  ierr = KSPSetFromOptions(ksp);CHKERRQ(ierr);
  ierr = KSPSetUp(ksp);CHKERRQ(ierr);

  ierr = PCFactorGetMatrix(pc, &F);  CHKERRQ(ierr);
  printf("PCFactorGetMatrix %i\n", ierr);

  ierr = MatMatSolve(F, B, X);CHKERRQ(ierr);
  printf("MatMatSolve %i\n", ierr);

  ierr = MatDestroy(&A);CHKERRQ(ierr);
  ierr = MatDestroy(&B);CHKERRQ(ierr);
  ierr = MatDestroy(&X);CHKERRQ(ierr);
  ierr = KSPDestroy(&ksp);CHKERRQ(ierr);
  ierr = PetscFinalize();
  return ierr;
}