!{\src2tex{textfont=tt}}
!!****f* ABINIT/prep_fourwf
!! NAME
!! prep_fourwf
!!
!! FUNCTION
!! this routine prepares the data to the call of fourwf.
!!
!! COPYRIGHT
!! Copyright (C) 1998-2014 ABINIT group (FBottin,MT,GZ,FDahm)
!! this file is distributed under the terms of the
!! gnu general public license, see ~abinit/COPYING
!! or http://www.gnu.org/copyleft/gpl.txt .
!! for the initials of contributors, see ~abinit/doc/developers/contributors.txt .
!!
!! INPUTS
!!  blocksize= size of block for FFT
!!  cwavef(2,npw*ndat)=planewave coefficients of wavefunction (one spinorial component?).
!!  dtfil <type(datafiles_type)>=variables related to files
!!  gs_hamk <type(gs_hamiltonian_type)>=all data for the hamiltonian at k
!!  gvnlc=matrix elements <G|Vnonlocal|C>
!!  kg_k(3,npw_k)=reduced planewave coordinates.
!!  lmnmax=if useylm=1, max number of (l,m,n) comp. over all type of psps
!!        =if useylm=0, max number of (l,n)   comp. over all type of psps
!!  mgfft=maximum size of 1d ffts
!!  mpi_enreg=informations about mpi parallelization
!!  mpsang= 1+maximum angular momentum for nonlocal pseudopotentials
!!  mpssoang= 1+maximum (spin*angular momentum) for nonlocal pseudopotentials
!!  natom=number of atoms in cell.
!!  nband_k=number of bands at this k point for that spin polarization
!!  npw_k=number of plane waves at this k point
!!  nspinor=number of spinorial components of the wavefunctions
!!  ntypat=number of types of atoms in unit cell.
!!  nvloc=final dimension of vlocal (usually 1, but 4 for non-collinear)
!!  n4,n5,n6 used for dimensionning of vlocal
!!  prtvol=control print volume and debugging output
!!  vlocal(n4,n5,n6,nvloc)= local potential in real space, on the augmented fft grid
!!
!! OUTPUT
!!  gwavef=(2,npw*ndat)=matrix elements <G|H|C>.
!!
!! SIDE EFFECTS
!!
!! PARENTS
!!      mkrho,vtowfk
!!
!! CHILDREN
!!      fourwf,gpu_fourwf,prep_index_wavef_bandpp,prep_wavef_sym_do
!!      prep_wavef_sym_undo,timab,xmpi_alltoallv
!!
!! SOURCE

#if defined HAVE_CONFIG_H
#include "config.h"
#endif

#include "abi_common.h"

subroutine prep_fourwf(rhoaug,blocksize,cwavef,wfraug,gs_hamk,iblock,ikpt,istwf_k,&
& mgfft,mpi_enreg,nband_k,npw_k,n4,n5,n6,occ_k,paral_kgb,wtk)

 use defs_basis
 use defs_abitypes
 use m_profiling
 use m_xmpi
 use m_bandfft_kpt
 use m_errors

 use m_hamiltonian, only : gs_hamiltonian_type

!This section has been created automatically by the script Abilint (TD).
!Do not modify the following lines by hand.
#undef ABI_FUNC
#define ABI_FUNC 'prep_fourwf'
 use interfaces_18_timing
 use interfaces_53_ffts
 use interfaces_66_wfs, except_this_one => prep_fourwf
!End of the abilint section

 implicit none

!Arguments ------------------------------------
!scalars
 integer,intent(in) :: blocksize,iblock,ikpt,mgfft,n4,n5,n6,nband_k
 integer,intent(in) :: npw_k,paral_kgb
 integer,intent(inout) :: istwf_k
 real(dp),intent(in) :: wtk
 type(gs_hamiltonian_type),intent(inout) :: gs_hamk
 type(mpi_type),intent(inout) :: mpi_enreg
