module petsc_stuff
use f90_kind

! Include PETSc files

#include "finclude/petscvecdef.h"
#include "finclude/petscmatdef.h"
#include "finclude/petsckspdef.h"
use petscvec
use petscmat
use petscksp
implicit none

contains


subroutine init_petsc
use f90_kind
implicit none

PetscErrorCode :: ierr

call PetscInitialize(PETSC_NULL_CHARACTER,ierr)

end subroutine init_petsc



subroutine stop_petsc
use f90_kind
implicit none

PetscErrorCode :: ierr

call PetscFinalize(ierr)

end subroutine stop_petsc








subroutine init_petsc_vector(petsc_vector,no_rows)
use f90_kind
implicit none

Vec :: petsc_vector
integer, intent(in) :: no_rows

PetscErrorCode :: ierr

call VecCreate(PETSC_COMM_WORLD,petsc_vector,ierr)
call VecSetSizes(petsc_vector,PETSC_DECIDE,no_rows,ierr)
call VecSetFromOptions(petsc_vector,ierr)
call VecZeroEntries(petsc_vector,ierr)

end subroutine init_petsc_vector






subroutine finalize_petsc_vector(petsc_vector)
! Turns the vector into its final form, i.e. no coefficients can be
! added after calling this subroutine.
use f90_kind
implicit none

Vec :: petsc_vector

PetscErrorCode :: ierr

call VecAssemblyBegin(petsc_vector,ierr)
call VecAssemblyEnd(petsc_vector,ierr)

end subroutine finalize_petsc_vector









subroutine init_petsc_matrix(petsc_matrix,no_rows,no_cols,no_entries_row)
use f90_kind
implicit none

Mat :: petsc_matrix
integer, intent(in) :: no_rows
integer, intent(in) :: no_cols
integer, intent(in) :: no_entries_row

integer, dimension(:), allocatable :: entries_per_rows
PetscErrorCode :: ierr

allocate(entries_per_rows(1:no_rows))
entries_per_rows = no_entries_row

call MatCreateSeqAIJ(PETSC_COMM_SELF,no_rows,no_cols, /*no_entries_row*/ 0,entries_per_rows,petsc_matrix,ierr)
call MatZeroEntries(petsc_matrix,ierr)

deallocate(entries_per_rows)

end subroutine init_petsc_matrix




subroutine finalize_petsc_matrix(petsc_matrix)
! Turns the matrix into its final form, i.e. no coefficients can be
! added after calling this subroutine.
use f90_kind
implicit none

Mat :: petsc_matrix

PetscErrorCode :: ierr

call MatAssemblyBegin(petsc_matrix,MAT_FINAL_ASSEMBLY,ierr)
call MatAssemblyEnd(petsc_matrix,MAT_FINAL_ASSEMBLY,ierr)

end subroutine finalize_petsc_matrix






subroutine finalize_petsc_solver(A_petsc,ksp_A)
! Turns the matrix into its final form, i.e. no coefficients can be
! added after calling this subroutine.
use f90_kind
implicit none

Mat :: A_petsc
KSP :: ksp_A

PetscErrorCode :: ierr
PC :: pcc

call KSPCreate(PETSC_COMM_WORLD,ksp_A,ierr)
call KSPSetOperators(ksp_A,A_petsc,A_petsc,DIFFERENT_NONZERO_PATTERN,ierr)  
call KSPSetType(ksp_A,KSPCG,ierr)
call KSPGetPC(ksp_A,pcc,ierr)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

! AMG preconditioner

!call PCSetType(pcc,PCGAMG,ierr)

! ICC preconditioner

call PCSetType(pcc,PCICC,ierr)
call PCFactorSetShiftType(pcc,MAT_SHIFT_POSITIVE_DEFINITE,ierr)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

call PCSetFromOptions(pcc,ierr)
call KSPSetFromOptions(ksp_A,ierr)
call KSPSetUp(ksp_A,ierr)

end subroutine finalize_petsc_solver










subroutine addcoef_petsc_vector(petsc_vector,row,coef)
use f90_kind
implicit none

Vec :: petsc_vector
integer, intent(in) :: row
real(dp), intent(in) :: coef

PetscErrorCode :: ierr

call VecSetValue(petsc_vector,row-1,coef,ADD_VALUES,ierr)

end subroutine addcoef_petsc_vector





subroutine addcoef_petsc_matrix(petsc_matrix,row,col,coef)
use f90_kind
implicit none

Mat :: petsc_matrix
integer, intent(in) :: row, col
real(dp), intent(in) :: coef
PetscInt :: m,n

PetscErrorCode :: ierr

call MatSetValue(petsc_matrix,row-1,col-1,coef,ADD_VALUES,ierr)

end subroutine addcoef_petsc_matrix



subroutine triple_mat_prod(A,B,C,s)
! calc C = s * A * B * A^T
use f90_kind
implicit none

Mat :: A
Mat :: B
Mat :: C
real(dp), intent(in) :: s

Mat :: D
PetscErrorCode :: ierr
PetscInt :: m,n

! D = B * A^T

call MatMatTransposeMult(B,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT_DOUBLE_PRECISION,D,ierr)

! C = A * D

call MatMatMult(A,D,Mat_initial_matrix,PETSC_DEFAULT_DOUBLE_PRECISION,C,ierr)

! C = s * C

call MatScale(C,s,ierr)

end subroutine triple_mat_prod






subroutine solve_petsc_system(ksp_A,rhs_petsc,x_petsc,toler,maxiter,iter_done)
! Solves the system of equations with a given right hand side
! Initial guess can be provided in x
!
! TODO: Implement the use of the initial guess
use f90_kind
implicit none

Vec :: rhs_petsc
Vec :: x_petsc
KSP :: ksp_A
real(dp), intent(in) :: toler
integer, intent(in) :: maxiter
integer, intent(out) :: iter_done

real(dp) :: resnorm
real(dp) :: rtol
real(dp) :: dtol

PetscErrorCode :: ierr
KSPConvergedReason :: reason

rtol = toler
dtol = 1.d100

call KSPSetTolerances(ksp_A,rtol,PETSC_DEFAULT_DOUBLE_PRECISION,dtol,maxiter,ierr)
call KSPSolve(ksp_A,rhs_petsc,x_petsc,ierr)
call KSPGetIterationNumber(ksp_A,iter_done,ierr)
call KSPGetConvergedReason(ksp_A,reason,ierr)
if (reason < 1) then
  WRITE (*,*) 'PANIC: solve_petsc_system: No convergence. reason ',reason
  iter_done = -1 * iter_done
endif
call KSPGetResidualNorm(ksp_A,resnorm,ierr)

end subroutine solve_petsc_system

end module petsc_stuff