#include <petsc.h>

int main(int argc, char** argv) {
    PetscInitialize(&argc, &argv, NULL, NULL);
    PetscMPIInt size;
    MPI_Comm_size(PETSC_COMM_WORLD, &size);
    if(size != 1)
        return 1;
    Mat J1, J2;
    MatCreate(PETSC_COMM_WORLD, &J1);
    MatSetType(J1, MATMPIAIJ);
    MatSetSizes(J1, 2, 2, 2, 2);
    PetscInt ij[5] = { 0, 1, 2, 0, 1 };
    PetscScalar C[2] = { 1, 2 };
    MatMPIAIJSetPreallocationCSR(J1, ij, ij + 3, C);
    // J1 = [[1, 0] ; [0, 2]]
    MatAssemblyBegin(J1, MAT_FINAL_ASSEMBLY);
    MatAssemblyEnd(J1, MAT_FINAL_ASSEMBLY);
    MatDuplicate(J1, MAT_COPY_VALUES, &J2);
    MatScale(J2, 4.0);
    // J2 = [[4, 0] ; [0, 8]]
    Mat N;
    {
        Mat nest[4] = { J1, NULL, NULL, J2 };
        MatCreateNest(PETSC_COMM_WORLD, 2, NULL, 2, NULL, nest, &N);
    }
    // N = [[1, 0, 0, 0] ; [0, 2, 0, 0] ; [0, 0, 4, 0] ; [0, 0, 0, 8]]
    KSP nest;
    KSPCreate(PETSC_COMM_WORLD, &nest);
    KSPSetFromOptions(nest);
    KSPSetOperators(nest, N, N);
    {
        PC pc;
        KSPGetPC(nest, &pc);
        PCSetType(pc, PCFIELDSPLIT);
        IS is[2];
        PetscInt idx[2] = { 0, 1 };
        ISCreateGeneral(PETSC_COMM_WORLD, 2, idx, PETSC_COPY_VALUES, is);
        idx[0] = 2;
        idx[1] = 3;
        ISCreateGeneral(PETSC_COMM_WORLD, 2, idx, PETSC_COPY_VALUES, is + 1);
        PCFieldSplitSetIS(pc, NULL, is[0]);
        PCFieldSplitSetIS(pc, NULL, is[1]);
        ISDestroy(is);
        ISDestroy(is + 1);
    }
    KSPSetUp(nest);
    KSP sub;
    {
        PC pc;
        KSPGetPC(nest, &pc);
        PetscInt nsplits;
        KSP* subksp;
        PCFieldSplitGetSubKSP(pc, &nsplits, &subksp);
        sub = subksp[nsplits - 1];
        PetscObjectReference((PetscObject)subksp[nsplits - 1]);
        KSPSetOperators(sub, J2, J2);
        PetscFree(subksp);
    }
    Mat S;
    {
        PC pc;
        KSPGetPC(sub, &pc);
        PCSetType(pc, PCFIELDSPLIT);
        IS is[2];
        PetscInt idx[1] = { 0 };
        ISCreateGeneral(PETSC_COMM_WORLD, 1, idx, PETSC_COPY_VALUES, is);
        idx[0] = 1;
        ISCreateGeneral(PETSC_COMM_WORLD, 1, idx, PETSC_COPY_VALUES, is + 1);
        PCFieldSplitSetIS(pc, NULL, is[0]);
        PCFieldSplitSetIS(pc, NULL, is[1]);
        ISDestroy(is);
        ISDestroy(is + 1);
        KSPSetUp(sub);
        {
            PetscInt nsplits;
            KSP* subksp;
            PCFieldSplitGetSubKSP(pc, &nsplits, &subksp);
            PetscInt ij[3] = { 0, 1, 0 };
            PetscScalar C[2] = { 0.00001 };
            MatCreate(PETSC_COMM_WORLD, &S);
            MatSetType(S, MATMPIAIJ);
            MatSetSizes(S, 1, 1, 1, 1);
            MatMPIAIJSetPreallocationCSR(S, ij, ij + 2, C);
            MatAssemblyBegin(S, MAT_FINAL_ASSEMBLY);
            MatAssemblyEnd(S, MAT_FINAL_ASSEMBLY);
            KSPSetOperators(subksp[nsplits - 1], S, S);
            // J2 = [[4, 0] ; [0, 0.00001]]
            PetscFree(subksp);
        }
        KSPSetUp(sub);
    }
    {
        Vec x, y;
        MatCreateVecs(J2, &x, &y);
        VecSet(y, 1.0);
        KSPSetFromOptions(sub);
        KSPSolve(sub, y, x);
        // x = [[4, 0] ; [0, 0.00001]]^-1 * [1 ; 1] = [1/4 ; 1/0.00001]
        VecView(x, PETSC_VIEWER_STDOUT_WORLD);
        VecDestroy(&x);
        VecDestroy(&y);
    }
    {
        Vec x, y;
        MatCreateVecs(N, &x, &y);
        VecSet(y, 1.0);
        KSPSolve(nest, y, x);
        // why is
        // x = [[1, 0, 0, 0] ; [0, 2, 0, 0] ; [0, 0, 4, 0] ; [0, 0, 0, 8]]^-1    *    [1 ; 1 ; 1 ; 1] = [1 ; 1/2 ; 1/4 ; 1/8]
        // and not
        // x = [[1, 0, 0, 0] ; [0, 2, 0, 0] ; [0, 0, 4, 0] ; [0, 0, 0, 0.00001]]^-1 * [1 ; 1 ; 1 ; 1] = [1 ; 1/2 ; 1/4 ; 1/0.00001]
        VecView(x, PETSC_VIEWER_STDOUT_WORLD);
        VecDestroy(&x);
        VecDestroy(&y);
    }
    KSPDestroy(&sub);
    KSPDestroy(&nest);
    MatDestroy(&S);
    MatDestroy(&N);
    MatDestroy(&J2);
    MatDestroy(&J1);
    PetscFinalize();
}