!> @file
!!  New routines for Poisson solver
!! @author
!! Copyright (C) 2002-2011 BigDFT group
!! This file is distributed under the terms of the
!! GNU General Public License, see ~/COPYING file
!! or http://www.gnu.org/copyleft/gpl.txt .
!! For the list of contributors, see ~/AUTHORS

!regroup the psolver from here -------------
subroutine apply_kernel(gpu,mesh,grid,wGPU,kernel,mu,&
     mpi_env,part_mpi,inplane_mpi,rho,offset,strten,zf,updaterho)
  use PSbase
  use wrapper_MPI
  use f_utils, only: f_zero
  use time_profiling, only: f_timing
  use dynamic_memory
  use workarrays
  use box
  use PStypes
  use PScore
  implicit none
  !>when true, the density is updated with the value of zf
  logical, intent(in) :: updaterho
  logical, intent(in) :: gpu !< logical variable controlling the gpu acceleration
  type(cell), intent(in) :: mesh
  type(FFT_metadata), intent(inout) :: grid
  type(mpi_environment), intent(in) :: mpi_env, part_mpi, inplane_mpi
  real(dp), dimension(grid%nd1,grid%nd2,grid%nd3/mpi_env%nproc), intent(inout) :: kernel
  real(dp), intent(in) :: mu
  type(GPU_workarrays), intent(in) :: wGPU
  !> Total integral on the supercell of the final potential on output
  real(dp), intent(in) :: offset
  real(dp), dimension(grid%m1,grid%m3*grid%n3p), intent(inout) :: rho
  !>work array for the usage of the main routine
  real(dp), dimension(grid%md1,grid%md3*2*(grid%md2/mpi_env%nproc)), intent(inout) :: zf
  real(dp), dimension(6), intent(out) :: strten !< stress tensor associated to the cell
  !local variables
  integer, dimension(3) :: n
  integer :: size1,size2,size3,switch_alg,i_stat,ierr
  !real(dp) :: pt,rh
  real(dp), dimension(:), allocatable :: zf1
  logical :: globalW
  type(GPU_workarrays) :: lGPU

  call f_routine(id='apply_kernel')

  !this routine builds the values for each process of the potential (zf), multiplying by scal
  !fill the array with the values of the charge density
  !no more overlap between planes
  !still the complex case should be defined
  call f_zero(strten)
  if (.not. gpu) then !CPU case
    call f_zero(zf)

    call to_fft_grid(grid, zf, rho, mpi_env%nproc)

     call f_timing(TCAT_PSOLV_COMPUT,'OF')
     call G_PoissonSolver(mpi_env%iproc,mpi_env%nproc,&
          part_mpi%mpi_comm,inplane_mpi%iproc,&
          inplane_mpi%mpi_comm,1,&
          grid%n1,grid%n2,grid%n3,&
          grid%nd1,grid%nd2,grid%nd3,&
          grid%md1,grid%md2,grid%md3,&
          kernel,zf,&
          grid%scal,mu**2,mesh,offset,strten)
     call f_timing(TCAT_PSOLV_COMPUT,'ON')

     if (updaterho) then
        call from_fft_grid(grid, rho, zf, mpi_env%nproc)
    end if

  else !GPU case
     call f_timing(TCAT_PSOLV_COMPUT,'ON')
     n(1)=grid%n1!kernel%ndims(1)*(2-kernel%geo(1))
     n(2)=grid%n3!kernel%ndims(2)*(2-kernel%geo(2))
     n(3)=grid%n2!kernel%ndims(3)*(2-kernel%geo(3))

     size1=grid%md1*grid%md2*grid%md3! nproc always 1 kernel%ndims(1)*kernel%ndims(2)*kernel%ndims(3)

     lGPU = wGPU
     globalW = allocated_inner_GPU_workarrays(lGPU)
     if (.not. globalW) then
        size2 = 2*grid%n1*grid%n2*grid%n3
        size3 = grid%m1*grid%n3p*grid%m3
        call alloc_inner_GPU_workarrays(lGPU, size2, size3, i_stat)
        if (i_stat /= 0) print *,'error cudamalloc',i_stat
     endif

     if (mpi_env%nproc > 1) then
        call f_zero(zf)
        call reset_gpu_data( grid%m1*grid%n3p*grid%m3,&
                            rho, lGPU%rho)

        call pad_data( lGPU%rho, lGPU%work1, grid%m1,&
                    grid%n3p,grid%m3, grid%md1,&
                    grid%md2,grid%md3);

        !TODO : gpudirect (even if this code is disabled right now)
        if (mpi_env%iproc == 0) zf1 = f_malloc0(size1,id='zf1')

        call mpi_gatherv(zf,grid%n3p*grid%md3*grid%md1,mpitype(zf),&
             zf1,lGPU%rhocounts,lGPU%rhodispls, &
             mpitype(zf),0,mpi_env%mpi_comm,ierr)

        if (mpi_env%iproc == 0) then
           !fill the GPU memory
           call reset_gpu_data(size1,zf1,lGPU%work1)

           switch_alg=0
           if (GPU_has_plan(lGPU)) then
              call cuda_3d_psolver_general(n,lGPU%plan,lGPU%work1,lGPU%work2, &
                   lGPU%k,switch_alg,lGPU%geo,grid%scal)
           else
              call cuda_3d_psolver_plangeneral(n,lGPU%work1,lGPU%work2, &
                   lGPU%k,lGPU%geo,grid%scal)
           endif

           !take data from GPU
           call get_gpu_data(size1,zf1,lGPU%work1)
        endif

        call MPI_Scatterv(zf1,lGPU%rhocounts,lGPU%rhodispls,&
             mpitype(zf1),zf,&
             grid%n3p*grid%md3*grid%md1, &
             mpitype(zf1),0,mpi_env%mpi_comm,ierr)

        if (mpi_env%iproc == 0) call f_free(zf1)
     else
        !fill the GPU memory
        if(lGPU%stay_on_gpu /= 1) then
        call reset_gpu_data( grid%m1*grid%n3p*grid%m3,&
                            rho, lGPU%rho)
        end if

        call pad_data( lGPU%rho, lGPU%work1, grid%m1,&
                    grid%n3p,grid%m3, grid%md1,&
                    grid%md2,grid%md3);

        switch_alg=0

        if (GPU_has_plan(lGPU)) then
           call cuda_3d_psolver_general(n,lGPU%plan,lGPU%work1,lGPU%work2, &
                lGPU%k,switch_alg,lGPU%geo,grid%scal)
        else
           call cuda_3d_psolver_plangeneral(n,lGPU%work1,lGPU%work2, &
                lGPU%k,lGPU%geo,grid%scal)
        endif


        !take data from GPU
        !call get_gpu_data(size1,zf,lGPU%work1)
     endif

    if (updaterho) then
        call unpad_data( lGPU%rho, lGPU%work1, grid%m1,&
                    grid%n3p,grid%m3, grid%md1,&
                    grid%md2,grid%md3);
        if(lGPU%stay_on_gpu /= 1) then
        call get_gpu_data(grid%m1*grid%n3p*grid%m3,&
                        rho,lGPU%rho)
       call synchronize()
        end if
     end if

     if (.not. globalW) then
        call cudafree(lGPU%work1)
        call cudafree(lGPU%work2)
     endif
