// Using PETSc MatIS, how to get local matrix (= one domain) before and after assembly (neumann/dirichlet).
//
// A 3x3 global matrix is built (diag: 1, 2, 1): it's made of 2 overlapping 2x2 local matrix (diag: 1, 1).
// Get non assembled local matrix : OK with MatISGetLocalMat.
//
// How to get assembled local matrix (initial local matrix + neigbhor contributions on the borders) ?
// Need to use MatGetLocalSubMatrix ? Not working ?! Expected result is Diag: 2, 1.
//
// ~> g++ -o matISLocalMat.exe matISLocalMat.cpp -lpetsc -lm; mpirun -n 2 matISLocalMat.exe

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

int main(int argc,char **argv) {
  PetscInitialize(&argc, &argv, NULL, NULL);
  int size = 0; MPI_Comm_size(MPI_COMM_WORLD, &size); if (size != 2) return 1;
  int rank = 0; MPI_Comm_rank(MPI_COMM_WORLD, &rank);

  PetscInt localSize = 2, globalSize = localSize*2 /*2 MPI*/;
  PetscInt localIdx[2] = {0, 0};
  if (rank == 0) {localIdx[0] = 0; localIdx[1] = 1;}
  else           {localIdx[0] = 1; localIdx[1] = 2;}
  ISLocalToGlobalMapping map;
  ISLocalToGlobalMappingCreate(PETSC_COMM_WORLD, 1, localSize, localIdx, PETSC_COPY_VALUES, &map);
  Mat A;
  MatCreateIS(PETSC_COMM_WORLD, 1, localSize, localSize, globalSize, globalSize, map, map, &A);
  ISLocalToGlobalMappingDestroy(&map);
  MatISSetPreallocation(A, localSize, NULL, localSize, NULL);
  PetscScalar localVal[4] = {1., 0., 0., 1.};
  MatSetValues(A, 2, localIdx, 2, localIdx, localVal, ADD_VALUES);
  MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY); MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY);
  MatView(A, PETSC_VIEWER_STDOUT_WORLD); PetscViewerFlush(PETSC_VIEWER_STDOUT_WORLD); // Diag: 1, 2, 1

  if (rank > 0) { // Do not pollute stdout: print only 1 proc
    std::cout << std::endl << "non assembled local matrix:" << std::endl << std::endl;
    Mat nonAssembledLocalMat;
    MatISGetLocalMat(A, &nonAssembledLocalMat);
    MatView(nonAssembledLocalMat, PETSC_VIEWER_STDOUT_SELF); // Diag: 1, 1

    std::cout << std::endl << "assembled local matrix:" << std::endl << std::endl;
    Mat assembledLocalMat;
    IS is; ISCreateGeneral(PETSC_COMM_SELF, localSize, localIdx, PETSC_COPY_VALUES, &is);
    MatGetLocalSubMatrix(A, is, is, &assembledLocalMat); // KO ?!...
    MatView(assembledLocalMat, PETSC_VIEWER_STDOUT_SELF); // Would like to get => Diag: 2, 1
  }

  PetscFinalize();
  return 0;
}