!arrays
 real(dp),intent(in) :: occ_k(nband_k),cwavef(2,npw_k*blocksize)
 real(dp),intent(inout) :: rhoaug(n4,n5,n6),wfraug(2,n4,n5,n6)

!Local variables-------------------------------
!local variables for mpialltoallv
!local variable for bandpp and inversion by symetry of time
!scalars
 integer :: bandpp,bandpp_sym,ier,iibandpp,ikpt_this_proc,ind_occ,ind_occ1,ind_occ2,ipw
 integer :: nproc_band,nproc_fft,npw_fft
 integer :: old_me_g0=0,spaceComm=0,tim_fourwf
 integer,pointer :: idatarecv0,ndatarecv,ndatarecv_tot,ndatasend_sym
 real(dp) :: weight,weight1,weight2
 logical :: flag_inv_sym,have_to_reequilibrate
!arrays
 integer,pointer :: indices_pw_fft(:),kg_k_fft(:,:),kg_k_gather(:,:),kg_k_gather_sym(:,:)
 integer,pointer :: rdispls(:),rdispls_sym(:)
 integer,pointer :: recvcounts(:),recvcount_fft(:),recvcounts_sym(:),recvcounts_sym_tot(:)
 integer,pointer :: recvdisp_fft(:),sdispls(:),sdispls_sym(:)
 integer,pointer :: sendcounts(:),sendcount_fft(:),sendcounts_sym(:),sendcounts_sym_all(:)
 integer,pointer :: senddisp_fft(:),tab_proc(:)
 integer,allocatable :: rdisplsloc(:)
 integer,allocatable :: recvcountsloc(:),sdisplsloc(:)
 integer,allocatable :: sendcountsloc(:)
 integer,pointer :: index_wavef_band(:),index_wavef_send(:)
 real(dp) :: dummy(2,1),tsec(2)
 real(dp),allocatable :: buff_wf(:,:),cwavef_alltoall(:,:),cwavef_fft(:,:), cwavef_fft_tr(:,:)
 real(dp),allocatable :: weight_t(:),weight1_t(:),weight2_t(:)
 real(dp),pointer :: ewavef_alltoall_sym(:,:)

!no_abirules
!correspondence with abinit. here for real wf but in complex mode
!this is the index of a given band

! *************************************************************************

!====================================================================================
 nproc_band = mpi_enreg%nproc_band
 nproc_fft  = mpi_enreg%nproc_fft
 bandpp     = mpi_enreg%bandpp

 flag_inv_sym = (istwf_k==2 .and. any(gs_hamk%ngfft(7) == [401,402,312]))

 if (flag_inv_sym) then
   istwf_k         = 1
   if (modulo(bandpp,2)==0) then
     bandpp_sym   = bandpp/2
   else
     bandpp_sym   = bandpp
   end if
 end if
