MODULE calculus
  ! Routine to evaluate gradients, interpolation schemes and integrals
  USE prec_const, ONLY: dp
  IMPLICIT NONE
  REAL(dp), dimension(-2:2) :: dz_usu = &
   (/  1._dp/12._dp, -2._dp/3._dp, 0._dp, 2._dp/3._dp, -1._dp/12._dp /) ! fd4 centered stencil
  REAL(dp), dimension(-2:1) :: dz_o2e = &
   (/ 1._dp/24._dp,-9._dp/8._dp, 9._dp/8._dp,-1._dp/24._dp /) ! fd4 odd to even stencil
  REAL(dp), dimension(-1:2) :: dz_e2o = &
   (/ 1._dp/24._dp,-9._dp/8._dp, 9._dp/8._dp,-1._dp/24._dp /) ! fd4 odd to even stencil
   REAL(dp), dimension(-2:2) :: dz2_usu = &
   (/-1._dp/12._dp, 4._dp/3._dp, -5._dp/2._dp, 4._dp/3._dp, -1._dp/12._dp /)! 2th derivative, 4th order (for parallel hypdiff)
   REAL(dp), dimension(-2:2) :: dz4_usu = &
   (/  1._dp, -4._dp, 6._dp, -4._dp, 1._dp /) ! 4th derivative, 2nd order (for parallel hypdiff)
   REAL(dp), dimension(-2:1) :: iz_o2e = &
   (/ -1._dp/16._dp, 9._dp/16._dp, 9._dp/16._dp, -1._dp/16._dp /) ! grid interpolation, 4th order, odd to even
   REAL(dp), dimension(-1:2) :: iz_e2o = &
   (/ -1._dp/16._dp, 9._dp/16._dp, 9._dp/16._dp, -1._dp/16._dp /) ! grid interpolation, 4th order, even to odd
  PUBLIC :: simpson_rule_z, interp_z, grad_z, grad_z4

CONTAINS

SUBROUTINE grad_z(target,local_nz,Ngz,inv_deltaz,f,ddzf)
  implicit none
  ! Compute the periodic boundary condition 4 points centered finite differences
  ! formula among staggered grid or not.
  ! not staggered : the derivative results must be on the same grid as the field
  !     staggered : the derivative is computed from a grid to the other
  INTEGER,  INTENT(IN) :: target, local_nz, Ngz
  REAL(dp), INTENT(IN) :: inv_deltaz
  COMPLEX(dp),dimension(local_nz+Ngz), INTENT(IN)  :: f
  COMPLEX(dp),dimension(local_nz),     INTENT(OUT) :: ddzf
  INTEGER :: iz
  IF(Ngz .GT. 3) THEN ! Cannot apply four points stencil on less than four points grid
    SELECT CASE(TARGET)
    CASE(1)
      CALL grad_z_o2e(local_nz,Ngz,inv_deltaz,f,ddzf)
    CASE(2)
      CALL grad_z_e2o(local_nz,Ngz,inv_deltaz,f,ddzf)
    CASE DEFAULT ! No staggered grid -> usual centered finite differences
      DO iz = 1,local_nz
       ddzf(iz) = dz_usu(-2)*f(iz  ) + dz_usu(-1)*f(iz+1) &
                 +dz_usu( 0)*f(iz+2) &
                 +dz_usu( 1)*f(iz+3) + dz_usu( 2)*f(iz+4)
      ENDDO
      ddzf(:) = ddzf(:) * inv_deltaz
    END SELECT
  ELSE
    ddzf = 0._dp
  ENDIF
