#include <petsc/finclude/petscmat.h>
program main
  use petscmat

  implicit none

  Mat A
  MatPartitioning part
  IS is
  PetscInt   ::     i, m, N
  PetscInt   ::     rstart, rend
  PetscInt, pointer, dimension(:) ::   emptyranks, bigranks, cols
  PetscScalar, pointer, dimension(:) :: vals
  PetscInt :: &
    nbigranks = 10, &
    nemptyranks = 10
  PetscMPIInt   ::  rank, sizef
  PetscErrorCode ierr
  PetscBool set
  PetscInt, parameter :: zero = 0, one = 1, two = 2, three = 3

  PetscCallA(PetscInitialize(ierr))

  PetscCallMPIA(MPI_Comm_rank(PETSC_COMM_WORLD, rank, ierr))
  PetscCallMPIA(MPI_Comm_size(PETSC_COMM_WORLD, sizef, ierr))

  allocate (emptyranks(nemptyranks))
  allocate (bigranks(nbigranks))

  PetscCallA(PetscOptionsGetIntArray(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-emptyranks', emptyranks, nemptyranks, set, ierr))
  PetscCallA(PetscOptionsGetIntArray(PETSC_NULL_OPTIONS, PETSC_NULL_CHARACTER, '-bigranks', bigranks, nbigranks, set, ierr))

  m = 1
  do i = 1, nemptyranks
    if (rank == emptyranks(i)) m = 0
  end do
  do i = 1, nbigranks
    if (rank == bigranks(i)) m = 5
  end do

  deallocate (emptyranks)
  deallocate (bigranks)

  PetscCallA(MatCreate(PETSC_COMM_WORLD, A, ierr))
  PetscCallA(MatSetsizes(A, m, m, PETSC_DECIDE, PETSC_DECIDE, ierr))
  PetscCallA(MatSetFromOptions(A, ierr))
  PetscCallA(MatSeqAIJSetPreallocation(A, three, PETSC_NULL_INTEGER_ARRAY, ierr))
  PetscCallA(MatMPIAIJSetPreallocation(A, three, PETSC_NULL_INTEGER_ARRAY, two, PETSC_NULL_INTEGER_ARRAY, ierr))
  PetscCallA(MatSeqBAIJSetPreallocation(A, one, three, PETSC_NULL_INTEGER_ARRAY, ierr))
  PetscCallA(MatMPIBAIJSetPreallocation(A, one, three, PETSC_NULL_INTEGER_ARRAY, two, PETSC_NULL_INTEGER_ARRAY, ierr))
  PetscCallA(MatSeqSBAIJSetPreallocation(A, one, two, PETSC_NULL_INTEGER_ARRAY, ierr))
  PetscCallA(MatMPISBAIJSetPreallocation(A, one, two, PETSC_NULL_INTEGER_ARRAY, one, PETSC_NULL_INTEGER_ARRAY, ierr))

  PetscCallA(MatGetSize(A, PETSC_NULL_INTEGER, N, ierr))
  PetscCallA(MatGetOwnershipRange(A, rstart, rend, ierr))

  allocate (cols(1:3))
  allocate (vals(1:3))
  do i = rstart, rend - 1

    cols = (/mod((i + N - 1), N), i, mod((i + 1), N)/)
    vals = [1.0, 1.0, 1.0]
    PetscCallA(MatSetValues(A, one, [i], three, cols, vals, INSERT_VALUES, ierr))
  end do
  deallocate (cols)
  deallocate (vals)
  PetscCallA(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY, ierr))
  PetscCallA(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY, ierr))
  PetscCallA(MatView(A, PETSC_VIEWER_STDOUT_WORLD, ierr))

  PetscCallA(MatPartitioningCreate(PETSC_COMM_WORLD, part, ierr))
  PetscCallA(MatPartitioningSetAdjacency(part, A, ierr))
  PetscCallA(MatPartitioningSetFromOptions(part, ierr))
  PetscCallA(MatPartitioningApply(part, is, ierr))
  PetscCallA(ISView(is, PETSC_VIEWER_STDOUT_WORLD, ierr))
  PetscCallA(ISDestroy(is, ierr))
  PetscCallA(MatPartitioningDestroy(part, ierr))
  PetscCallA(MatDestroy(A, ierr))
  PetscCallA(PetscFinalize(ierr))

end program

!/*TEST
!
!   test:
!      nsize: 8
!      args: -emptyranks 0,2,4 -bigranks 1,3,7 -mat_partitioning_type average
!      output_file: output/ex17_1.out
!      # cannot test with external package partitioners since they produce different results on different systems
!
!TEST*/