!====================================================================================

 spaceComm=mpi_enreg%comm_band
 ikpt_this_proc=mpi_enreg%my_kpttab(ikpt)

 ABI_ALLOCATE(sendcountsloc,(nproc_band))
 ABI_ALLOCATE(sdisplsloc   ,(nproc_band))
 ABI_ALLOCATE(recvcountsloc,(nproc_band))
 ABI_ALLOCATE(rdisplsloc   ,(nproc_band))

 recvcounts   =>bandfft_kpt(ikpt_this_proc)%recvcounts(:)
 sendcounts   =>bandfft_kpt(ikpt_this_proc)%sendcounts(:)
 rdispls      =>bandfft_kpt(ikpt_this_proc)%rdispls   (:)
 sdispls      =>bandfft_kpt(ikpt_this_proc)%sdispls   (:)
 ndatarecv    =>bandfft_kpt(ikpt_this_proc)%ndatarecv

 kg_k_gather          =>bandfft_kpt(ikpt_this_proc)%kg_k_gather(:,:)
 gs_hamk%gbound(:,:)  = bandfft_kpt(ikpt_this_proc)%gbound(:,:)

 if (flag_inv_sym ) then
   idatarecv0           =>bandfft_kpt(ikpt_this_proc)%idatarecv0
   ndatarecv_tot        =>bandfft_kpt(ikpt_this_proc)%ndatarecv_tot
   ndatasend_sym        =>bandfft_kpt(ikpt_this_proc)%ndatasend_sym
   kg_k_gather_sym      =>bandfft_kpt(ikpt_this_proc)%kg_k_gather_sym(:,:)
   rdispls_sym          =>bandfft_kpt(ikpt_this_proc)%rdispls_sym(:)
   recvcounts_sym       =>bandfft_kpt(ikpt_this_proc)%recvcounts_sym(:)
   recvcounts_sym_tot   =>bandfft_kpt(ikpt_this_proc)%recvcounts_sym_tot(:)
   sdispls_sym          =>bandfft_kpt(ikpt_this_proc)%sdispls_sym(:)
   sendcounts_sym       =>bandfft_kpt(ikpt_this_proc)%sendcounts_sym(:)
   sendcounts_sym_all   =>bandfft_kpt(ikpt_this_proc)%sendcounts_sym_all(:)
   tab_proc             =>bandfft_kpt(ikpt_this_proc)%tab_proc(:)
 end if

 ABI_ALLOCATE(cwavef_alltoall,(2,ndatarecv*bandpp))

 recvcountsloc(:)=recvcounts(:)*2*bandpp
 rdisplsloc(:)=rdispls(:)*2*bandpp
 sendcountsloc(:)=sendcounts(:)*2
 sdisplsloc(:)=sdispls(:)*2

 call timab(547,1,tsec)
 call xmpi_alltoallv(cwavef,sendcountsloc,sdisplsloc,cwavef_alltoall,&
& recvcountsloc,rdisplsloc,spaceComm,ier)
 call timab(547,2,tsec)

!If me_fft==0, I have the G=0 vector, but keep for the record the old value
 if (mpi_enreg%me_fft==0) then
   old_me_g0=mpi_enreg%me_g0
   mpi_enreg%me_g0=1
 end if

 tim_fourwf=16

!Eventually adjust load balancing for FFT (by changing FFT distrib)
 have_to_reequilibrate = bandfft_kpt(ikpt_this_proc)%have_to_reequilibrate
 if(have_to_reequilibrate) then
   npw_fft =  bandfft_kpt(ikpt_this_proc)%npw_fft
   sendcount_fft  => bandfft_kpt(ikpt_this_proc)%sendcount_fft(:)
   recvcount_fft  => bandfft_kpt(ikpt_this_proc)%recvcount_fft(:)
   senddisp_fft   => bandfft_kpt(ikpt_this_proc)%senddisp_fft(:)
   recvdisp_fft   => bandfft_kpt(ikpt_this_proc)%recvdisp_fft(:)
   indices_pw_fft => bandfft_kpt(ikpt_this_proc)%indices_pw_fft(:)
   kg_k_fft       => bandfft_kpt(ikpt_this_proc)%kg_k_fft(:,:)
   ABI_ALLOCATE( buff_wf, (2,ndatarecv*bandpp) ) ! for sorting cgwavef
   ABI_ALLOCATE( cwavef_fft, (2,npw_fft*bandpp) )
   if(bandpp>1) then
     ABI_ALLOCATE( cwavef_fft_tr, (2,npw_fft*bandpp))
   end if
 end if

!====================================================================
 if ((.not.(flag_inv_sym)) .and. (bandpp==1)) then

!  Compute the index of the band
   ind_occ = (iblock-1)*blocksize + mpi_enreg%me_band + 1

   if(abs(occ_k(ind_occ)) >=tol8) then

!    Compute the weight of the band
     weight=occ_k(ind_occ)*wtk/gs_hamk%ucvol

     if(have_to_reequilibrate) then
