#include <petsc.h>

#define N 100

int main(int argc, char *argv[])
{
   Mat A;
   Vec u, b;
   PetscErrorCode ierr;
   KSP ksp;
   PC pc;
   unsigned int i;
   
   PetscInitialize(&argc, &argv, (char*)0, "void");
   ierr = MatCreate(PETSC_COMM_WORLD, &A);CHKERRQ(ierr);
   ierr = MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, N, N);CHKERRQ(ierr);
   ierr = MatSetType(A, MATAIJ);CHKERRQ(ierr);
   ierr = MatSetUp(A);CHKERRQ(ierr);
   
   for (i = 0; i < N-1; ++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 = MatSetValue(A, N-1, N-1, 2, INSERT_VALUES);CHKERRQ(ierr);
   ierr = MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
   ierr = MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
   
   ierr = VecCreate(PETSC_COMM_WORLD, &u);CHKERRQ(ierr);
   ierr = VecSetSizes(u, PETSC_DECIDE, N);CHKERRQ(ierr);
   ierr = VecSetFromOptions(u);CHKERRQ(ierr);
   ierr = VecDuplicate(u, &b);CHKERRQ(ierr);
   
   ierr = VecSet(u, 1);CHKERRQ(ierr);
   ierr = MatMult(A, u, b);CHKERRQ(ierr);
   
   // Create Index Sets
   PetscInt P_A_indices[N/4];
   PetscInt P_B_indices[N/4];
   PetscInt T_indices[N/2];
   PetscInt P_indices[N/2];
   
   for (i = 0; i < N/2; ++i)
   {
      P_indices[i] = i;
      T_indices[i] = i+N/2;
   }
   for (i = 0; i < N/4; ++i)
   {
      P_A_indices[i] = i;
      P_B_indices[i] = i+N/4;
   }

   IS P_A_IS, P_B_IS, P_IS, T_IS;
   ierr = ISCreateGeneral(PETSC_COMM_WORLD, N/2, P_indices, PETSC_COPY_VALUES, &P_IS);CHKERRQ(ierr);
   ierr = ISCreateGeneral(PETSC_COMM_WORLD, N/2, T_indices, PETSC_COPY_VALUES, &T_IS);CHKERRQ(ierr);
   ierr = ISCreateGeneral(PETSC_COMM_WORLD, N/4, P_A_indices, PETSC_COPY_VALUES, &P_A_IS);CHKERRQ(ierr);
   ierr = ISCreateGeneral(PETSC_COMM_WORLD, N/4, P_B_indices, PETSC_COPY_VALUES, &P_B_IS);CHKERRQ(ierr);
   
   // Solve the system
   ierr = KSPCreate(PETSC_COMM_WORLD, &ksp);CHKERRQ(ierr);
   ierr = KSPSetType(ksp, KSPGMRES);CHKERRQ(ierr);
   ierr = KSPSetOperators(ksp, A, A);CHKERRQ(ierr);
   ierr = KSPSetFromOptions(ksp);CHKERRQ(ierr);
   ierr = KSPGetPC(ksp, &pc);CHKERRQ(ierr);
   
   // Define the fieldsplit for the global matrix
   ierr = PCFieldSplitSetIS(pc, "P", P_IS);CHKERRQ(ierr);
   ierr = PCFieldSplitSetIS(pc, "T", T_IS);CHKERRQ(ierr);
   ierr = KSPSetUp(ksp);CHKERRQ(ierr);
   
   // fieldsplit for submatrix P:
   KSP *ksp_all, ksp_P;
   PetscInt dummy;
   ierr = PCFieldSplitGetSubKSP(pc, &dummy, &ksp_all);CHKERRQ(ierr);
   
   ksp_P = ksp_all[0];
   PC pc_P;
   ierr = KSPGetPC(ksp_P, &pc_P);CHKERRQ(ierr);
   ierr = PCFieldSplitSetIS(pc_P, "A", P_A_IS);CHKERRQ(ierr);
   ierr = PCFieldSplitSetIS(pc_P, "B", P_B_IS);CHKERRQ(ierr);
   
   ierr = KSPSolve(ksp, b, u);CHKERRQ(ierr);
   PetscFinalize();
   return 0;
}