CONTAINS
  SUBROUTINE grad_z_o2e(local_nz,Ngz,inv_deltaz,fo,ddzfe) ! Paruta 2018 eq (27)
    ! gives the gradient of a field from the odd grid to the even one
    implicit none
    INTEGER,  INTENT(IN) :: local_nz, Ngz
    REAL(dp), INTENT(IN) :: inv_deltaz
    COMPLEX(dp),dimension(local_nz+Ngz), INTENT(IN)  :: fo
    COMPLEX(dp),dimension(local_nz),     INTENT(OUT) :: ddzfe !
    INTEGER :: iz
    DO iz = 1,local_nz
     ddzfe(iz) = dz_o2e(-2)*fo(iz  ) + dz_o2e(-1)*fo(iz+1) &
                +dz_o2e( 0)*fo(iz+2) + dz_o2e( 1)*fo(iz+3)
    ENDDO
    ddzfe(:) = ddzfe(:) * inv_deltaz
  END SUBROUTINE grad_z_o2e

  SUBROUTINE grad_z_e2o(local_nz,Ngz,inv_deltaz,fe,ddzfo) ! n2v for Paruta 2018 eq (28)
    ! gives the gradient of a field from the even grid to the odd one
    implicit none
    INTEGER,  INTENT(IN) :: local_nz, Ngz
    REAL(dp), INTENT(IN) :: inv_deltaz
    COMPLEX(dp),dimension(local_nz+Ngz), INTENT(IN)  :: fe
    COMPLEX(dp),dimension(local_nz),     INTENT(OUT) :: ddzfo
    INTEGER :: iz
    DO iz = 1,local_nz
     ddzfo(iz) = dz_e2o(-1)*fe(iz+1) + dz_e2o(0)*fe(iz+2) &
                +dz_e2o( 1)*fe(iz+3) + dz_e2o(2)*fe(iz+4)
    ENDDO
    ddzfo(:) = ddzfo(:) * inv_deltaz
  END SUBROUTINE grad_z_e2o
END SUBROUTINE grad_z

SUBROUTINE grad_z2(local_nz,Ngz,inv_deltaz,f,ddz2f)
  ! Compute the second order fourth derivative for periodic boundary condition
  implicit none
  INTEGER, INTENT(IN)  :: local_nz, Ngz
  REAL(dp), INTENT(IN) :: inv_deltaz
  COMPLEX(dp),dimension(local_nz+Ngz), INTENT(IN)  :: f
  COMPLEX(dp),dimension(local_nz),     INTENT(OUT) :: ddz2f
  INTEGER :: iz
  IF(Ngz .GT. 3) THEN ! Cannot apply four points stencil on less than four points grid
      DO iz = 1,local_nz
       ddz2f(iz) = dz2_usu(-2)*f(iz  ) + dz2_usu(-1)*f(iz+1) &
                  +dz2_usu( 0)*f(iz+2)&
                  +dz2_usu( 1)*f(iz+3) + dz2_usu( 2)*f(iz+4)
      ENDDO
  ELSE
    ddz2f = 0._dp
  ENDIF
  ddz2f = ddz2f * inv_deltaz**2
END SUBROUTINE grad_z2


SUBROUTINE grad_z4(local_nz,Ngz,inv_deltaz,f,ddz4f)
  ! Compute the second order fourth derivative for periodic boundary condition
  implicit none
  INTEGER,  INTENT(IN) :: local_nz, Ngz
  REAL(dp), INTENT(IN) :: inv_deltaz
  COMPLEX(dp),dimension(local_nz+Ngz), INTENT(IN)  :: f
  COMPLEX(dp),dimension(local_nz),     INTENT(OUT) :: ddz4f
  INTEGER :: iz
  IF(Ngz .GT. 3) THEN ! Cannot apply four points stencil on less than four points grid
      DO iz = 1,local_nz
       ddz4f(iz) = dz4_usu(-2)*f(iz  ) + dz4_usu(-1)*f(iz+1) &
                  +dz4_usu( 0)*f(iz+2)&
                  +dz4_usu( 1)*f(iz+3) + dz4_usu( 2)*f(iz+4)
      ENDDO
  ELSE
    ddz4f = 0._dp
  ENDIF
  ddz4f = ddz4f * inv_deltaz**4
END SUBROUTINE grad_z4


SUBROUTINE interp_z(target,local_nz,Ngz,f_in,f_out)
  ! Function meant to interpolate one field defined on a even/odd z into
  !  the other odd/even z grid.
  ! If Staggered Grid flag (SG) is false, returns identity
  implicit none
  INTEGER, INTENT(IN) :: local_nz, Ngz
  INTEGER, intent(in) :: target ! target grid : 0 for even grid, 1 for odd
  COMPLEX(dp),dimension(local_nz+Ngz), INTENT(IN)  :: f_in
  COMPLEX(dp),dimension(local_nz),     INTENT(OUT) :: f_out
  SELECT CASE(TARGET)
  CASE(1) ! output on even grid
    CALL interp_o2e_z(local_nz,Ngz,f_in,f_out)
  CASE(2) ! output on odd grid
    CALL interp_e2o_z(local_nz,Ngz,f_in,f_out)
  CASE DEFAULT ! No staggered grid -> usual centered finite differences
    f_out = f_in
  END SELECT
