// Building MatIS with sparse or dense local matrix ?
//
// ~> g++ -o matISCSRDenseRect.exe matISCSRDenseRect.cpp -lpetsc -lm; mpirun -n 2 matISCSRDenseRect.exe

#include <iostream>
#include <string>
#include "petsc.h"

using namespace std;

int main(int argc,char **argv) {
  if (argc != 2) {cout << "error: need arg = csr or dense" << endl; return 1;}
  string matType = argv[1];
  if (matType != "csr" && matType != "dense") {cout << "error: need arg = csr or dense" << endl; return 1;}

  PetscInitialize(&argc, &argv, NULL, NULL);
  int size = 0; MPI_Comm_size(MPI_COMM_WORLD, &size); if (size != 2) {cout << "error: mpi != 2" << endl; return 1;}
  int rank = 0; MPI_Comm_rank(MPI_COMM_WORLD, &rank);

  PetscInt localIdx[2] = {0, 0};
  if (rank == 0) {localIdx[0] = 0; localIdx[1] = 1;}
  else           {localIdx[0] = 1; localIdx[1] = 2;}
  ISLocalToGlobalMapping rmap;
  ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, 1, &rank, PETSC_COPY_VALUES, &rmap);
  ISLocalToGlobalMapping cmap;
  ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, 2, localIdx, PETSC_COPY_VALUES, &cmap);

  Mat A;
  MatCreateIS(PETSC_COMM_WORLD, 1, PETSC_DECIDE, PETSC_DECIDE, 2, 3, rmap, cmap, &A);

  Mat Aloc;
  if (matType == "csr") {cout << matType << endl; MatCreateSeqAIJ(PETSC_COMM_SELF, 1, 2, 2, NULL, &Aloc);}
  else {cout << matType << endl; MatCreateSeqDense(PETSC_COMM_SELF, 1, 2, NULL, &Aloc);}
  localIdx[0] = 0; localIdx[1] = 1;
  PetscScalar localVal[2] = {1., 0.}; PetscInt oneLocalRow = 0;
  MatSetValues(Aloc, 1, &oneLocalRow, 2, localIdx, localVal, ADD_VALUES); // Add local row
  MatISSetLocalMat(A, Aloc);

  MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);

  Mat assembledGlobalMat;
  MatISGetMPIXAIJ(A, MAT_INITIAL_MATRIX, &assembledGlobalMat); // Collective call.

  MatView(A, PETSC_VIEWER_STDOUT_WORLD);

  PetscFinalize();
  return 0;
}