! This file is part of GTC version 4.4 
! GTC version 4.4 is released under the 3-Clause BSD license:

! Copyright (c) 2002,2010,2016, GTC Team (team leader: Zhihong Lin, zhihongl@uci.edu)
! All rights reserved.

! Redistribution and use in source and binary forms, with or without 
! modification, are permitted provided that the following conditions are met:

! 1. Redistributions of source code must retain the above copyright notice, 
!    this list of conditions and the following disclaimer.

! 2. Redistributions in binary form must reproduce the above copyright notice, 
!    this list of conditions and the following disclaimer in the documentation 
!    and/or other materials provided with the distribution.

! 3. Neither the name of the GTC Team nor the names of its contributors may be 
!    used to endorse or promote products derived from this software without 
!    specific prior written permission.
! ==============================================================================
! X. Wei: This implementation combines the new and old PETSc interfaces
! The new intefaces are based on Jian Bao's implementation(11e0bf3)
! fix some PETSc versioning problems
#include "petscversion.h"
#if PETSC_VERSION_LT(3,8,0)
#include "petsc_old.F90"
#else

subroutine petsc_init
#include "petsc/finclude/petsc.h"
  use petsc
  use system_env, only: ierror, GTC_ABORT
  use global_parameters,only:mype,partd_comm,gtcout,psi0
  use petsc_array,only: newcomm,nproc_newcomm,myrank_newcomm

  implicit none 

  PetscErrorCode ierr

  call PetscInitialize(PETSC_NULL_CHARACTER,ierr)
  call MPI_Comm_dup(partd_comm,newcomm,ierror)

  call mpi_comm_size(newcomm,nproc_newcomm,ierror)

  call mpi_comm_rank(newcomm,myrank_newcomm,ierror)

  if(ierror /= 0) then
    write(*,*)mype,'main_petsc: newcomm error',ierror
    call GTC_ABORT 
  endif
        
  return
end subroutine petsc_init

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

subroutine petsc_destroy(psdata)
#include "petsc/finclude/petsc.h"
  use petsc
  use petsc_array
  implicit none

  type(petscData) :: psdata
  PetscErrorCode ierr

  call VecDestroy(psdata%x,ierr)
  call VecDestroy(psdata%b,ierr)
  call VecDestroy(psdata%x_recv,ierr)
  call VecScatterDestroy(psdata%gather,ierr)
  call ISDestroy(psdata%is_from,ierr)
  call ISDestroy(psdata%is_to,ierr)
  call MatDestroy(psdata%Apetsc,ierr)
  call KSPDestroy(psdata%ksp,ierr)

  return
end subroutine petsc_destroy
subroutine petsc_final
#include "petsc/finclude/petsc.h"
  use petsc
  implicit none
  PetscErrorCode ierr

  call PetscFinalize(ierr)
  return
end subroutine petsc_final

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

subroutine laplacian_core(psdata,luserx_output,luserb_input)
#include "petsc/finclude/petsc.h"
  use petsc
  use precision, only: lk
  use global_parameters,only:istep,irk,mype
  use petsc_array
  implicit none

  type(petscData) :: psdata
  real(lk),dimension(0:),intent(inout) :: luserx_output
  real(lk),dimension(0:),intent(in) :: luserb_input

  PetscScalar temp
  PetscErrorCode ierr
  PetscInt i,ilow,ihigh,isize,ilo,ihi
!  PetscInt isi
!  PetscOffset idx
!  PetscScalar vec_x(1)
  PetscScalar, pointer :: vec_x(:)

  call VecGetLocalSize(psdata%b,isize,ierr)
  call VecGetOwnershipRange(psdata%b,ilow,ihigh,ierr)
  do i=ilow,ihigh-1
    temp=luserb_input(i)
    call VecSetValue(psdata%b,i,temp,INSERT_VALUES,ierr)
  enddo
  call VecAssemblyBegin(psdata%b,ierr)
  call VecAssemblyEnd(psdata%b,ierr)

!     if(istep==1 .and. irk==1 .and. mype==0)then
!       write(*,*)'checking vector b:'
!       call VecView(psdata%b,PETSC_VIEWER_STDOUT_SELF,ierr)
!       CHKERRQ(ierr)
!     endif
!     if(istep==1 .and. irk==1 .and. mype==0)then
!       write(*,*)'checking matrix a:'
!       call MatView(psdata%Apetsc,PETSC_VIEWER_STDOUT_SELF,ierr)
!       CHKERRQ(ierr)
!     endif

  call MatMult(psdata%Apetsc,psdata%b,psdata%x,ierr)

