#include <iostream>
#include "petscmat.h"
#include "petscviewer.h"

// Compiling with: mpic++ -g3 -Wall -I ~/software/petsc/include -I ~/software/petsc/arch-linux2-c-debug/include -L ~/software/petsc/arch-linux2-c-debug/lib -lpetsc test.cpp
// Running without mpirun.

int main(int argc, char **args)
{
  PetscInitialize(&argc, &args, "", NULL);

  PetscErrorCode ierr = 0;
  int num_rows = 10;
  Mat matrix;
  Vec vector;

  // Create dense matrix, but code should work for sparse, too (I hope).
  ierr = MatCreateDense(PETSC_COMM_WORLD, PETSC_DECIDE, PETSC_DECIDE, num_rows, 4, NULL, &matrix); CHKERRQ(ierr);
  ierr = MatAssemblyBegin(matrix, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  ierr = MatAssemblyEnd(matrix, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);

  ierr = VecCreate(PETSC_COMM_WORLD, &vector); CHKERRQ(ierr);
  ierr = VecSetSizes(vector, PETSC_DECIDE, num_rows); CHKERRQ(ierr);
  ierr = VecSetFromOptions(vector); CHKERRQ(ierr);

  // Init vector with 0, 1, ... , num_rows-1
  PetscScalar *a;
  PetscInt range_start, range_end, pos = 0;
  VecGetOwnershipRange(vector, &range_start, &range_end);
  ierr = VecGetArray(vector, &a); CHKERRQ(ierr);
  for (PetscInt i = range_start; i < range_end; i++) {
    a[pos] = pos + range_start;
    pos++;
  }
  VecRestoreArray(vector, &a);
  // VecAssemblyBegin(vector); VecAssemblyEnd(vector);  I don't think it's needed here, changes nothing.
  ierr = VecView(vector, PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);

  PetscInt irow[num_rows];
  const PetscInt col = 2;
  const PetscScalar *vec;

  for (int i = 0; i < num_rows; i++) {
    irow[i] = i;
  }
  ierr = VecGetArrayRead(vector, &vec); CHKERRQ(ierr);
  // MatSetValuesLocal(Mat mat,PetscInt nrow,const PetscInt irow[],PetscInt ncol,const PetscInt icol[],const PetscScalar y[],InsertMode addv)
  ierr = MatSetValuesLocal(matrix, num_rows, irow, 1, &col, vec, INSERT_VALUES); CHKERRQ(ierr);
  // Works fine with MatSetValues
  ierr = VecRestoreArrayRead(vector, &vec); CHKERRQ(ierr);
  ierr = MatAssemblyBegin(matrix, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);
  ierr = MatAssemblyEnd(matrix, MAT_FINAL_ASSEMBLY); CHKERRQ(ierr);

  ierr = MatView(matrix, PETSC_VIEWER_STDOUT_WORLD); CHKERRQ(ierr);

  PetscFinalize();
  return 0;
}