CONTAINS
  SUBROUTINE interp_o2e_z(local_nz, Ngz,fo,fe)
   ! gives the value of a field from the odd grid to the even one
   implicit none
   INTEGER, INTENT(IN) :: local_nz, Ngz
   COMPLEX(dp),dimension(local_nz+Ngz), INTENT(IN)  :: fo
   COMPLEX(dp),dimension(local_nz),     INTENT(OUT) :: fe
   INTEGER :: iz
   ! 4th order interp
   DO iz = 1,local_nz
     fe(iz) = iz_o2e(-2)*fo(iz )  + iz_o2e(-1)*fo(iz+1) &
            + iz_o2e( 0)*fo(iz+2) + iz_o2e( 1)*fo(iz+3)
   ENDDO
  END SUBROUTINE interp_o2e_z

  SUBROUTINE interp_e2o_z(local_nz, Ngz,fe,fo)
   ! gives the value of a field from the even grid to the odd one
   implicit none
   INTEGER, INTENT(IN) :: local_nz, Ngz
   COMPLEX(dp),dimension(local_nz+Ngz), INTENT(IN)  :: fe
   COMPLEX(dp),dimension(local_nz),     INTENT(OUT) :: fo
   INTEGER :: iz
   ! 4th order interp
   DO iz = 1,local_nz
     fo(iz) = iz_e2o(-1)*fe(iz+1) + iz_e2o( 0)*fe(iz+2) &
            + iz_e2o( 1)*fe(iz+3) + iz_e2o( 2)*fe(iz+4)
   ENDDO
  END SUBROUTINE interp_e2o_z

END SUBROUTINE interp_z

SUBROUTINE simpson_rule_z(local_nz,dz,f,intf)
 ! integrate f(z) over z using the simpon's rule. Assume periodic boundary conditions (f(ize+1) = f(izs))
 !from molix BJ Frei
 USE prec_const, ONLY: dp, onethird
 USE parallel,   ONLY: num_procs_z, rank_z, comm_z, manual_0D_bcast
 USE mpi
 implicit none
 INTEGER, INTENT(IN) :: local_nz
 REAL(dp),INTENT(IN) :: dz
 complex(dp),dimension(local_nz), intent(in) :: f
 COMPLEX(dp), intent(out) :: intf
 COMPLEX(dp)              :: buffer, local_int
 INTEGER :: root, i_, iz, ierr

 ! Buil local sum using the weights of composite Simpson's rule
 local_int = 0._dp
 DO iz = 1,local_nz
   IF(MODULO(iz,2) .EQ. 1) THEN ! odd iz
     local_int = local_int + 2._dp*onethird*dz*f(iz)
   ELSE ! even iz
     local_int = local_int + 4._dp*onethird*dz*f(iz)
   ENDIF
 ENDDO
 buffer = local_int
 root = 0
 !Gather manually among the rank_z=0 processes and perform the sum
 intf = 0._dp
 IF (num_procs_z .GT. 1) THEN
     !! Everyone sends its local_sum to root = 0
     IF (rank_z .NE. root) THEN
         CALL MPI_SEND(buffer, 1 , MPI_DOUBLE_COMPLEX, root, 5678, comm_z, ierr)
     ELSE
         ! Recieve from all the other processes
         DO i_ = 0,num_procs_z-1
             IF (i_ .NE. rank_z) &
                 CALL MPI_RECV(buffer, 1 , MPI_DOUBLE_COMPLEX, i_, 5678, comm_z, MPI_STATUS_IGNORE, ierr)
                 intf = intf + buffer
         ENDDO
     ENDIF
     CALL manual_0D_bcast(intf)
 ELSE
   intf = local_int
 ENDIF

END SUBROUTINE simpson_rule_z

END MODULE calculus