!      filling of sorted send buffers before exchange
       do ipw = 1 ,ndatarecv
         buff_wf(1:2, indices_pw_fft(ipw) ) = cwavef_alltoall(1:2,ipw)
       end do
       call xmpi_alltoallv(buff_wf,2*sendcount_fft,2*senddisp_fft,  &
&       cwavef_fft,2*recvcount_fft, 2*recvdisp_fft, mpi_enreg%comm_fft,ier)

       call fourwf(1,rhoaug,cwavef_fft,dummy,wfraug,&
&       gs_hamk%gbound,gs_hamk%gbound,&
&       istwf_k,kg_k_fft,kg_k_fft,mgfft,mpi_enreg,1,&
&       gs_hamk%ngfft,npw_fft,1,n4,n5,n6,1,paral_kgb,tim_fourwf,weight,weight,&
&       use_gpu_cuda=gs_hamk%use_gpu_cuda)
     else
       call fourwf(1,rhoaug,cwavef_alltoall,dummy,wfraug,&
&       gs_hamk%gbound,gs_hamk%gbound,&
&       istwf_k,kg_k_gather,kg_k_gather,mgfft,mpi_enreg,1,&
&       gs_hamk%ngfft,ndatarecv,1,n4,n5,n6,1,paral_kgb,tim_fourwf,weight,weight,&
&       use_gpu_cuda=gs_hamk%use_gpu_cuda)
     end if
   end if

 else if ((.not.(flag_inv_sym)) .and. (bandpp>1) ) then

!  -------------------------------------------------------------
!  Computation of the index to class the waves functions below bandpp
!  -------------------------------------------------------------
   call prep_index_wavef_bandpp(nproc_band,bandpp,&
&   1,ndatarecv,&
&   recvcounts,rdispls,&
&   index_wavef_band)

!  -------------------------------------------------------
!  Sorting of the wave functions below bandpp
!  -------------------------------------------------------
   cwavef_alltoall(:,:) = cwavef_alltoall(:,index_wavef_band)

   if(have_to_reequilibrate) then
!    filling of sorted send buffers before exchange
     do iibandpp=1,bandpp
       do ipw = 1 ,ndatarecv
         buff_wf(1:2, iibandpp + bandpp*(indices_pw_fft(ipw)-1)) = &
&         cwavef_alltoall(1:2,ipw + ndatarecv*(iibandpp-1))
       end do
     end do
     call xmpi_alltoallv(buff_wf,2*bandpp*sendcount_fft,2*bandpp*senddisp_fft,  &
&     cwavef_fft_tr,2*bandpp*recvcount_fft, 2*bandpp*recvdisp_fft, mpi_enreg%comm_fft,ier)
     do iibandpp=1,bandpp
       do ipw = 1 ,npw_fft
         cwavef_fft(1:2,  ipw + npw_fft*(iibandpp -1)) = cwavef_fft_tr(1:2,  iibandpp + bandpp*(ipw-1))
       end do
     end do
   end if

!  -------------------
!  Fourier calculation
!  -------------------

   if(gs_hamk%use_gpu_cuda==1) then
     ABI_ALLOCATE(weight_t,(bandpp))
     do iibandpp=1,bandpp
!      Compute the index of the band
       ind_occ = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + iibandpp
!      Compute the weight of the band
       weight_t(iibandpp)=occ_k(ind_occ)*wtk/gs_hamk%ucvol
       if(abs(occ_k(ind_occ)) < tol8) then
         weight_t(iibandpp) = 0
       end if
     end do
!    Accumulate time because it is not done in gpu_fourwf
     call timab(240+tim_fourwf,1,tsec)
