#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 != 2)
  {
    printf("must run with 2 processes");
    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)
    {
      PetscInt d_nnz[] = {1, 1, 1};
      PetscInt o_nnz[] = {3, 3, 3};
    err =
      MatSetSizes(A, 3, 4, 6, 7);
    CHKERRQ(err);
      err = MatMPIAIJSetPreallocation(A, 0, d_nnz, 0, o_nnz);
      CHKERRQ(err);
    }
    else
    {
      PetscInt d_nnz[] = {3, 3, 3};
      PetscInt o_nnz[] = {1, 1, 1};
    err =
      MatSetSizes(A, 3, 3, 6, 7);
    CHKERRQ(err);
      err = MatMPIAIJSetPreallocation(A, 0, d_nnz, 0, o_nnz);
      CHKERRQ(err);
    }
  }
  err = MatSetUp(A);
  CHKERRQ(err);

  // set values inside the matrix
  for (PetscInt row = 0; row < 3; ++row)
  {
    for (PetscInt col = 3; col < 7; ++col)
    {
      err = MatSetValue(A, 3 * rank + row, col, 1, INSERT_VALUES);
      CHKERRQ(err);
    }
  }

  err = MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY);
  CHKERRQ(err);
  err = MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);
  CHKERRQ(err);

  err = MatView(A, PETSC_VIEWER_STDOUT_WORLD);
  CHKERRQ(err);

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

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