!* 
  call VecScatterBegin(psdata%gather,psdata%x,psdata%x_recv,INSERT_VALUES, &
    SCATTER_FORWARD,ierr)
  call VecScatterEnd(psdata%gather,psdata%x,psdata%x_recv,INSERT_VALUES, &
    SCATTER_FORWARD,ierr)
!*
  call VecGetOwnershipRange(psdata%x_recv,ilo,ihi,ierr)
!  call VecGetLocalSize(psdata%x_recv,isi,ierr)
!  call VecGetArray(psdata%x_recv,vec_x,idx,ierr)
  call VecGetArrayF90(psdata%x_recv,vec_x,ierr)
  do i=ilo,ihi-1
    luserx_output(i)=vec_x(i-ilo+1)
  enddo
!  call VecRestoreArray(psdata%x,vec_x,idx,ierr)
  call VecRestoreArrayF90(psdata%x_recv,vec_x,ierr)

  return
end subroutine laplacian_core

!********************************YX following

subroutine lapmat_pscreate(mindexlap,matrixsize,psdata,lusera_input,luserp_input,lusers_input)
#include "petsc/finclude/petsc.h"
  use petsc
  use precision, only:lk
  use global_parameters,only: mgridlow,mgridhigh,mype,istep,gtcout
  use petsc_array
  implicit none

  integer mindexlap,matrixsize
  type(petscData) :: psdata
  real(lk),dimension(:),intent(in) :: lusera_input
  integer,dimension(:),intent(in)  :: luserp_input
  integer,dimension(-1:),intent(in)  ::lusers_input

  double precision tol,value(mindexlap)
  PetscErrorCode ierr
  PetscInt n,i,j,ij,ilow,ihigh,nz
  PetscInt idx(1)
  PetscBool flg
  integer col(mindexlap)

  n=matrixsize
! PETSc 3.7 adds one additional argument to this function. uncomment the next
! line for PETSc >3.7.0  (modified by L. Shi, 01/21/2017)
  call PetscOptionsGetInt(PETSC_NULL_OPTIONS,PETSC_NULL_CHARACTER,'-n',n,flg,ierr)
! For PETSc versions before 3.7:
!  call PetscOptionsGetInt(PETSC_NULL_CHARACTER,'-n',n,flg,ierr)
  call VecCreate(newcomm,psdata%x,ierr)
  call VecSetSizes(psdata%x,PETSC_DECIDE,n,ierr)
  call VecSetFromOptions(psdata%x,ierr)
  call VecDuplicate(psdata%x,psdata%b,ierr)

!*
  call VecCreate(PETSC_COMM_SELF,psdata%x_recv,ierr)
  call VecSetSizes(psdata%x_recv,n,n,ierr)
  call VecSetType(psdata%x_recv,VECSEQ,ierr)
  call VecSetFromOptions(psdata%x_recv,ierr)

  call ISCreateStride(newcomm,n,0,1,psdata%is_from,ierr)
  call ISCreateStride(PETSC_COMM_SELF,n,0,1,psdata%is_to,ierr)

  call VecScatterCreate(psdata%x,psdata%is_from,psdata%x_recv,psdata%is_to,psdata%gather,ierr)
!*

!     call MatCreate(newcomm,psdata%Apetsc,ierr)
!     CHKERRQ(ierr)
!     call MatSetSizes(psdata%Apetsc,PETSC_DECIDE,PETSC_DECIDE,n,n,ierr)
!     CHKERRQ(ierr)
!     call MatSetType(psdata%Apetsc,MATMPIAIJ,ierr)
!     CHKERRQ(ierr)
!*optimized matrix creation
  call MatCreateAIJ(newcomm,PETSC_DECIDE,PETSC_DECIDE,n,n, &
     mindexlap,PETSC_NULL_INTEGER,mindexlap,PETSC_NULL_INTEGER, &
     psdata%Apetsc,ierr)

  call MatSetFromOptions(psdata%Apetsc,ierr)
  call MatGetOwnershipRange(psdata%Apetsc,ilow,ihigh,ierr)

  ij=lusers_input(ilow-1)
  do i=ilow,ihigh-1
    nz=lusers_input(i)-lusers_input(i-1)
    do j=1,nz
      ij=ij+1
      col(j) = luserp_input(ij)
      value(j) = lusera_input(ij)
    enddo
    idx(1)=i
    call MatSetValues(psdata%Apetsc,1,idx,nz,col,value,INSERT_VALUES,ierr)
  enddo

  call MatAssemblyBegin(psdata%Apetsc,MAT_FINAL_ASSEMBLY,ierr)
  call MatAssemblyEnd(psdata%Apetsc,MAT_FINAL_ASSEMBLY,ierr)
