#include <petsc.h>

int main(int argc, char *argv[])
{
   Mat A;
   Vec u, b;
   unsigned int n = 100;
   
   KSP ksp;
   PC pc;
   
   PetscInitialize(&argc, &argv, (char*)0, "void");
   MatCreate(PETSC_COMM_WORLD, &A);
   MatSetSizes(A, PETSC_DECIDE, PETSC_DECIDE, n, n);
   MatSetType(A, MATAIJ);
   MatSetUp(A);

   for (unsigned int i = 0; i < n-1; ++i)
   {
      MatSetValue(A, i, i, 2, INSERT_VALUES);
      MatSetValue(A, i+1, i, -1, INSERT_VALUES);
      MatSetValue(A, i, i+1, -1, INSERT_VALUES);
   }
   MatSetValue(A, n-1, n-1, 2, INSERT_VALUES);
   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
   
   VecCreate(PETSC_COMM_WORLD, &u);
   VecSetSizes(u, PETSC_DECIDE, n);
   VecSetFromOptions(u);
   VecDuplicate(u, &b);
   
   VecSet(u, 1);
   MatMult(A, u, b);
   
   // Create Index Sets
   PetscInt *P_A_indices, *P_B_indices, *T_indices, *P_indices;
   P_A_indices = new PetscInt[n/4];
   P_B_indices = new PetscInt[n/4];
   T_indices = new PetscInt[n/2];
   P_indices = new PetscInt[n/2];
   
   for (unsigned int i = 0; i < n/2; ++i)
   {
      P_indices[i] = i;
      T_indices[i] = i+n/2;
   }
   for (unsigned int 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;
   ISCreateGeneral(PETSC_COMM_WORLD, n/2, P_indices, PETSC_COPY_VALUES, &P_IS);
   ISCreateGeneral(PETSC_COMM_WORLD, n/2, T_indices, PETSC_COPY_VALUES, &T_IS);
   ISCreateGeneral(PETSC_COMM_WORLD, n/4, P_A_indices, PETSC_COPY_VALUES, &P_A_IS);
   ISCreateGeneral(PETSC_COMM_WORLD, n/4, P_B_indices, PETSC_COPY_VALUES, &P_B_IS);

   // Solve the system
   KSPCreate(PETSC_COMM_WORLD, &ksp);
   KSPSetType(ksp, KSPGMRES);
   KSPSetOperators(ksp, A, A);
   KSPSetFromOptions(ksp);
   KSPSetUp(ksp);
   KSPGetPC(ksp, &pc);

   // Define the fieldsplit for the global matrix
   PCFieldSplitSetIS(pc, "P", P_IS);
   PCFieldSplitSetIS(pc, "T", T_IS);

   // fieldsplit for submatrix P:
   KSP *ksp_all, ksp_P;
   PetscInt dummy;
   PCFieldSplitGetSubKSP(pc, &dummy, &ksp_all);

   ksp_P = ksp_all[0];
   PC pc_P;
   KSPGetPC(ksp_P, &pc_P); // This should be the preconditioner for fieldsplit P
   PCFieldSplitSetIS(pc_P, "A", P_A_IS);
   PCFieldSplitSetIS(pc_P, "B", P_B_IS);

   KSPSolve(ksp, b, u);
   PetscFinalize();
}