module workarrays
  use f_precisions, dp => f_double, gp => f_double
  implicit none
  private

  type, public :: PCG_workarrays
     real(dp), dimension(:,:), pointer :: res,z,p,q
  end type PCG_workarrays
  private :: nullify_PCG_workarrays, free_PCG_workarrays
  
  type, public :: PI_workarrays
     real(dp), dimension(:,:), pointer :: rho
  end type PI_workarrays
  private :: nullify_PI_workarrays, free_PI_workarrays

  type, public :: low_level_workarrays
     ! Allocate the required one, depending on used method.
     type(PCG_workarrays), pointer :: PCG => null()
     type(PI_workarrays), pointer :: PI => null()
  end type low_level_workarrays
  public :: alloc_PCG_workarrays, alloc_PI_workarrays, free_low_level_workarrays

  type, public :: PCG_GPU_workarrays
     logical :: used
     
     integer(f_address) :: r,z,p,q,x
     integer(f_address) :: oneoeps,corr
     !> GPU scalars. Event if they are scalars of course their address is needed
     integer(f_address) :: alpha, beta, kappa, beta0, reduc
  end type PCG_GPU_workarrays
  public :: nullify_PCG_GPU_workarrays
  public :: alloc_PCG_GPU_workarrays, free_PCG_GPU_workarrays
  public :: allocated_PCG_GPU, init_PCG_GPU_workarrays
  
  type, public :: dielectric_arrays
     !> Order of accuracy for the finite difference nabla, used for nonvacuum calculations (only in SC).
     !! Used also for the additional contribution to the KS potential for an implicit solvation run.
     !! Represents the total number of points at left and right of the x0 where we want to calculate the derivative.
     integer :: nord

     !> dielectric function epsilon, continuous and differentiable in the whole
     !domain, to be used in the case of Preconditioned Conjugate Gradient (PCG) method
     real(dp), dimension(:,:), pointer :: eps
     !> logaritmic derivative of the dielectric function epsilon,
     !! to be used in the case of Polarization Iteration (SC) method
     real(dp), dimension(:,:,:,:), pointer :: dlogeps
     !> inverse of the dielectric function
     !! in the case of Polarization Iteration (SC) method
     !! inverse of the square root of epsilon
     !! in the case of the Preconditioned Conjugate Gradient (PCG)
     real(dp), dimension(:,:), pointer :: oneoeps
     !> correction term, given in terms of the multiplicative factor of nabla*eps*nabla
     !! to be used for Preconditioned Conjugate Gradient
     real(dp), dimension(:,:), pointer :: corr
     !> inner rigid cavity to be integrated in the sccs method to avoit inner
     !! cavity discontinuity due to near-zero edens near atoms
     real(dp), dimension(:,:), pointer :: epsinnersccs

     integer :: nat !< dimensions of the atomic based cavity. Zero if unused
     !>positions of the atoms of the cavity in the simulation box
     real(dp), dimension(:,:), pointer :: rxyz !< atomic coordinates
     real(dp), dimension(:), pointer :: radii !<radii of the cavity per atom

     real(dp) :: IntVol     !< Volume integral needed for the non-electrostatic energy contributions
     real(dp) :: IntSur     !< Surface integral needed for the non-electrostatic energy contributions
  end type dielectric_arrays
  public :: nullify_dielectric_arrays, free_dielectric_arrays
  public :: alloc_dielectric_arrays_forPCG, alloc_dielectric_arrays_forPI
  public :: set_dielectric_arrays_forPCG, set_dielectric_arrays_forPI
  public :: dielectric_arrays_forPI, dielectric_arrays_forPCG
  public :: dielectric_PCG_set_at, dielectric_PI_set_at
  public :: dielectric_PCG_reduce_beta, dielectric_PCG_reduce_kappa
  public :: dielectric_PCG_reduce_norm

  type, public :: grid_observables
     !> input guess vectors to be preserved for future use
     !!or work arrays, might be of variable dimension
     !!(either full of distributed)
     real(dp), dimension(:,:), pointer :: pot
     !> Polarization charge vector for print purpose only.
     real(dp), dimension(:,:), pointer :: rho_pol
     !> ionic density, in the case of a Poisson-Boltzmann (PB) approach
     real(dp), dimension(:,:), pointer :: rho_ions
  end type grid_observables
  public :: nullify_grid_observables, free_grid_observables
  public :: ensure_grid_observables_allocations

  type, public :: GPU_workarrays
     integer, dimension(3) :: geo
     integer :: stay_on_gpu

     integer(f_address) :: work1, work2, rho, pot_ion, k !<addresses for the GPU memory
     !> GPU scalars. Event if they are scalars of course their address is needed
     integer(f_address) :: eexctX, reduc, ehart

     !> Arrays needed to gather the information of the Poisson solver on multiple gpus
     integer, dimension(:), pointer :: rhocounts, rhodispls

     integer, dimension(5) :: plan
  end type GPU_workarrays
  public :: nullify_GPU_workarrays, free_GPU_workarrays
  public :: alloc_GPU_workarrays, alloc_inner_GPU_workarrays, allocated_inner_GPU_workarrays
  public :: create_GPU_rho_distribution, create_GPU_plans
  public :: GPU_has_plan
  