!     call MatTranspose(psdata%Apetsc,Bpetsc1,ierr)
!     call MatAXPY(Bpetsc1,1.d0,psdata%Apetsc,DIFFERENT_NONZERO_PATTERN,ierr)
!     call MatPartitioningSetAdjacency(mpart, Bpetsc1, ierr )
!     call MatCopy(psdata%Apetsc,Bpetsc1,DIFFERENT_NONZERO_PATTERN,ierr)

  call KSPCreate(newcomm,psdata%ksp,ierr)
  call KSPSetOperators(psdata%ksp,psdata%Apetsc,psdata%Apetsc,ierr)
!     call KSPGetPC(psdata%ksp,psdata%pc,ierr)
!     call PCSetType(psdata%pc,PCJACOBI,ierr)
! 1.d-4 seems to be adequate
!  tol=1.d-4
  tol=1.d-6
!     tol=1.d-8
  call KSPSetTolerances(psdata%ksp,tol,PETSC_DEFAULT_REAL, &
       PETSC_DEFAULT_REAL,PETSC_DEFAULT_INTEGER,ierr)
  call KSPSetFromOptions(psdata%ksp,ierr)

  if(mype==0 .and. istep==1)then
    write(gtcout,*)'Finish create lap-petsc matrix'
    flush(gtcout)
  endif


  return
end subroutine lapmat_pscreate
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
subroutine lapmat_pssolver(psdata,luserx_output,luserb_input)
#include "petsc/finclude/petsc.h"
  use petsc
  use precision, only:lk
  use global_parameters,only:istep,irk,mype
  use petsc_array
  implicit none

  type (petscData),intent(in) :: psdata
  real(lk),dimension(0:),intent(inout) :: luserx_output
  real(lk),dimension(0:),intent(in) :: luserb_input

  PetscScalar temp
  PetscErrorCode ierr
  PetscInt i,ilow,ihigh,isize,ilo,ihi
!  PetscInt isi
!  PetscOffset idx
!  PetscScalar vec_x(1)
  PetscScalar, pointer :: vec_x(:)

  call VecGetLocalSize(psdata%b,isize,ierr)
  call VecGetOwnershipRange(psdata%b,ilow,ihigh,ierr)
  do i=ilow,ihigh-1
    temp=luserb_input(i)
    call VecSetValue(psdata%b,i,temp,INSERT_VALUES,ierr)
  enddo
  call VecAssemblyBegin(psdata%b,ierr)
  call VecAssemblyEnd(psdata%b,ierr)

!!!!!!!!!!!!!!!!!!!diag petsc  !!!!!!!
!!!!     if(istep==1 .and. mype==0)then 
!!!   write(*,*)'checking vector b:'
!   call VecView(psdata%b,PETSC_VIEWER_STDOUT_SELF,ierr)
!   CHKERRQ(ierr)
!!!     endif
!!!     if(istep==1 .and. mype==0)then
!!!   write(*,*)'checking matrix a:'
!!!   call MatView(psdata%Apetsc,PETSC_VIEWER_STDOUT_SELF,ierr)
!!!   CHKERRQ(ierr)
!!     endif

  call KSPSetInitialGuessNonzero(psdata%ksp,Petsc_True,ierr)
  call KSPSolve(psdata%ksp,psdata%b,psdata%x,ierr)

!* 
  call VecScatterBegin(psdata%gather,psdata%x,psdata%x_recv,INSERT_VALUES, &
       SCATTER_FORWARD,ierr)
  call VecScatterEnd(psdata%gather,psdata%x,psdata%x_recv,INSERT_VALUES, &
       SCATTER_FORWARD,ierr)
!*
  call VecGetOwnershipRange(psdata%x_recv,ilo,ihi,ierr)
!  call VecGetLocalSize(psdata%x_recv,isi,ierr)
!  call VecGetArray(psdata%x_recv,vec_x,idx,ierr)
  call VecGetArrayF90(psdata%x_recv,vec_x,ierr)
  do i=ilo,ihi-1
    luserx_output(i)=vec_x(i-ilo+1)
  enddo
!  call VecRestoreArray(psdata%x,vec_x,idx,ierr)
  call VecRestoreArrayF90(psdata%x_recv,vec_x,ierr)

  return
end subroutine lapmat_pssolver
#endif !PETSC_VERSION_LT(3,8,0)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