#if defined HAVE_GPU_CUDA
     call gpu_fourwf(1,rhoaug,&
&     cwavef_alltoall,&
&     dummy,wfraug,&
&     gs_hamk%gbound,gs_hamk%gbound,&
&     istwf_k,kg_k_gather,kg_k_gather,mgfft,mpi_enreg,bandpp,&
&     gs_hamk%ngfft,ndatarecv,1,n4,n5,n6,1,paral_kgb,&
&     tim_fourwf,weight_t,weight_t)
#endif
     call timab(240+tim_fourwf,2,tsec)
     ABI_DEALLOCATE(weight_t)
   else

     do iibandpp=1,bandpp

!      Compute the index of the band
       ind_occ = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + iibandpp
!      Compute the weight of the band
       weight=occ_k(ind_occ)*wtk/gs_hamk%ucvol

       if(abs(occ_k(ind_occ)) >=tol8) then
         if(have_to_reequilibrate) then
           call fourwf(1,rhoaug, &
&           cwavef_fft(:,(npw_fft*(iibandpp-1))+1:(npw_fft*iibandpp)), &
&           dummy,wfraug,&
&           gs_hamk%gbound,gs_hamk%gbound,&
&           istwf_k,kg_k_fft,kg_k_fft,mgfft,mpi_enreg,1,&
&           gs_hamk%ngfft,npw_fft,1,n4,n5,n6,1,paral_kgb,tim_fourwf,weight,weight,&
&           use_gpu_cuda=gs_hamk%use_gpu_cuda)
         else
           call fourwf(1,rhoaug,&
&           cwavef_alltoall(:,(ndatarecv*(iibandpp-1))+1:(ndatarecv*iibandpp)),&
&           dummy,wfraug,&
&           gs_hamk%gbound,gs_hamk%gbound,&
&           istwf_k,kg_k_gather,kg_k_gather,mgfft,mpi_enreg,1,&
&           gs_hamk%ngfft,ndatarecv,1,n4,n5,n6,1,paral_kgb,&
&           tim_fourwf,weight,weight)
         end if
       end if
     end do

   end if ! (use_gpu_cuda==1)

!  -----------------------------------------------------
!  Sorting waves functions below the processors
!  -----------------------------------------------------
   cwavef_alltoall(:,index_wavef_band) = cwavef_alltoall(:,:)
   ABI_DEALLOCATE(index_wavef_band)


 else if (flag_inv_sym) then

!  -------------------------------------------------------------
!  Computation of the index to class the waves functions below bandpp
!  -------------------------------------------------------------
   call prep_index_wavef_bandpp(nproc_band,bandpp,&
&   1,ndatarecv,&
&   recvcounts,rdispls,&
&   index_wavef_band)

!  -------------------------------------------------------
!  Sorting the wave functions below bandpp
!  -------------------------------------------------------
   cwavef_alltoall(:,:) = cwavef_alltoall(:,index_wavef_band)

!  ------------------------------------------------------------
!  We associate the waves functions by two
!  ------------------------------------------------------------
   call prep_wavef_sym_do(mpi_enreg,bandpp,1,&
&   ndatarecv,&
&   ndatarecv_tot,ndatasend_sym,tab_proc,&
&   cwavef_alltoall,&
&   sendcounts_sym,sdispls_sym,&
&   recvcounts_sym,rdispls_sym,&
&   ewavef_alltoall_sym,&
&   index_wavef_send)

!  ------------------------------------------------------------
!  Fourier calculation
!  ------------------------------------------------------------
   if(gs_hamk%use_gpu_cuda==1) then
     ABI_ALLOCATE(weight1_t,(bandpp_sym))
     ABI_ALLOCATE(weight2_t,(bandpp_sym))
     do iibandpp=1,bandpp_sym
!      Compute the index of the bands
       if (bandpp/=1) then
         ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp-1)
         ind_occ2 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp  )
       else
         ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + 1
         ind_occ2 = ind_occ1
       end if