contains

  pure subroutine nullify_PCG_workarrays(w)
    implicit none
    type(PCG_workarrays), intent(out) :: w
    nullify(w%res)
    nullify(w%z)
    nullify(w%p)
    nullify(w%q)
  end subroutine nullify_PCG_workarrays

  subroutine free_PCG_workarrays(w)
    use dynamic_memory
    implicit none
    type(PCG_workarrays), intent(inout) :: w    
    call f_free_ptr(w%res)
    call f_free_ptr(w%q)
    call f_free_ptr(w%p)
    call f_free_ptr(w%z)
  end subroutine free_PCG_workarrays

  subroutine alloc_PCG_workarrays(w, rho, n1, n23)
    use dynamic_memory
    implicit none
    type(low_level_workarrays), intent(out) :: w
    integer, intent(in) :: n1, n23
    real(dp), dimension(n1, n23), intent(in) :: rho

    allocate(w%PCG)
    call nullify_PCG_workarrays(w%PCG)
    
    w%PCG%res = f_malloc_ptr([n1, n23], id = 'res')
    call f_memcpy(src = rho, dest = w%PCG%res)

    w%PCG%q = f_malloc0_ptr([n1, n23], id = 'q')
    w%PCG%p = f_malloc0_ptr([n1, n23], id = 'p')
    w%PCG%z = f_malloc_ptr([n1, n23], id = 'z')
  end subroutine alloc_PCG_workarrays
     
  pure subroutine nullify_PI_workarrays(w)
    implicit none
    type(PI_workarrays), intent(out) :: w
    nullify(w%rho)
  end subroutine nullify_PI_workarrays

  subroutine free_PI_workarrays(w)
    use dynamic_memory
    implicit none
    type(PI_workarrays), intent(inout) :: w    
    call f_free_ptr(w%rho)
  end subroutine free_PI_workarrays

  subroutine alloc_PI_workarrays(w, ndims)
    use dynamic_memory
    implicit none
    type(low_level_workarrays), intent(out) :: w
    integer, dimension(3), intent(in) :: ndims

    allocate(w%PI)
    call nullify_PI_workarrays(w%PI)
    w%PI%rho = f_malloc0_ptr([ndims(1), ndims(2)*ndims(3)], id = 'rho')
  end subroutine alloc_PI_workarrays

  subroutine free_low_level_workarrays(w)
    implicit none
    type(low_level_workarrays), intent(inout) :: w

    if (associated(w%PCG)) then
       call free_PCG_workarrays(w%PCG)
       nullify(w%PCG)
    end if
    if (associated(w%PI)) then
       call free_PI_workarrays(w%PI)
       nullify(w%PI)
    end if
  end subroutine free_low_level_workarrays

  pure subroutine nullify_PCG_GPU_workarrays(w)
    use f_utils, only: f_zero
    implicit none
    type(PCG_GPU_workarrays), intent(out) :: w

    w%used = .false.
    call f_zero(w%r)
    call f_zero(w%z)
    call f_zero(w%p)
    call f_zero(w%q)
    call f_zero(w%x)
    call f_zero(w%corr)
    call f_zero(w%oneoeps)
    call f_zero(w%alpha)
    call f_zero(w%beta)
    call f_zero(w%kappa)
    call f_zero(w%beta0)
    call f_zero(w%reduc)
  end subroutine nullify_PCG_GPU_workarrays
  
  subroutine free_PCG_GPU_workarrays(w)
    implicit none
    type(PCG_GPU_workarrays), intent(inout) :: w    

    if (.not. w%used) return

    call cudafree(w%z)
    call cudafree(w%r)
    call cudafree(w%p)
    call cudafree(w%q)
    call cudafree(w%x)
    call cudafree(w%oneoeps)
    call cudafree(w%corr)
    call cudafree(w%alpha)
    call cudafree(w%beta)
    call cudafree(w%beta0)
    call cudafree(w%kappa)
    call cudafree(w%reduc)
    call nullify_PCG_GPU_workarrays(w)
    w%used = .true.
  end subroutine free_PCG_GPU_workarrays

  subroutine alloc_PCG_GPU_workarrays(w, nsize, i_stat)
    use dynamic_memory
    implicit none
    type(PCG_GPU_workarrays), intent(inout) :: w
    integer, intent(in) :: nsize
    integer, intent(out) :: i_stat
    real(dp) :: alpha

    if (.not. w%used) return

    call cudamalloc(nsize, w%z, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(nsize, w%r, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(nsize, w%p, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(nsize, w%q, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(nsize, w%x, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(nsize, w%oneoeps, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(nsize, w%corr, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(sizeof(alpha), w%alpha, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(sizeof(alpha), w%beta, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(sizeof(alpha), w%beta0, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(sizeof(alpha), w%kappa, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(sizeof(alpha), w%reduc, i_stat)
    if (i_stat /= 0) return
  end subroutine alloc_PCG_GPU_workarrays

  pure function allocated_PCG_GPU(w) result(alloc)
    implicit none
    type(PCG_GPU_workarrays), intent(in) :: w
    logical :: alloc

    ! Check only the first one, since allocated by bunch.
    alloc = (w%z /= 0)
  end function allocated_PCG_GPU

  subroutine init_PCG_GPU_workarrays(w, diel, res)
    implicit none
    type(PCG_GPU_workarrays), intent(in) :: w
    type(dielectric_arrays), intent(in) :: diel
    real(dp), dimension(:,:), intent(in) :: res

    integer :: i_stat, nsize

    if (.not. w%used) return
    
    nsize = product(shape(diel%oneoeps))
    call cudamemset(w%p, 0, nsize, i_stat)
    if (i_stat /= 0) print *,'error cudamemset p',i_stat
    call cudamemset(w%q, 0, nsize, i_stat)
    if (i_stat /= 0) print *,'error cudamemset q',i_stat
    call cudamemset(w%x, 0, nsize, i_stat)
    if (i_stat /= 0) print *,'error cudamemset x',i_stat
    call reset_gpu_data(nsize, res, w%r)
    call reset_gpu_data(nsize, diel%oneoeps, w%oneoeps)
    call reset_gpu_data(nsize, diel%corr, w%corr)
  end subroutine init_PCG_GPU_workarrays
     
  pure subroutine nullify_dielectric_arrays(w)
    implicit none
    type(dielectric_arrays), intent(out) :: w
    nullify(w%eps)
    nullify(w%dlogeps)
    nullify(w%oneoeps)
    nullify(w%corr)
    nullify(w%epsinnersccs)
    w%nat = 0
    nullify(w%radii)
    nullify(w%rxyz)
    w%IntVol = 0._dp
    w%IntSur = 0._dp
  end subroutine nullify_dielectric_arrays

  subroutine free_dielectric_arrays(w)
    use dynamic_memory
    implicit none
    type(dielectric_arrays), intent(inout) :: w
    call f_free_ptr(w%eps)
    call f_free_ptr(w%dlogeps)
    call f_free_ptr(w%oneoeps)
    call f_free_ptr(w%corr)
    call f_free_ptr(w%epsinnersccs)
    call f_free_ptr(w%radii)
    call f_free_ptr(w%rxyz)
  end subroutine free_dielectric_arrays

  subroutine alloc_dielectric_cavity(w, nat, rxyz, radii)
    use dynamic_memory
    implicit none
    type(dielectric_arrays), intent(inout) :: w
    integer, intent(in) :: nat
    real(gp), dimension(3, nat), intent(in) :: rxyz
    real(gp), dimension(nat), intent(in) :: radii

    w%IntVol = 0._dp
    w%IntSur = 0._dp
    w%nat = nat
    w%rxyz = f_malloc_ptr((/ 3, nat /), id = 'rxyz')
    if (nat > 0) call f_memcpy(n = 3 * nat, src = rxyz(1, 1), dest = w%rxyz(1, 1))
    w%radii = f_malloc_ptr(nat, id = 'radii')
    if (nat > 0) call f_memcpy(n = nat, src = radii(1), dest = w%radii(1))
  end subroutine alloc_dielectric_cavity

  subroutine alloc_dielectric_arrays_forPCG(w, n1, n23)
    use dynamic_memory
    implicit none
    type(dielectric_arrays), intent(out) :: w
    integer, intent(in) :: n1, n23

    call nullify_dielectric_arrays(w)
    !check the dimensions of the associated arrays
    w%eps = f_malloc_ptr((/ n1, n23 /), id = "w%oneoeps")
    w%oneoeps = f_malloc_ptr((/ n1, n23 /), id = "w%oneoeps")
    w%corr = f_malloc_ptr((/ n1, n23 /), id = "w%corr")
    w%epsinnersccs = f_malloc_ptr((/ n1, n23 /), id = 'epsinnersccs')
  end subroutine alloc_dielectric_arrays_forPCG

  subroutine normalize_dielectric_integrals(w, epsm1, hh)
    implicit none
    type(dielectric_arrays), intent(inout) :: w
    real(gp), intent(in) :: epsm1, hh

    if (epsm1 /= 0._dp) then
       w%IntVol = w%IntVol * hh / epsm1
       w%IntSur = w%IntSur * hh / epsm1
    else
       w%IntVol = 0._dp
       w%IntSur = 0._dp
    end if
  end subroutine normalize_dielectric_integrals

  subroutine set_dielectric_arrays_forPCG(w, n1, n23, i3s, n3p, mesh, &
       & eps, oneosqrteps, corr, nat, rxyz, radii, cavity, hh)
    use dictionaries
    use yaml_strings
    use box
    use FdDer
    use psolver_environment, only: cavity_data, rigid_cavity_arrays, vacuum_eps
    use numerics, only: pi
    use dynamic_memory
    implicit none
    type(dielectric_arrays), intent(inout) :: w
    integer, intent(in) :: n1, n23, i3s, n3p
    type(cell), intent(in) :: mesh
    !> dielectric function. Needed for non VAC methods, given in full dimensions
    real(dp), dimension(:,:,:), intent(in), optional :: eps
    !! if absent, it will be calculated from the array of epsilon
    real(dp), dimension(:,:,:), intent(in), optional :: corr
    !> inverse square root of epsilon. Needed for PCG method.
    !! if absent, it will be calculated from the array of epsilon
    real(dp), dimension(:,:,:), intent(in), optional :: oneosqrteps
    !> number of atoms, can be inserted to define the rigid cavity
    integer, intent(in), optional :: nat
    !> if nat is present, also the positions of the atoms have to be defined
    !! (dimension 3,nat). The rxyz values are defined accordingly to the
    !! position of the atoms in the grid starting from the point [1,1,1] to ndims(:)
    real(gp), dimension(*), intent(in), optional :: rxyz
    !> and the radii aroud each atoms also have to be defined (dimension nat)
    real(gp), dimension(*), intent(in), optional :: radii
    type(cavity_data), intent(in), optional :: cavity
    real(gp), intent(in), optional :: hh

    integer :: i1, i2, i3, i23
    real(dp) :: cc,ep,depsr,kk
    real(dp), dimension(3) :: v,dleps
    real(dp), dimension(:,:,:), allocatable :: de2,ddeps
    real(dp), dimension(:,:,:,:), allocatable :: deps

    if (present(corr)) then
       !check the dimensions (for the moment no parallelism)
       if (any(shape(corr) /= mesh%ndims)) &
            call f_err_throw('Error in the dimensions of the array corr,'//&
            trim(yaml_toa(shape(corr))))
       call f_memcpy(n=n1*n23,src=corr(1,1,i3s),dest=w%corr)
    else if (present(eps)) then
       !check the dimensions (for the moment no parallelism)
       if (any(shape(eps) /= mesh%ndims)) &
            call f_err_throw('Error in the dimensions of the array epsilon,'//&
            trim(yaml_toa(shape(eps))))
       !allocate work arrays
       deps=f_malloc([mesh%ndims(1),mesh%ndims(2),mesh%ndims(3),3],id='deps')
       de2 =f_malloc(mesh%ndims,id='de2')
       ddeps=f_malloc(mesh%ndims,id='ddeps')

       call nabla_u_and_square(mesh,eps,deps,de2,w%nord)

       call div_u_i(mesh,deps,ddeps,w%nord)
       i23=1
       do i3=i3s,i3s+n3p-1!kernel%ndims(3)
          do i2=1,mesh%ndims(2)
             do i1=1,mesh%ndims(1)
                w%corr(i1,i23)=(-0.125d0/pi)*&
                     (0.5d0*de2(i1,i2,i3)/eps(i1,i2,i3)-ddeps(i1,i2,i3))
             end do
             i23=i23+1
          end do
       end do
       call f_free(deps)
       call f_free(ddeps)
       call f_free(de2)
    else if (present(nat) .and. present(rxyz) .and. present(radii) .and. present(cavity)) then
       call alloc_dielectric_cavity(w, nat, rxyz, radii)
       do i3=i3s,n3p+i3s-1
          v(3)=cell_r(mesh,i3,dim=3)
          do i2=1,mesh%ndims(2)
             v(2)=cell_r(mesh,i2,dim=2)
             i23=i2+(i3-i3s)*mesh%ndims(2)
             do i1=1,mesh%ndims(1)
                v(1)=cell_r(mesh,i1,dim=1)
                call rigid_cavity_arrays(cavity,mesh%dom,v,&
                     nat,rxyz,radii,ep,depsr,dleps,cc,kk)
                w%eps(i1,i23)=ep
                w%corr(i1,i23)=cc
                w%IntVol = w%IntVol + (cavity%epsilon0 - ep)
                w%IntSur = w%IntSur + depsr
             end do
          end do
       end do
       call normalize_dielectric_integrals(w, cavity%epsilon0 - vacuum_eps, hh)
    else
       call f_err_throw('For method "PCG" the arrays corr or epsilon should be present')
    end if

    if (present(oneosqrteps)) then
       !check the dimensions (for the moment no parallelism)
       if (any(shape(oneosqrteps) /= mesh%ndims)) &
            call f_err_throw('Error in the dimensions of the array oneosqrteps,'//&
            trim(yaml_toa(shape(oneosqrteps))))
       call f_memcpy(n=n1*n23,src=oneosqrteps(1,1,i3s),dest=w%oneoeps)
    else if (present(eps)) then
       i23=1
       do i3=i3s,i3s+n3p-1!kernel%ndims(3)
          do i2=1,mesh%ndims(2)
             do i1=1,mesh%ndims(1)
                w%oneoeps(i1,i23)=1.0_dp/sqrt(eps(i1,i2,i3))
             end do
             i23=i23+1
          end do
       end do
    else if (present(nat) .and. present(rxyz) .and. present(radii)) then
       do i23=1,n23
          do i1=1,n1
             w%oneoeps(i1,i23)=1.d0/sqrt(w%eps(i1,i23)) !here square root is correct
          end do
       end do
    else
       call f_err_throw('For method "PCG" the arrays oneosqrteps or epsilon should be present')
    end if
    
    if (present(eps)) then
       call f_memcpy(n=n1*n23,src=eps(1,1,i3s),dest=w%eps)
    else if (.not. present(nat) .or. .not. present(rxyz) .or. .not. present(radii)) then
       call f_err_throw('For method "PCG" the arrays eps should be present')
    end if
  end subroutine set_dielectric_arrays_forPCG

  function dielectric_PCG_reduce_beta(diel, PCG, GPU) result(beta)
    implicit none
    type(dielectric_arrays), intent(in) :: diel
    type(PCG_workarrays), intent(inout) :: PCG
    type(PCG_GPU_workarrays), intent(in) :: GPU
    real(dp) :: beta, zeta, rval
    integer :: i1, i23, i_stat

    beta = 0._dp
    if (.not. GPU%used) then
       !$omp parallel do default(shared) private(i1,i23,rval,zeta) &
       !$omp reduction(+:beta)
       do i23=1,size(PCG%z, 2)
          do i1=1,size(PCG%z, 1)
             zeta=PCG%z(i1,i23)
             zeta=zeta*diel%oneoeps(i1,i23)
             rval=PCG%res(i1,i23)
             rval=rval*zeta
             beta=beta+rval
             PCG%z(i1,i23)=zeta
          end do
       end do
       !$omp end parallel do
    else
       call reset_gpu_data(size(PCG%z, 1) * size(PCG%z, 2), PCG%z, GPU%z)
       call cudamemset(GPU%beta, 0, sizeof(beta),i_stat)
       call first_reduction_kernel(size(PCG%z, 1),size(PCG%z, 2),GPU%p,GPU%q,GPU%r,&
            GPU%x,GPU%z,GPU%corr, GPU%oneoeps, GPU%alpha,&
            GPU%beta, GPU%beta0, GPU%kappa, GPU%reduc, beta)
    end if
  end function dielectric_PCG_reduce_beta

  function dielectric_PCG_reduce_kappa(diel, PCG, GPU, beta, beta0) result(kappa)
    implicit none
    type(dielectric_arrays), intent(in) :: diel
    type(PCG_workarrays), intent(inout) :: PCG
    type(PCG_GPU_workarrays), intent(in) :: GPU
    real(dp), intent(in) :: beta, beta0
    real(dp) :: kappa, zeta, rval, epsc, pval, qval, beta_tilde
    integer :: i1, i23, i_stat

    kappa=0.d0
    if (.not. GPU%used) then
       beta_tilde = beta / beta0
       !$omp parallel do default(shared) private(i1,i23,epsc,zeta)&
       !$omp private(pval,qval,rval) reduction(+:kappa)
       do i23=1,size(PCG%z, 2)
          do i1=1,size(PCG%z, 1)
             zeta=PCG%z(i1,i23)
             epsc=diel%corr(i1,i23)
             pval=PCG%p(i1,i23)
             qval=PCG%q(i1,i23)
             rval=PCG%res(i1,i23)
             pval = zeta+beta_tilde*pval
             qval = zeta*epsc+rval+beta_tilde*qval
             PCG%p(i1,i23) = pval
             PCG%q(i1,i23) = qval
             rval=pval*qval
             kappa = kappa+rval 
          end do
       end do
       !$omp end parallel do
    else
       call cudamemset(GPU%kappa, 0, sizeof(kappa), i_stat)

       call reset_gpu_data(sizeof(beta), beta, GPU%beta)
       call reset_gpu_data(sizeof(beta0), beta0, GPU%beta0)

       call second_reduction_kernel(size(PCG%z, 1),size(PCG%z, 2),GPU%p,GPU%q,GPU%r,&
            GPU%x,GPU%z,GPU%corr, GPU%oneoeps, GPU%alpha,&
            GPU%beta, GPU%beta0, GPU%kappa, GPU%reduc, kappa)

       !  call get_gpu_data(size1,p,kernel%p_GPU)
       !  call get_gpu_data(size1,q,kernel%q_GPU)
    end if
  end function dielectric_PCG_reduce_kappa
  
  function dielectric_PCG_reduce_norm(diel, x, PCG, GPU, alpha) result(normr)
    implicit none
    type(dielectric_arrays), intent(in) :: diel
    real(dp), dimension(:,:), intent(inout) :: x
    type(PCG_workarrays), intent(inout) :: PCG
    type(PCG_GPU_workarrays), intent(in) :: GPU
    real(dp), intent(in) :: alpha
    real(dp) :: normr
    integer :: i1, i23

    normr=0.d0
    if (.not. GPU%used) then
       !$omp parallel do default(shared) private(i1,i23) &
       !$omp reduction(+:normr)
       do i23=1,size(PCG%z, 2)
          do i1=1,size(PCG%z, 1)
             x(i1,i23) = x(i1,i23) + alpha*PCG%p(i1,i23)
             PCG%res(i1,i23) = PCG%res(i1,i23) - alpha*PCG%q(i1,i23)
             PCG%z(i1,i23) = PCG%res(i1,i23)*diel%oneoeps(i1,i23)
             normr=normr+PCG%res(i1,i23)*PCG%res(i1,i23)
          end do
       end do
       !$omp end parallel do
    else
       call reset_gpu_data(sizeof(alpha), alpha, GPU%alpha)
       call third_reduction_kernel(size(PCG%z, 1),size(PCG%z, 2),GPU%p,GPU%q,GPU%r,&
            GPU%x,GPU%z,GPU%corr, GPU%oneoeps, GPU%alpha,&
            GPU%beta, GPU%beta0, GPU%kappa, GPU%reduc, normr)
       call get_gpu_data(size(PCG%z, 1) * size(PCG%z, 2), PCG%z, GPU%z)
    end if
  end function dielectric_PCG_reduce_norm
  
  subroutine alloc_dielectric_arrays_forPI(w, n1, n23, ndims)
    use dynamic_memory
    implicit none
    type(dielectric_arrays), intent(out) :: w
    integer, intent(in) :: n1, n23
    integer, dimension(3), intent(in) :: ndims

    call nullify_dielectric_arrays(w)
    !check the dimensions of the associated arrays
    w%eps = f_malloc_ptr((/ n1, n23 /), id = "w%oneoeps")
    w%oneoeps = f_malloc_ptr((/ n1, n23 /), id = "w%oneoeps")
    w%dlogeps = f_malloc_ptr((/ 3, ndims(1), ndims(2), ndims(3) /), id = "w%dlogeps")
    w%epsinnersccs = f_malloc_ptr((/ n1, n23 /), id = 'epsinnersccs')
  end subroutine alloc_dielectric_arrays_forPI

  subroutine set_dielectric_arrays_forPI(w, n1, n23, i3s, n3p, mesh, &
       & eps, oneoeps, dlogeps, nat, rxyz, radii, cavity, hh)
    use dictionaries
    use yaml_strings
    use box
    use psolver_environment, only: cavity_data, rigid_cavity_arrays, vacuum_eps
    use FdDer
    use numerics, only: pi
    use dynamic_memory
    implicit none
    type(dielectric_arrays), intent(inout) :: w
    integer, intent(in) :: n1, n23, i3s, n3p
    type(cell), intent(in) :: mesh
    !> dielectric function. Needed for non VAC methods, given in full dimensions
    real(dp), dimension(:,:,:), intent(in), optional :: eps
    !> logarithmic derivative of epsilon. Needed for SC method.
    !! if absent, it will be calculated from the array of epsilon
    real(dp), dimension(:,:,:,:), intent(in), optional :: dlogeps
    !> inverse of epsilon. Needed for SC method.
    !! if absent, it will be calculated from the array of epsilon
    real(dp), dimension(:,:,:), intent(in), optional :: oneoeps
    !> number of atoms, can be inserted to define the rigid cavity
    integer, intent(in), optional :: nat
    !> if nat is present, also the positions of the atoms have to be defined
    !! (dimension 3,nat). The rxyz values are defined accordingly to the
    !! position of the atoms in the grid starting from the point [1,1,1] to ndims(:)
    real(gp), dimension(*), intent(in), optional :: rxyz
    !> and the radii aroud each atoms also have to be defined (dimension nat)
    real(gp), dimension(*), intent(in), optional :: radii
    type(cavity_data), intent(in), optional :: cavity
    real(gp), intent(in), optional :: hh

    integer :: i1, i2, i3, i23
    real(dp) :: cc,ep,depsr,kk,intVol,intSurf
    real(dp), dimension(3) :: v,dleps
    real(dp), dimension(:,:,:,:), allocatable :: deps

    intVol = 0._dp
    intSurf = 0._dp
    if (present(dlogeps)) then
       !check the dimensions (for the moment no parallelism)
       if (any(shape(dlogeps) /= &
            [3,mesh%ndims(1),mesh%ndims(2),mesh%ndims(3)])) &
            call f_err_throw('Error in the dimensions of the array dlogeps,'//&
            trim(yaml_toa(shape(dlogeps))))
       call f_memcpy(src=dlogeps,dest=w%dlogeps)
    else if (present(eps)) then
       !check the dimensions (for the moment no parallelism)
       if (any(shape(eps) /= mesh%ndims)) &
            call f_err_throw('Error in the dimensions of the array epsilon,'//&
            trim(yaml_toa(shape(eps))))
       !allocate arrays
       deps=f_malloc([mesh%ndims(1),mesh%ndims(2),mesh%ndims(3),3],id='deps')
       call nabla_u(mesh,eps,deps,w%nord)
       do i3=1,mesh%ndims(3)
          do i2=1,mesh%ndims(2)
             do i1=1,mesh%ndims(1)
                !switch and create the logarithmic derivative of epsilon
                w%dlogeps(1,i1,i2,i3)=deps(i1,i2,i3,1)/eps(i1,i2,i3)
                w%dlogeps(2,i1,i2,i3)=deps(i1,i2,i3,2)/eps(i1,i2,i3)
                w%dlogeps(3,i1,i2,i3)=deps(i1,i2,i3,3)/eps(i1,i2,i3)
             end do
          end do
       end do
       call f_free(deps)
    else if (present(nat) .and. present(rxyz) .and. present(radii) .and. present(cavity)) then
       call alloc_dielectric_cavity(w, nat, rxyz, radii)
       do i3=1,mesh%ndims(3)
          v(3)=cell_r(mesh,i3,dim=3)
          do i2=1,mesh%ndims(2)
             v(2)=cell_r(mesh,i2,dim=2)
             i23=i2+(i3-i3s)*mesh%ndims(2)
             do i1=1,mesh%ndims(1)
                v(1)=cell_r(mesh,i1,dim=1)
                call rigid_cavity_arrays(cavity,mesh%dom,v,nat,&
                     rxyz,radii,ep,depsr,dleps,cc,kk)
                if (i23 <= n23 .and. i23 >=1) then
                   w%eps(i1,i23)=ep
                   w%IntVol = w%IntVol + (cavity%epsilon0 - ep)
                   w%IntSur = w%IntSur + depsr
                end if
                w%dlogeps(:,i1,i2,i3)=dleps
             end do
          end do
       end do
       call normalize_dielectric_integrals(w, cavity%epsilon0 - vacuum_eps, hh)
    else
       call f_err_throw('For method "PI" the arrays dlogeps or epsilon should be present')
    end if

    if (present(oneoeps)) then
       !check the dimensions (for the moment no parallelism)
       if (any(shape(oneoeps) /= mesh%ndims)) &
            call f_err_throw('Error in the dimensions of the array oneoeps,'//&
            trim(yaml_toa(shape(oneoeps))))
       call f_memcpy(n=n1*n23,src=oneoeps(1,1,i3s),dest=w%oneoeps)
    else if (present(eps)) then
       i23=1
       do i3=i3s,i3s+n3p-1!kernel%ndims(3)
          do i2=1,mesh%ndims(2)
             do i1=1,mesh%ndims(1)
                w%oneoeps(i1,i23)=1.0_dp/eps(i1,i2,i3)
             end do
             i23=i23+1
          end do
       end do
    else if (present(nat) .and. present(rxyz) .and. present(radii)) then
       do i23=1,n23
          do i1=1,n1
             w%oneoeps(i1,i23)=1.d0/w%eps(i1,i23)
          end do
       end do
    else
       call f_err_throw('For method "PI" the arrays oneoeps or epsilon should be present')
    end if

    if (present(eps)) then
       call f_memcpy(n=n1*n23,src=eps(1,1,i3s),dest=w%eps)
    else if (.not. present(nat) .or. .not. present(rxyz) .or. .not. present(radii)) then
       call f_err_throw('For method "PI" the arrays eps should be present')
    end if
  end subroutine set_dielectric_arrays_forPI

  pure function dielectric_arrays_forPI(w) result(ok)
    implicit none
    type(dielectric_arrays), intent(in) :: w
    logical :: ok

    ok = associated(w%oneoeps) .and. associated(w%dlogeps)
  end function dielectric_arrays_forPI

  pure function dielectric_arrays_forPCG(w) result(ok)
    implicit none
    type(dielectric_arrays), intent(in) :: w
    logical :: ok

    ok = associated(w%oneoeps) .and. associated(w%corr)
  end function dielectric_arrays_forPCG

  subroutine dielectric_PCG_set_at(w, i1, i23, innervalue, cavity, edens, ddt_edens, nabla_nrm2, cc, epscur, depsdrho, dsurfdrho)
    use psolver_environment
    implicit none
    type(dielectric_arrays), intent(inout) :: w
    integer, intent(in) :: i1, i23
    real(gp), intent(in) :: innervalue
    type(cavity_data), intent(in) :: cavity
    real(dp), intent(in) :: edens, ddt_edens, nabla_nrm2, cc
    real(dp), intent(out) :: epscur, depsdrho, dsurfdrho

    ! Check for inner sccs cavity value to fix as vacuum
    if (w%epsinnersccs(i1,i23).gt.innervalue) then
       w%oneoeps(i1,i23)=1.d0
       w%corr(i1,i23)=0.d0

       epscur = 1._dp
       depsdrho = 0._dp
       dsurfdrho = 0._dp
    else
       w%oneoeps(i1,i23) = oneosqrteps(edens, cavity)
       w%corr(i1,i23) = corr_term(edens, nabla_nrm2, ddt_edens, cavity)

       epscur = eps(edens, cavity)
       depsdrho = epsprime(edens, cavity)
       dsurfdrho = -surf_term(edens,nabla_nrm2,ddt_edens,cc,cavity)
    end if
  end subroutine dielectric_PCG_set_at

  subroutine dielectric_PI_set_at(w, i1, i23, i2, i3, innervalue, cavity, &
       & edens, ddt_edens, nabla_edens, nabla_nrm2, cc, epscur, depsdrho, dsurfdrho)
    use psolver_environment
    implicit none
    type(dielectric_arrays), intent(inout) :: w
    integer, intent(in) :: i1, i23, i2, i3
    real(gp), intent(in) :: innervalue
    type(cavity_data), intent(in) :: cavity
    real(dp), intent(in) :: edens, ddt_edens, nabla_nrm2, cc
    real(dp), dimension(3), intent(in) :: nabla_edens
    real(dp), intent(out) :: epscur, depsdrho, dsurfdrho

    real(dp) :: logepspr

    ! Check for inner sccs cavity value to fix as vacuum
    if (w%epsinnersccs(i1,i23).gt.innervalue) then
       w%oneoeps(i1,i23)=1.d0
       w%dlogeps(1,i1,i2,i3)=0.d0
       w%dlogeps(2,i1,i2,i3)=0.d0
       w%dlogeps(3,i1,i2,i3)=0.d0
       
       epscur=1.d0
       depsdrho=0.d0
       dsurfdrho=0.d0
    else
       w%oneoeps(i1,i23)=oneoeps(edens,cavity)
       logepspr = logepsprime(edens,cavity)
       w%dlogeps(1,i1,i2,i3)=nabla_edens(1) * logepspr
       w%dlogeps(2,i1,i2,i3)=nabla_edens(2) * logepspr
       w%dlogeps(3,i1,i2,i3)=nabla_edens(3) * logepspr
       
       epscur=eps(edens,cavity)
       depsdrho=epsprime(edens,cavity)
       dsurfdrho=surf_term(edens,nabla_nrm2,ddt_edens,cc,cavity)
    end if
  end subroutine dielectric_PI_set_at

  pure subroutine nullify_grid_observables(w)
    implicit none
    type(grid_observables), intent(out) :: w
    
    nullify(w%rho_pol)
    nullify(w%pot)
    nullify(w%rho_ions)
  end subroutine nullify_grid_observables

  subroutine free_grid_observables(w)
    use dynamic_memory
    implicit none
    type(grid_observables), intent(inout) :: w

    call f_free_ptr(w%rho_pol)
    call f_free_ptr(w%pot)
    call f_free_ptr(w%rho_ions)
  end subroutine free_grid_observables

  subroutine ensure_grid_observables_allocations(w, n1, n23_rho, n23_pot)
    use dynamic_memory
    implicit none
    type(grid_observables), intent(inout) :: w
    integer, intent(in) :: n1, n23_rho, n23_pot
    
    if (.not. associated(w%pot)) w%pot=f_malloc0_ptr([n1,n23_pot],id='pot')
    if (.not. associated(w%rho_pol)) w%rho_pol=f_malloc_ptr([n1,n23_rho],id='rho_pol')
  end subroutine ensure_grid_observables_allocations

  pure subroutine nullify_GPU_workarrays(w)
    use f_utils, only: f_zero
    implicit none
    type(GPU_workarrays), intent(out) :: w
    w%geo=(/0,0,0/)
    w%stay_on_gpu = 0
    call f_zero(w%work1)
    call f_zero(w%work2)
    call f_zero(w%rho)
    call f_zero(w%pot_ion)
    call f_zero(w%k)
    call f_zero(w%eexctX)
    call f_zero(w%ehart)
    call f_zero(w%reduc)
    nullify(w%rhocounts)
    nullify(w%rhodispls)
    w%plan = (/ 0, 0, 0, 0, 0 /)
  end subroutine nullify_GPU_workarrays

  subroutine free_GPU_workarrays(w)
    use dictionaries
    use dynamic_memory
    implicit none
    type(GPU_workarrays), intent(inout) :: w
    integer :: i_stat

    call cudadestroystream(i_stat)
    if (i_stat /= 0) call f_err_throw('error freeing stream ')
    call cudadestroycublashandle()
    
    call cudafree(w%k)
    call cudafree(w%ehart)
    call cudafree(w%eexctX)
    call cudafree(w%reduc)

    call cudafree(w%work1)
    call cudafree(w%work2)
    call cudafree(w%rho)
    call cudafree(w%pot_ion)

    call f_free_ptr(w%rhocounts)
    call f_free_ptr(w%rhodispls)

    if (w%plan(1) /= 0) call cufftDestroy(w%plan(1))
    if (w%plan(2) /= 0) call cufftDestroy(w%plan(2))
    if (w%plan(3) /= 0) call cufftDestroy(w%plan(3))
    if (w%plan(4) /= 0) call cufftDestroy(w%plan(4))
  end subroutine free_GPU_workarrays

  subroutine alloc_inner_GPU_workarrays(w, nsize2, nsize3, i_stat)
    type(GPU_workarrays), intent(out) :: w
    integer, intent(in) :: nsize2, nsize3
    integer, intent(out) :: i_stat

    call cudamalloc(nsize2, w%work1, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(nsize2, w%work2, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(nsize3, w%rho, i_stat)
    if (i_stat /= 0) return
  end subroutine alloc_inner_GPU_workarrays

  pure function allocated_inner_GPU_workarrays(w) result(alloc)
    implicit none
    type(GPU_workarrays), intent(in) :: w
    logical :: alloc

    alloc = (w%work1 /=0 .and. w%work2 /= 0)
  end function allocated_inner_GPU_workarrays

  subroutine alloc_GPU_workarrays(w, nsizek, geocode, i_stat)
    use dynamic_memory
    implicit none
    type(GPU_workarrays), intent(out) :: w
    integer, intent(in) :: nsizek
    character(len = 1), intent(in) :: geocode
    integer, intent(out) :: i_stat

    select case(geocode)
    case('P')
       w%geo=[1,1,1]
    case('S')
       w%geo=[1,0,1]
    case('W')
       w%geo=[0,0,1]
    case('F')
       w%geo=[0,0,0]
    end select
    
    call cudacreatestream(i_stat)
    if (i_stat /= 0) return
    call cudacreatecublashandle()
    
    call cudamalloc(nsizek, w%k, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(1, w%ehart, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(1, w%eexctX, i_stat)
    if (i_stat /= 0) return
    call cudamemset(w%eexctX, 0, 1, i_stat)
    if (i_stat /= 0) return
    call cudamalloc(64, w%reduc, i_stat)
    if (i_stat /= 0) return
  end subroutine alloc_GPU_workarrays

  subroutine create_GPU_rho_distribution(w, nproc, m2, md1, md2, md3)
    use dynamic_memory
    implicit none
    type(GPU_workarrays), intent(inout) :: w
    integer, intent(in) :: nproc, m2, md1, md2, md3
    integer :: istart, jend, displ, jproc

    displ = 0
    w%rhocounts = f_malloc_ptr([0.to.nproc-1], id='rhocounts')
    w%rhodispls = f_malloc_ptr([0.to.nproc-1], id='rhodispls')
    do jproc = 0, nproc-1
      w%rhodispls(jproc) = displ
      istart = jproc * (md2 / nproc)
      jend = min((jproc + 1) * md2 / nproc, m2)
      if (istart <= m2-1) then
         w%rhocounts(jproc) = (jend - istart) * md3 * md1
      else
         w%rhocounts(jproc) = 0
      end if
      displ = displ + w%rhocounts(jproc)
    end do
  end subroutine create_GPU_rho_distribution

  subroutine create_GPU_plans(w, n)
    implicit none
    type(GPU_workarrays), intent(inout) :: w
    integer, dimension(3), intent(in) :: n

    integer :: switch_alg
    
    call cuda_3d_psolver_general_plan(n, w%plan, switch_alg, w%geo)
  end subroutine create_GPU_plans

  pure function GPU_has_plan(w) result(hasPlan)
    implicit none
    type(GPU_workarrays), intent(in) :: w
    logical :: hasPlan

    hasPlan = (w%plan(1) /= 0 .or. w%plan(2) /= 0 .or. w%plan(3) /= 0 .or. w%plan(4) /= 0 .or. w%plan(5) /= 0)
  end function GPU_has_plan

end module workarrays
