#include <petsc.h>

/*#define N 100*/
#define N 120

int main(int argc, char *argv[])
{
   Mat A;
   Vec u,b;
   PetscErrorCode ierr;
   KSP ksp;
   PC pc;
   unsigned int i;
   PetscInt Istart,Iend,local_m,local_n;
   int rank;

   PetscInitialize(&argc,&argv,(char*)0,"void");
   ierr = MatCreate(PETSC_COMM_WORLD,&A);CHKERRQ(ierr);
   ierr = MatSetSizes(A,PETSC_DECIDE,N,N,N);CHKERRQ(ierr);
   ierr = MatSetType(A,MATMPIAIJ);CHKERRQ(ierr);
   ierr = MatSetUp(A);CHKERRQ(ierr);

   MPI_Comm_rank(MPI_COMM_WORLD,&rank);
   ierr = MatGetOwnershipRange(A,&Istart,&Iend);CHKERRQ(ierr);
   printf("%d: Istart = %d, Iend = %d\n",rank,Istart,Iend);

   for (i = Istart; i < Iend; ++i)
   {
       ierr = MatSetValue(A,i,i,2,INSERT_VALUES);CHKERRQ(ierr);
       ierr = MatSetValue(A,i+1,i,-1,INSERT_VALUES);CHKERRQ(ierr);
       ierr = MatSetValue(A,i,i+1,-1,INSERT_VALUES);CHKERRQ(ierr);
   }

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

   ierr = MatGetLocalSize(A,&local_m,&local_n);CHKERRQ(ierr);
   printf("%d: local_m = %d, local_n = %d\n",rank,local_m,local_n);

   // Create Index Sets
   PetscInt* A_indices;
   PetscInt* B_indices;
   PetscInt A_size,B_size;

   if(rank == 0)
   {
       A_size = Iend-Istart;
       B_size = 0;
       A_indices = (PetscInt*) calloc(A_size,sizeof(PetscInt));
       B_indices = (PetscInt*) calloc(B_size,sizeof(PetscInt));
       for(i = 0; i < A_size; ++i)
          A_indices[i] = Istart + i;
   }
   else if(rank == 1)
   {
       A_size = (Iend-Istart)/2;
       B_size = (Iend-Istart)/2;
       A_indices = (PetscInt*) calloc(A_size,sizeof(PetscInt));
       B_indices = (PetscInt*) calloc(B_size,sizeof(PetscInt));
       for(i = 0; i < A_size; ++i)
          A_indices[i] = Istart + i;
       for(i = 0; i < B_size; ++i)
          B_indices[i] = Istart + i + A_size;
   }

   IS A_IS, B_IS;
   ierr = ISCreateGeneral(PETSC_COMM_WORLD,A_size,A_indices,PETSC_COPY_VALUES,&A_IS);CHKERRQ(ierr);
   ierr = ISCreateGeneral(PETSC_COMM_WORLD,B_size,B_indices,PETSC_COPY_VALUES,&B_IS);CHKERRQ(ierr);
   printf("%d: A_size = %d, B_size = %d\n",rank,A_size,B_size);
   printf("%d: ISCreateGeneral completed\n",rank);

   // Solve the system
   ierr = KSPCreate(PETSC_COMM_WORLD,&ksp);CHKERRQ(ierr);
   ierr = KSPSetType(ksp,KSPGMRES);CHKERRQ(ierr);
   ierr = KSPSetOperators(ksp,A,A);CHKERRQ(ierr);
   printf("%d: KSPSetOperators completed\n",rank);

   // Define the fieldsplit for the global matrix
   ierr = KSPGetPC(ksp,&pc);CHKERRQ(ierr);
   ierr = PCSetType(pc,PCFIELDSPLIT);CHKERRQ(ierr);
   ierr = PCFieldSplitSetIS(pc,"a",A_IS);CHKERRQ(ierr);
   ierr = PCFieldSplitSetIS(pc,"b",B_IS);CHKERRQ(ierr);
   printf("%d: PCFieldSplitSetIS completed\n",rank);
   ierr = ISSetBlockSize(A_IS,2);CHKERRQ(ierr);
   ierr = ISSetBlockSize(B_IS,2);CHKERRQ(ierr);
   ierr = KSPSetFromOptions(ksp);CHKERRQ(ierr);
   ierr = KSPSetUp(ksp);CHKERRQ(ierr);

   free(A_indices);
   free(B_indices);
   ierr = ISDestroy(&A_IS);CHKERRQ(ierr);
   ierr = ISDestroy(&B_IS);CHKERRQ(ierr);
   ierr = KSPDestroy(&ksp);CHKERRQ(ierr);

   PetscFinalize();
   return 0;
}