!      Compute the weight of the band
       weight1_t(iibandpp) = occ_k(ind_occ1)*wtk/gs_hamk%ucvol
       weight2_t(iibandpp) = occ_k(ind_occ2)*wtk/gs_hamk%ucvol
     end do
     call timab(240+tim_fourwf,1,tsec)
#if defined HAVE_GPU_CUDA
     call gpu_fourwf(1,rhoaug,&
&     ewavef_alltoall_sym,&
&     dummy,wfraug,&
&     gs_hamk%gbound,gs_hamk%gbound,&
&     istwf_k,kg_k_gather_sym,kg_k_gather_sym,mgfft,mpi_enreg,bandpp_sym,&
&     gs_hamk%ngfft,ndatarecv_tot,1,n4,n5,n6,1,paral_kgb,&
&     tim_fourwf,weight1_t,weight2_t)
#endif
     call timab(240+tim_fourwf,2,tsec)
     ABI_DEALLOCATE(weight1_t)
     ABI_DEALLOCATE(weight2_t)
   else

     do iibandpp=1,bandpp_sym

!      Compute the index of the bands
       if (bandpp/=1) then
         ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp-1)
         ind_occ2 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + (2*iibandpp  )
       else
         ind_occ1 = (iblock-1)*blocksize + (mpi_enreg%me_band * bandpp) + 1
         ind_occ2 = ind_occ1
       end if

       weight1 = occ_k(ind_occ1)*wtk/gs_hamk%ucvol
       weight2 = occ_k(ind_occ2)*wtk/gs_hamk%ucvol

       call fourwf(1,rhoaug,&
&       ewavef_alltoall_sym(:,(ndatarecv_tot*(iibandpp-1))+1:(ndatarecv_tot*iibandpp)),&
&       dummy,wfraug,&
&       gs_hamk%gbound,gs_hamk%gbound,&
&       istwf_k,kg_k_gather_sym,kg_k_gather_sym,mgfft,mpi_enreg,1,&
&       gs_hamk%ngfft,ndatarecv_tot,1,n4,n5,n6,1,paral_kgb,&
&       tim_fourwf,weight1,weight2)

     end do

   end if ! (use_gpu_cuda==1)

!  ------------------------------------------------------------
!  We dissociate each wave function in two waves functions
!  gwavef is classed below of bandpp
!  ------------------------------------------------------------
   call prep_wavef_sym_undo(mpi_enreg,bandpp,1,&
&   ndatarecv,&
&   ndatarecv_tot,ndatasend_sym,idatarecv0,&
&   cwavef_alltoall,&
&   sendcounts_sym,sdispls_sym,&
&   recvcounts_sym,rdispls_sym,&
&   ewavef_alltoall_sym,&
&   index_wavef_send)

   ABI_DEALLOCATE(ewavef_alltoall_sym)
   ABI_DEALLOCATE(index_wavef_send)

!  -------------------------------------------------------
!  Sorting waves functions below the processors
!  -------------------------------------------------------
   cwavef_alltoall(:,index_wavef_band) = cwavef_alltoall(:,:)

   ABI_DEALLOCATE(index_wavef_band)

 end if

!====================================================================
 ! Restore the input value of istwf_k
 if (flag_inv_sym) istwf_k = 2
!====================================================================

 if (mpi_enreg%me_fft==0) mpi_enreg%me_g0=old_me_g0
 if(have_to_reequilibrate) then
   ABI_DEALLOCATE( buff_wf  )
   ABI_DEALLOCATE( cwavef_fft )
   if(bandpp > 1) then
     ABI_DEALLOCATE(cwavef_fft_tr)
   end if
 end if
 ABI_DEALLOCATE(sendcountsloc)
 ABI_DEALLOCATE(sdisplsloc)
 ABI_DEALLOCATE(recvcountsloc)
 ABI_DEALLOCATE(rdisplsloc)
 ABI_DEALLOCATE(cwavef_alltoall)

end subroutine prep_fourwf
!!***
