PROGRAM test_ksp

#include <petsc/finclude/petscdmdef.h> 
#include <petsc/finclude/petscmatdef.h>
#include <petsc/finclude/petscvecdef.h>
#include <petsc/finclude/petsckspdef.h>

  USE petscdmda
  USE petscsys

  IMPLICIT NONE

  INTEGER, PARAMETER :: pr = 8

  DM                :: decomp
  DMBoundaryType    :: BType(3)
  INTEGER           :: N
  INTEGER           :: is, js ,ks, nxl, nyl, nzl  ! subdomain index
  INTEGER           :: rank, nproc
  INTEGER           :: i, j, k
  REAL(pr)          :: h                          ! space step
  Vec               :: x,b
  REAL(pr), POINTER :: ptr_vec(:,:,:) => NULL()
  Mat               :: A
  MatNullSpace      :: nullspace
  MatStencil        :: row(4), col(4,7)
  PetscInt          :: i1 = 1, i7 = 7
  REAL(pr)          :: value(7)
  KSP               :: ksp
  PetscErrorCode    :: ierr
  REAL(pr)          :: q2 = 0.5_pr
  REAL(pr)          :: erri = 0, err1 = 0
  LOGICAL           :: is_set
  

  CALL PetscInitialize( PETSC_NULL_OBJECT, ierr )
  CALL MPI_Comm_rank( PETSC_COMM_WORLD, rank, ierr )
  CALL MPI_Comm_size( PETSC_COMM_WORLD, nproc, ierr )

  N    = 64
  h    = 1.0_pr / N
  BType = DM_BOUNDARY_PERIODIC

  CALL PetscOptionsGetInt( PETSC_NULL_OBJECT, PETSC_NULL_CHARACTER, '-N', N, is_set, ierr )

  CALL DMDACreate3d( PETSC_COMM_WORLD, BType(1), BType(2), BType(3), &
                   & DMDA_STENCIL_STAR, -N, -N, -N, PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, 1, 1, &
                   & PETSC_NULL_INTEGER,  PETSC_NULL_INTEGER, PETSC_NULL_INTEGER, &
                   & decomp, ierr ) 
  CALL DMDAGetCorners( decomp, is, js, ks, nxl, nyl, nzl, ierr) 

  ! create vector for rhs and solution
  CALL DMCreateGlobalVector( decomp, b, ierr )
  CALL VecDuplicate( b, x, ierr )
  CALL VecSet( x, 0.0_pr, ierr )

  ! set rhs
  CALL DMDAVecGetArrayF90( decomp, b, ptr_vec, ierr )
  DO i = is, is+nxl-1
    DO j = js, js+nyl-1
      DO k = ks, ks+nzl-1
        ptr_vec(i,j,k) = -3 * SIN((i+q2)*h) * SIN((j+q2)*h) * SIN((k+q2)*h)
      END DO
    END DO
  END DO
  ptr_vec = - h**2 * ptr_vec
  CALL DMDAVecRestoreArrayF90( decomp, b, ptr_vec, ierr )

  ! assembly rhs
  CALL VecAssemblyBegin( b, ierr )
  CALL VecAssemblyEnd( b, ierr )

  ! create matrix
  CALL DMSetMatType( decomp, MATAIJ, ierr )
  CALL DMCreateMatrix( decomp, A, ierr )

  ! set matrix
  DO i = is, is+nxl-1
    DO j = js, js+nyl-1
      DO k = ks, ks+nzl-1

        ! row index of current point
        row(MatStencil_i) = i
        row(MatStencil_j) = j
        row(MatStencil_k) = k   

        ! column index of current point and its neighbors
        col(MatStencil_i,1) = i;    col(MatStencil_j,1) = j;    col(MatStencil_k,1) = k
        col(MatStencil_i,2) = i-1;  col(MatStencil_j,2) = j;    col(MatStencil_k,2) = k     
        col(MatStencil_i,3) = i+1;  col(MatStencil_j,3) = j;    col(MatStencil_k,3) = k 
        col(MatStencil_i,4) = i;    col(MatStencil_j,4) = j-1;  col(MatStencil_k,4) = k 
        col(MatStencil_i,5) = i;    col(MatStencil_j,5) = j+1;  col(MatStencil_k,5) = k 
        col(MatStencil_i,6) = i;    col(MatStencil_j,6) = j;    col(MatStencil_k,6) = k-1 
        col(MatStencil_i,7) = i;    col(MatStencil_j,7) = j;    col(MatStencil_k,7) = k+1 

        ! set values at current point and its neighbors
        value(1 ) =  6.0_pr
        value(2:) = -1.0_pr

        ! set the matrix's elements
         CALL MatSetValuesStencil( A, i1, row, i7, col, value, INSERT_VALUES, ierr ) 

      END DO
    END DO
  END DO

  ! assembly the matrix
  CALL MatAssemblyBegin( A, MAT_FINAL_ASSEMBLY, ierr )
  CALL MatAssemblyEnd( A, MAT_FINAL_ASSEMBLY, ierr )

  ! create nullspace, periodic bc give the matrix a nullspace
  CALL MatNullSpaceCreate( PETSC_COMM_WORLD, PETSC_TRUE, PETSC_NULL_INTEGER, &
  & PETSC_NULL_INTEGER, nullspace, ierr )

  CALL MatSetNullSpace( A, nullspace, ierr )
  CALL MatNullSpaceDestroy( nullspace, ierr )

  ! remove the nullspace
  CALL MatGetNullSpace( A, nullspace, ierr )
  CALL MatNullSpaceRemove( nullspace, b, PETSC_NULL_OBJECT, ierr )

  ! create ksp solver
  CALL KSPCreate( PETSC_COMM_WORLD, ksp, ierr )
  CALL KSPSetDM( ksp, decomp, ierr )
  CALL KSPSetDMActive( ksp, PETSC_FALSE, ierr )
  CALL KSPSetFromOptions( ksp, ierr )

  ! Solve system
  CALL KSPSetOperators( ksp, A, A, ierr )
  CALL KSPSolve( ksp, b, x, ierr )

  ! get solution and check the error
  CALL DMDAVecGetArrayF90( decomp, x, ptr_vec, ierr )
  DO i = is, is+nxl-1
    DO j = js, js+nyl-1
      DO k = ks, ks+nzl-1
        erri = MAX(erri, ABS(ptr_vec(i,j,k) - SIN((i+q2)*h)*SIN((j+q2)*h)*SIN((k+q2)*h)))
        err1 = err1 + ABS(ptr_vec(i,j,k) - SIN((i+q2)*h)*SIN((j+q2)*h)*SIN((k+q2)*h))
      END DO
    END DO
  END DO
  CALL DMDAVecRestoreArrayF90( decomp, x, ptr_vec, ierr )

  !CALL MPI_Allreduce( MPI_IN_PLACE, erri, 1, MPI_REAL8, MPI_MAX, MPI_COMM_WORLD, ierr )
  !CALL MPI_Allreduce( MPI_IN_PLACE, err1, 1, MPI_REAL8, MPI_MIN, MPI_COMM_WORLD, ierr )
  CALL MPI_Allreduce( erri, MPI_IN_PLACE, 1, MPI_REAL8, MPI_MAX, PETSC_COMM_WORLD, ierr )
  CALL MPI_Allreduce( err1, MPI_IN_PLACE, 1, MPI_REAL8, MPI_SUM, PETSC_COMM_WORLD, ierr )

  IF( rank == 0 ) THEN 
    PRINT*, 'norm1 error:   ', err1 / N**3
    PRINT*, 'norm inf error:', erri
  END IF

  CALL VecDestroy( x, ierr )
  CALL VecDestroy( b, ierr ) 
  CALL MatDestroy( A, ierr )
  CALL KSPDestroy( ksp, ierr )
  CALL DMDestroy( decomp, ierr )

  CALL PetscFinalize( ierr )

END PROGRAM test_ksp