#include <petscmat.h>

int main(int argc, char **argv)
{
  Mat            L;
  PetscScalar    vals[1];
  PetscInt       cols[1];
  PetscScalar    one, zero;
  PetscInt       n = 10, start, end, row;
  PetscErrorCode ierr;

  ierr = PetscInitialize(&argc, &argv, NULL, NULL); if (ierr) return ierr;
  #ifdef PETSC_USE_COMPLEX
    one = 1.0 + PETSC_i;
  #else
    one = 1.0;
  #endif
  zero = 0.0;

  ierr = MatCreate(PETSC_COMM_WORLD, &L);CHKERRQ(ierr);
  ierr = MatSetType(L, MATAIJ);CHKERRQ(ierr);
  ierr = MatSetSizes(L, PETSC_DECIDE, PETSC_DECIDE, n, n);CHKERRQ(ierr);

  ierr = MatSeqAIJSetPreallocation(L, 1, NULL);CHKERRQ(ierr);
  ierr = MatMPIAIJSetPreallocation(L, 1, NULL, 0, NULL);CHKERRQ(ierr); /* No allocated non-zero in off-diagonal part */
  ierr = MatSetOption(L, MAT_IGNORE_ZERO_ENTRIES, PETSC_TRUE);CHKERRQ(ierr);
  ierr = MatSetOption(L, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE);CHKERRQ(ierr);
  ierr = MatSetOption(L, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE);CHKERRQ(ierr);

  ierr = MatGetOwnershipRange(L, &start, &end);CHKERRQ(ierr);
  /* assembling a diagonal matrix */
  for (row = start; row < end; ++row) {
    cols[0] = row;
    vals[0] = one;
    ierr = MatSetValues(L, 1, &row, 1, cols, vals, ADD_VALUES);CHKERRQ(ierr);
  }
  ierr = MatAssemblyBegin(L, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(L, MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

  ierr = PetscPrintf(PETSC_COMM_WORLD, "After first assembly:\n");CHKERRQ(ierr);
  ierr = MatView(L, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);

  ierr = MatSetOption(L, MAT_NEW_NONZERO_LOCATION_ERR, PETSC_TRUE);CHKERRQ(ierr);

  //ierr = MatZeroEntries(L);CHKERRQ(ierr);

  ierr = PetscPrintf(PETSC_COMM_WORLD, "After second assembly:\n");CHKERRQ(ierr);
  ierr = MatView(L, PETSC_VIEWER_STDOUT_WORLD);CHKERRQ(ierr);

  /* assembling a diagonal matrix, adding a zero value to non-diagonal part */
  for (row = start; row < end; ++row) {
    if (!row) {
      cols[0] = n-1;
      vals[0] = zero;
#ifdef PETSC_USE_COMPLEX
      ierr = PetscPrintf(PETSC_COMM_SELF, "row %D col: %D val: %g + %g i\n", row, cols[0], (double) PetscRealPart(vals[0]), (double) PetscRealPart(vals[0]));CHKERRQ(ierr);
#else
      ierr = PetscPrintf(PETSC_COMM_SELF, "row %D col: %D val: %g\n", row, cols[0], (double) vals[0]);CHKERRQ(ierr);
#endif
      ierr = MatSetValues(L, 1, &row, 1, cols, vals, ADD_VALUES);CHKERRQ(ierr);
    }
    cols[0] = row;
    vals[0] = one;
    ierr = MatSetValues(L, 1, &row, 1, cols, vals, ADD_VALUES);CHKERRQ(ierr);
  }
  ierr = MatAssemblyBegin(L,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
  ierr = MatAssemblyEnd(L,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);

  ierr = MatDestroy(&L);CHKERRQ(ierr);
  ierr = PetscFinalize();
  return ierr;
}