program test_fieldsplit

#include "petsc/finclude/petscksp.h"
	use petscksp

    implicit none 

    Mat A

    Vec x, rhs

    MatType myType

    VecType myVec

    KSP ksp

    KSPType  mySolver

    PC  mypc
!------------------------------------------------------------------------------------------------------------------
        
    integer :: i, lm, M, ierr

    integer :: conn(9)

    real :: b(9*3), blockValues(3,3), relTol, absTol

    real :: ufields(2), pfields(1)
!------------------------------------------------------------------------------------------------------------------

    call PetscInitialize(PETSC_NULL_CHARACTER,ierr)

    lm = 3 * 9
    M = 3 * 9

    conn = (/0,1,2,3,4,5,6,7,8/)

    mytype = MATMPIBAIJ
    myVec = VECMPI

    call MatCreate(PETSC_COMM_WORLD, A, ierr)

    call MatSetSizes(A, lm, lm, M, M, ierr)

    call MatSetType(A, myType, ierr)

    call MatSetBlockSize(A, 3, ierr)

    call MatMPIBAIJSetPreallocation(A, 3, 9, PETSC_NULL_INTEGER, 9, PETSC_NULL_INTEGER, ierr)

    call MatSetOption(A, MAT_ROW_ORIENTED, PETSC_FALSE, ierr)

    call MatSetUp(A, ierr)

    call MatZeroEntries(A, ierr)

    blockValues = 0.0
    blockValues(1,1) = 1
    blockValues(2,2) = 1
    blockValues(3,3) = 1

    do i=1,9
        call MatSetValuesBlocked(A, 1, conn(i), 1, conn(i), blockValues, INSERT_VALUES, ierr)
    enddo

    call MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY, ierr)
    call MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY, ierr)
    
    !call MatView(this%A, PETSC_VIEWER_STDOUT_WORLD,ierr)

    ! Create RHS and solution vector
    call MatCreateVecs(A, rhs, x, ierr);

    b = 1
    call VecSetValuesBlocked(rhs, 9, conn, b, INSERT_VALUES, ierr);

    call VecAssemblyBegin(rhs, ierr)
    call VecAssemblyEnd(rhs, ierr)

    !  KSP

    relTol = 0.01
    absTol = 0.000001
    mySolver = "bcgs"

    call KSPCreate(PETSC_COMM_WORLD, ksp, ierr)
    call KSPSetOptionsPrefix(ksp, "test" ,ierr)
    call KSPSetType(ksp, mySolver, ierr)
    call KSPSetTolerances(ksp, relTol, absTol, PETSC_DEFAULT_REAL, 10000, ierr)
    call KSPSetErrorIfNotConverged(ksp, PETSC_TRUE, ierr)

    ! test this, just to see I can exit first from the linear system loop
    call KSPSetNormType(ksp, KSP_NORM_UNPRECONDITIONED, ierr)

    ! Set this preconditioner
    call KSPGetPC(ksp, mypc, ierr)

    ! FIELD SPLIT
    call PCSetType(mypc, PCFIELDSPLIT, ierr)  

    call PCFieldSplitSetBlockSize(mypc, 3, ierr)

    ufields(1) = 0
    ufields(2) = 1
    pfields(1) = 2

    call PCFieldSplitSetFields(mypc, "u", 2, ufields, ufields, ierr)
    call PCFieldSplitSetFields(mypc, "p", 1, pfields, pfields, ierr)

    call KSPSetOperators(ksp, A, A, ierr)

    call KSPSolve(ksp, rhs, x, ierr)

    call PetscFinalize(ierr)

end program test_fieldsplit