!     call synchronize()
     call f_timing(TCAT_PSOLV_COMPUT,'OF')
  endif

  call f_release_routine()

end subroutine apply_kernel

!> calculate the integral between the array in zf and the array in rho
!! copy the zf array in pot and sum with the array pot_ion if needed
subroutine finalize_hartree_results(sumpion,gpu,wGPU,resetGPU,pot_ion,m1,m2,m3p,&
     md1,md2,md3p,rho,zf,pot,eh)
  use f_enums
  use PSbase
  use workarrays, only: GPU_workarrays
  implicit none
  !if .true. the array pot is zf+pot_ion
  !if .false., pot is only zf
  logical, intent(in) :: sumpion
  logical, intent(in) :: gpu !< logical variable controlling the gpu acceleration
  logical, intent(in) :: resetGPU
  type(GPU_workarrays), intent(in) :: wGPU
  integer, intent(in) :: m1,m2,m3p !< dimension of the grid
  integer, intent(in) :: md1,md2,md3p !< dimension of the zf array
  !> original density and final potential (can point to the same array)
  real(dp), dimension(m1,m2*m3p), intent(in) :: rho
  real(dp), dimension(m1,m2*m3p), intent(out) :: pot
  !> ionic potential, to be added to the potential
  real(dp), dimension(m1,m2*m3p), intent(in) :: pot_ion
  !> work array for the usage of the main routine
  !!(might have the same dimension of rho, but generally is bigger)
  real(dp), dimension(md1,md2*md3p), intent(in) :: zf
  !>hartree energy, being \int d^3 x \rho(x) zf(x) (no volume element)
  real(dp), intent(out) :: eh
  !local variables
  integer :: i1,i23,j23,j3,n3delta, i_stat
  real(dp) :: pt,rh


  eh=0.0_dp


  if(gpu) then

     !in VAC case, rho and zf are already on the card and untouched
     if (resetGPU .and. wGPU%stay_on_gpu /= 1) then
        call reset_gpu_data(m1*m2*m3p,rho,wGPU%rho)
        call reset_gpu_data(m1*m2*m3p,zf,wGPU%work1)
     end if

     if (sumpion) then
        if (wGPU%pot_ion==0.d0)&
             call cudamalloc(m1*m2*m3p,wGPU%pot_ion,i_stat)
        call reset_gpu_data(m1*m2*m3p,pot_ion,wGPU%pot_ion)
     end if

     call unpad_data( wGPU%work2, wGPU%work1, m1,&
          m3p,m2, md1, md3p,md2)


     if(wGPU%stay_on_gpu /= 1) then
       call finalize_reduction_kernel(sumpion,m1,m2*m3p,md1,md2*md3p, wGPU%work2,&
             wGPU%rho, wGPU%pot_ion, wGPU%reduc, wGPU%ehart,1)
       call get_gpu_data(m1*m2*m3p,pot,wGPU%rho)
       call synchronize()
     else
       !delay results retrieval to after communications have been finished
       call finalize_reduction_kernel(sumpion,m1,m2*m3p,md1,md2*md3p, wGPU%work2, &
             wGPU%rho, wGPU%pot_ion, wGPU%reduc, wGPU%ehart,0)
     end if

  else
     !write(*,*) 'iproc, m1, md2, m2, m3p', iproc, m1, md2, m2, m3p
     !write(*,*) 'm1, md1, md2, m2, md3p, m3p', m1, md1, md2, m2, md3p, m3p
     !recollect the final data
     n3delta=md2-m2 !this is the y dimension
     if (sumpion) then
        !$omp parallel do default(shared) private(i1,i23,j23,j3,pt,rh) &
        !$omp reduction(+:eh)
        do i23=1,m2*m3p
           j3=(i23-1)/m2
           j23=i23+n3delta*j3
           eh=eh+my_dot(m1,zf(1,j23),rho(1,i23))
           call my_copy_add(m1,zf(1,j23),pot_ion(1,i23),pot(1,i23))
           !!do i1=1,m1
           !!   pt=zf(i1,j23)
           !!   rh=rho(i1,i23)
           !!   rh=rh*pt
           !!   eh=eh+rh
           !!   pot(i1,i23)=pt+pot_ion(i1,i23)
           !!end do
        end do
        !$omp end parallel do
     else
        !$omp parallel do default(shared) private(i1,i23,j23,j3,pt,rh) &
        !$omp reduction(+:eh)
        do i23=1,m2*m3p
           j3=(i23-1)/m2
           j23=i23+n3delta*j3
           eh=eh+my_dot(m1,zf(1,j23),rho(1,i23))
           call my_copy(m1,zf(1,j23),pot(1,i23))
           !!do i1=1,m1
           !!   pt=zf(i1,j23)
           !!   rh=rho(i1,i23)
           !!   rh=rh*pt
           !!   eh=eh+rh
           !!   pot(i1,i23)=pt
           !!end do
        end do
        !$omp end parallel do
     end if

  end if


  contains

     pure function my_dot(n,x,y) result(tt)
       implicit none
       integer, intent(in) :: n
       double precision :: tt
       double precision, dimension(n), intent(in) :: x,y
       !local variables
       integer :: i
       tt=0.d0
       do i=1,n
          tt=tt+x(i)*y(i)
       end do
     end function my_dot

     pure subroutine my_copy(n,x,y)
       implicit none
       integer, intent(in) :: n
       double precision, dimension(n), intent(in) :: x
       double precision, dimension(n), intent(out) :: y
       !local variables
       integer :: i
       do i=1,n
          y(i)=x(i)
       end do
     end subroutine my_copy

     pure subroutine my_copy_add(n,x,y,z)
       implicit none
       integer, intent(in) :: n
       double precision, dimension(n), intent(in) :: x, y
       double precision, dimension(n), intent(out) :: z
       !local variables
       integer :: i
       do i=1,n
          z(i)=x(i)+y(i)
       end do
     end subroutine my_copy_add

end subroutine finalize_hartree_results
