program mattry

  use petscksp
  implicit none
#include <petsc/finclude/petsckspdef.h>

  PetscInt :: n=4   ! setting 4 cells per process

  PetscErrorCode         :: ierr
  PetscInt               :: size,rank,i
  Mat                    :: A,A02
  MatType                :: type
  IS                     :: isg(3)
  IS                     :: isl(3)
  ISLocalToGlobalMapping :: map

  integer, allocatable, dimension(:) :: idx

  call PetscInitialize(PETSC_NULL_CHARACTER,ierr); CHKERRQ(ierr)
  call MPI_Comm_size(PETSC_COMM_WORLD,size,ierr); CHKERRQ(ierr)
  call MPI_Comm_rank(PETSC_COMM_WORLD,rank,ierr);CHKERRQ(ierr)

  ! local index sets for 3 fields
  allocate(idx(n))
  idx=(/ (i-1, i=1,n) /)
  call ISCreateGeneral(PETSC_COMM_WORLD,n,idx,PETSC_COPY_VALUES,isl(1),ierr);CHKERRQ(ierr)
  call ISCreateGeneral(PETSC_COMM_WORLD,n,idx+n,PETSC_COPY_VALUES,isl(2),ierr);CHKERRQ(ierr)
  call ISCreateGeneral(PETSC_COMM_WORLD,n,idx+2*n,PETSC_COPY_VALUES,isl(3),ierr);CHKERRQ(ierr)
!  call ISView(isl3,PETSC_VIEWER_STDOUT_WORLD,ierr); CHKERRQ(ierr)
  deallocate(idx)

  ! global index sets for 3 fields
  allocate(idx(n))
  idx=(/ (i-1+rank*3*n, i=1,n) /)
  call ISCreateGeneral(PETSC_COMM_WORLD,n,idx,PETSC_COPY_VALUES,isg(1),ierr);CHKERRQ(ierr)
  call ISCreateGeneral(PETSC_COMM_WORLD,n,idx+n,PETSC_COPY_VALUES,isg(2),ierr); CHKERRQ(ierr)
  call ISCreateGeneral(PETSC_COMM_WORLD,n,idx+2*n,PETSC_COPY_VALUES,isg(3),ierr); CHKERRQ(ierr)
!  call ISView(isg3,PETSC_VIEWER_STDOUT_WORLD,ierr); CHKERRQ(ierr)
  deallocate(idx)

  ! local-to-global mapping
  allocate(idx(3*n))
  idx=(/ (i-1+rank*3*n, i=1,3*n) /)
  call ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD,1,3*n,idx,PETSC_COPY_VALUES,map,ierr); CHKERRQ(ierr)
!  call ISLocalToGlobalMappingView(map,PETSC_VIEWER_STDOUT_WORLD,ierr); CHKERRQ(ierr)
  deallocate(idx)

  ! create the 3-by-3 block matrix
  call MatCreate(PETSC_COMM_WORLD,A,ierr); CHKERRQ(ierr)
  call MatSetSizes(A,3*n,3*n,PETSC_DECIDE,PETSC_DECIDE,ierr); CHKERRQ(ierr)
  call MatSetUp(A,ierr); CHKERRQ(ierr)
  call MatSetOptionsPrefix(A,"A_",ierr); CHKERRQ(ierr)
  call MatSetLocalToGlobalMapping(A,map,map,ierr); CHKERRQ(ierr)
  call MatSetFromOptions(A,ierr); CHKERRQ(ierr)

  ! setup nest
  call MatNestSetSubMats(A,3,isg,3,isg,PETSC_NULL_OBJECT,ierr); CHKERRQ(ierr)

  ! set diagonal of block A02 to 0.65
  call MatGetLocalSubmatrix(A,isl(1),isl(3),A02,ierr); CHKERRQ(ierr)
  do i=1,n
     call MatSetValuesLocal(A02,1,i-1,1,i-1,0.65d0,INSERT_VALUES,ierr); CHKERRQ(ierr)
  end do
  call MatRestoreLocalSubMatrix(A,isl(1),isl(3),A02,ierr); CHKERRQ(ierr)
  call MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY,ierr); CHKERRQ(ierr)
  call MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY,ierr); CHKERRQ(ierr)

  ! verify
  call MatGetSubmatrix(A,isg(1),isg(3),MAT_INITIAL_MATRIX,A02,ierr); CHKERRQ(ierr)
  call MatView(A02,PETSC_VIEWER_STDOUT_WORLD,ierr);CHKERRQ(ierr)

  call PetscFinalize(ierr)

end program mattry