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

int main(int argc, char** argv)
{
  PetscErrorCode err;
  err = PetscInitialize(&argc, &argv, NULL, "help");
  CHKERRQ(err);

  int rank, size;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &size);
  if(size != 4)
  {
    printf("must run with 4 processes\n");
    MPI_Abort(MPI_COMM_WORLD, -1);
  }

  // create a sparse AIJ matrix distributed across MPI
  Mat A;
  err = MatCreate(MPI_COMM_WORLD, &A);
  CHKERRQ(err);
  err = MatSetType(A, MATMPIAIJ);
  CHKERRQ(err);
  // setup pre-allocation for matrix space
  {
    if(rank == 0)
    {
      err = MatSetSizes(A, 3, 4, 6, 7);
      CHKERRQ(err);
      PetscInt d_nnz[] = {0, 0, 0};
      PetscInt o_nnz[] = {1, 1, 1};
      err = MatMPIAIJSetPreallocation(A, 0, d_nnz, 0, o_nnz);
      CHKERRQ(err);
    }
    else if(rank == 1)
    {
      err = MatSetSizes(A, 3, 3, 6, 7);
      CHKERRQ(err);
      PetscInt d_nnz[] = {3, 3, 3};
      PetscInt o_nnz[] = {0, 0, 0};
      err = MatMPIAIJSetPreallocation(A, 0, d_nnz, 0, o_nnz);
      CHKERRQ(err);
    }
    else if(rank == 2)
    {
      err = MatSetSizes(A, 3, 4, 6, 7);
      CHKERRQ(err);
      PetscInt d_nnz[] = {0, 0, 0};
      PetscInt o_nnz[] = {0, 0, 0};
      err = MatMPIAIJSetPreallocation(A, 0, d_nnz, 0, o_nnz);
      CHKERRQ(err);
    }
    else
    {
      err = MatSetSizes(A, 3, 3, 6, 7);
      CHKERRQ(err);
      PetscInt d_nnz[] = {3, 3, 3};
      PetscInt o_nnz[] = {0, 0, 0};
      err = MatMPIAIJSetPreallocation(A, 0, d_nnz, 0, o_nnz);
      CHKERRQ(err);
    }
  }
  err = MatSetUp(A);
  CHKERRQ(err);
  
  PetscInt row_start, row_stop;
  err = MatGetOwnershipRange(A, &row_start, &row_stop);
  CHKERRQ(err);

  PetscInt col_start, col_stop;
  err = MatGetOwnershipRangeColumn(A, &col_start, &col_stop);
  CHKERRQ(err);

  printf("rank %d owns rows/cols: [%d,%d), [%d, %d)\n", rank, row_start,
    row_stop, col_start, col_stop);

  // free memory
  err = MatDestroy(&A);
  CHKERRQ(err);

  // cleanup any internal PETSc data at end of program
  err = PetscFinalize();
  CHKERRQ(err);
  return 0;
}