! Module to compute electrostatic Coulomb potential
module MW_coulomb_keq0
      use MW_kinds, only: wp
      use MW_ewald, only: MW_ewald_t
      use MW_box, only: MW_box_t
      use MW_constants, only: pi
      use MW_electrode, only: MW_electrode_t
      use MW_ion, only: MW_ion_t
      use MW_localwork, only: MW_localwork_t
   implicit none
   private

   ! Public subroutines
   ! ------------------
   public :: diag_gradQelec_potential !grad_Q [grad_Q [U^k=0_QQ]]
   public :: gradQelec_potential !grad_Q [grad_Q [U^k=0_QQ]]
   public :: qmelt2Qelec_potential !grad_Q [U^k=0_qQ]
   public :: Qelec2Qelec_potential !grad_Q [U^k=0_QQ]
   public :: melt_forces !-grad_r [(U^k=0_qq) + (U^k=0_QQ) + (U^k=0_qQ)]
   public :: energy !(U^k=0_qq) + (U^k=0_QQ) + (U^k=0_qQ)

contains

   !================================================================================
   ! Compute the Coulomb potential felt by each electrode atoms due to all electrode atoms
   subroutine diag_gradQelec_potential(localwork, ewald, box, &
         electrodes, diag_gradQ_V)
      implicit none

      ! Parameters in
      ! -------------
      type(MW_localwork_t), intent(in) :: localwork  !< Localwork work distribution
      type(MW_ewald_t),    intent(in) :: ewald     !< Ewald summation parameters
      type(MW_box_t),      intent(in) :: box       !< Simulation box parameters

      type(MW_electrode_t), intent(in) :: electrodes(:)  !< electrode parameters

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: diag_gradQ_V(:)       !< Coulomb potential

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: i
      real(wp) :: volfactor
      real(wp) :: pot_ii
      real(wp) :: sqrpialpha
      integer :: itype, iatom

      num_elec_types = size(electrodes,1)

      ! Precompute some factors
      ! -----------------------
      sqrpialpha = sqrt(pi) / ewald%alpha

      ! Long range contribution from k=0 mode
      ! -------------------------------------
      ! Pair interactions are considered by block of pair_block_size x pair_block_size
      ! Diagonal block have a special treatment
      volfactor = 2.0_wp/box%area(3)
      pot_ii = volfactor*sqrpialpha

      do itype = 1, num_elec_types
         !$acc loop vector
         do i = 1, localwork%count_atoms(itype)
            iatom = localwork%offset_atoms(itype) + i
            diag_gradQ_V(iatom) = diag_gradQ_V(iatom) - pot_ii
         end do
         !$acc end loop
      end do
   end subroutine diag_gradQelec_potential

   !================================================================================
   ! Compute the Coulomb potential felt by each electrode atoms due to all electrode atoms
   subroutine gradQelec_potential(localwork, ewald, box, &
         electrodes, xyz_atoms, gradQ_V)
      implicit none

      ! Parameters in
      ! -------------
      type(MW_localwork_t), intent(in) :: localwork  !< Localwork work distribution
      type(MW_ewald_t),    intent(in) :: ewald     !< Ewald summation parameters
      type(MW_box_t),      intent(in) :: box       !< Simulation box parameters

      type(MW_electrode_t), intent(in) :: electrodes(:)  !< electrode parameters
      real(wp),             intent(in) :: xyz_atoms(:,:)  !< Electrode atoms positions

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: gradQ_V(:,:)       !< Coulomb potential

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: i, j
      real(wp) :: volfactor
      real(wp) :: zij, pot_ij, zijsq
      real(wp) :: alpha, alphasq, alphaconst, sqrpialpha, alphapi
      integer :: num_block_diag, num_block_full
      integer :: iblock, iblock_offset, istart, iend, jstart, jend, itype, jtype
      integer :: count_n, offset_n, count_m, offset_m

      num_elec_types = size(electrodes,1)

      ! Precompute some factors
      ! -----------------------
      alpha = ewald%alpha
      alphasq = alpha * alpha
      alphaconst = -1.0_wp / (4.0_wp*alphasq)
      sqrpialpha = sqrt(pi) / alpha
      alphapi = 2.0_wp * alpha / sqrt(pi)

      ! Long range contribution from k=0 mode
      ! -------------------------------------
      ! Pair interactions are considered by block of pair_block_size x pair_block_size
      ! Diagonal block have a special treatment
      volfactor = 2.0_wp/box%area(3)

#ifdef _OPENACC
      ! On GPU it is faster to loop on the whole interaction matrix than only on the half
      ! because coelescent memory access are faster

      !$acc data present(xyz_atoms(:,3), gradQ_V(:,:))
      ! Blocks on the diagonal compute only half of the pair interactions
      istart = 1
      iend = size(gradQ_V,1)
      jstart = 1
      jend = size(gradQ_V,2)
      !$acc parallel loop private(zij, zijsq, pot_ij)
      do i = istart, iend
         !$acc loop 
         do j = jstart, jend
            zij = xyz_atoms(j,3) - xyz_atoms(i,3)
            zijsq = zij*zij
            pot_ij = volfactor * (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
            !$acc atomic update
            gradQ_V(j,i) = gradQ_V(j,i) - pot_ij
         end do
         !$acc end loop
      end do
      !$acc end parallel loop
      !$acc end data
#else
      ! Blocks on the diagonal compute only half of the pair interactions
      do itype = 1, num_elec_types
         count_n = electrodes(itype)%count
         offset_n = electrodes(itype)%offset
         do jtype = 1, itype-1
            count_m = electrodes(jtype)%count
            offset_m = electrodes(jtype)%offset

            ! Number of blocks with full interactions (below the diagonal
            num_block_full = localwork%pair_atom2atom_num_block_full_local(itype,jtype)
            iblock_offset = localwork%pair_atom2atom_full_iblock_offset(itype,jtype)
            do iblock = 1, num_block_full
               call update_other_block_boundaries(iblock_offset+iblock, &
                     offset_n, count_n, offset_m, count_m, &
                     istart, iend, jstart, jend)
               do i = istart, iend
                  do j = jstart, jend
                     zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                     zijsq = zij*zij
                     pot_ij = volfactor * (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                     gradQ_V(j,i) = gradQ_V(j,i) - pot_ij
                  end do
               end do
            end do
         end do

         num_block_diag = localwork%pair_atom2atom_num_block_diag_local(itype)
         iblock_offset = localwork%pair_atom2atom_diag_iblock_offset(itype)

         do iblock = 1, num_block_diag
            call update_diag_block_boundaries(iblock_offset+iblock, offset_n, count_n, istart, iend)
            do i = istart, iend
               do j = istart, iend
                  zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                  zijsq = zij*zij
                  pot_ij = volfactor * (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                  gradQ_V(j,i) = gradQ_V(j,i) - pot_ij
               end do
            end do
         end do

         ! Number of blocks with full interactions (below the diagonal
         num_block_full = localwork%pair_atom2atom_num_block_full_local(itype,itype)
         iblock_offset = localwork%pair_atom2atom_full_iblock_offset(itype,itype)
         do iblock = 1, num_block_full
            call update_tri_block_boundaries(iblock_offset+iblock, offset_n, count_n, &
                  istart, iend, jstart, jend)
            do i = istart, iend
               do j = jstart, jend
                  zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                  zijsq = zij*zij
                  pot_ij = volfactor * (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                  gradQ_V(j,i) = gradQ_V(j,i) - pot_ij
               end do
            end do
         end do
      end do
#endif
   end subroutine gradQelec_potential

   !================================================================================
   ! Compute the Coulomb potential felt by each electrode atoms due to all melt ions
   !DEC$ ATTRIBUTES NOINLINE :: melt2elec_potential
   subroutine qmelt2Qelec_potential(localwork, ewald, box, &
         ions, xyz_ions, electrodes, xyz_atoms, V)
      implicit none

      ! Parameters in
      ! -------------
      type(MW_localwork_t), intent(in) :: localwork !< Local work assignment
      type(MW_ewald_t), intent(in)    :: ewald    !< Ewald summation parameters
      type(MW_box_t),   intent(in)    :: box      !< Simulation box parameters

      type(MW_ion_t), intent(in) :: ions(:)        !< ions parameters
      real(wp),       intent(in) :: xyz_ions(:,:)  !< ions xyz positions

      type(MW_electrode_t), intent(in) :: electrodes(:)   !< electrode parameters
      real(wp),             intent(in) :: xyz_atoms(:,:)  !< electrode atoms xyz positions

      ! Parameters out
      ! --------------
      real(wp),                 intent(inout) :: V(:)       !< Coulomb potential

      ! Local
      ! -----
      integer :: num_ion_types, num_elec_types
      integer :: i, j, itype, jtype
      real(wp) :: qj !< ion charge
      real(wp) :: volfactor
      real(wp) :: zij, zijsq, pot_ij
      real(wp) :: alpha, alphasq, sqrpialpha
      integer :: num_block_full
      integer :: iblock, iblock_offset, istart, iend, jstart, jend
      integer :: count_n, offset_n, count_m, offset_m
      real(wp) :: vi

      num_ion_types = size(ions,1)
      num_elec_types = size(electrodes,1)

      ! Precompute some factors
      ! -----------------------
      alpha = ewald%alpha
      alphasq = alpha * alpha
      sqrpialpha = sqrt(pi) / alpha

      ! Long range contribution from k=0 mode
      ! -------------------------------------
      volfactor = 2.0_wp/box%area(3)

      !$acc data present(xyz_ions(:,3), xyz_atoms(:,3), V(:))
      do itype = 1, num_elec_types
         count_n = electrodes(itype)%count
         offset_n = electrodes(itype)%offset
         do jtype = 1, num_ion_types
            count_m = ions(jtype)%count
            offset_m = ions(jtype)%offset

            ! Number of blocks with full interactions (below the diagonal
            num_block_full = localwork%pair_ion2atom_num_block_full_local(itype,jtype)
            iblock_offset = localwork%pair_ion2atom_full_iblock_offset(itype,jtype)
            qj = ions(jtype)%charge
            do iblock = 1, num_block_full
               call update_other_block_boundaries(iblock_offset+iblock, &
                     offset_n, count_n, offset_m, count_m, &
                     istart, iend, jstart, jend)
               !$acc parallel loop private(zij, zijsq, pot_ij, vi)
               do i = istart, iend
                  vi = 0.0_wp
                  !$acc loop reduction(+:vi)
                  do j = jstart, jend
                     zij = xyz_ions(j,3) - xyz_atoms(i,3)
                     zijsq = zij*zij
                     pot_ij = (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                     vi = vi + volfactor * qj * pot_ij
                  end do
                  !$acc end loop
                  V(i) = V(i) - vi
               end do
               !$acc end parallel loop
            end do
         end do
      end do
      !$acc end data
   end subroutine qmelt2Qelec_potential

   !================================================================================
   ! Compute the Coulomb potential felt by each electrode atoms due to all electrode atoms
   !DEC$ ATTRIBUTES NOINLINE :: elec2elec_potential
   subroutine Qelec2Qelec_potential(localwork, ewald, box, &
         electrodes, xyz_atoms, q_elec, V)
      implicit none
      ! Parameters in
      ! -------------
      type(MW_localwork_t), intent(in) :: localwork  !< Localwork work distribution
      type(MW_ewald_t),    intent(in) :: ewald     !< Ewald summation parameters
      type(MW_box_t),      intent(in) :: box       !< Simulation box parameters

      type(MW_electrode_t), intent(in) :: electrodes(:)  !< electrode parameters
      real(wp),             intent(in) :: q_elec(:)      !< Electrode atoms charge
      real(wp),             intent(in) :: xyz_atoms(:,:)  !< Electrode atoms positions

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: V(:)       !< Coulomb potential

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: i, j
      real(wp) :: volfactor, vi_sum
      real(wp) :: zij, pot_ij, zijsq
      real(wp) :: alpha, alphasq, alphaconst, sqrpialpha, alphapi
      integer :: num_block_diag, num_block_full
      integer :: iblock, iblock_offset, istart, iend, jstart, jend, itype, jtype
      integer :: count_n, offset_n, count_m, offset_m
      num_elec_types = size(electrodes,1)

      ! Precompute some factors
      ! -----------------------
      alpha = ewald%alpha
      alphasq = alpha * alpha
      alphaconst = -1.0_wp / (4.0_wp*alphasq)
      sqrpialpha = sqrt(pi) / alpha
      alphapi = 2.0_wp * alpha / sqrt(pi)

      ! Long range contribution from k=0 mode
      ! -------------------------------------
      ! Pair interactions are considered by block of pair_block_size x pair_block_size
      ! Diagonal block have a special treatment
      volfactor = 2.0_wp/box%area(3)

      !$acc data present(xyz_atoms(:,3), q_elec(:), V(:))

      ! Blocks on the diagonal compute only half of the pair interactions
      do itype = 1, num_elec_types
         count_n = electrodes(itype)%count
         offset_n = electrodes(itype)%offset
         do jtype = 1, itype-1
            count_m = electrodes(jtype)%count
            offset_m = electrodes(jtype)%offset

            ! Number of blocks with full interactions (below the diagonal
            num_block_full = localwork%pair_atom2atom_num_block_full_local(itype,jtype)
            iblock_offset = localwork%pair_atom2atom_full_iblock_offset(itype,jtype)
            do iblock = 1, num_block_full
               call update_other_block_boundaries(iblock_offset+iblock, &
                     offset_n, count_n, offset_m, count_m, &
                     istart, iend, jstart, jend)
               !$acc parallel loop  private(zij, zijsq, pot_ij, vi_sum)
               do i = istart, iend
                  vi_sum = 0.0_wp
                  !$acc loop reduction(+:vi_sum)
                  do j = jstart, jend
                     zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                     zijsq = zij*zij
                     pot_ij = volfactor * (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                     !$acc atomic update
                     V(j) = V(j) - q_elec(i) * pot_ij
                     vi_sum = vi_sum + q_elec(j) * pot_ij
                  end do
                  !$acc end loop
                  V(i) = V(i) - vi_sum
               end do
               !$acc end parallel loop
            end do
         end do

         num_block_diag = localwork%pair_atom2atom_num_block_diag_local(itype)
         iblock_offset = localwork%pair_atom2atom_diag_iblock_offset(itype)

         do iblock = 1, num_block_diag
            call update_diag_block_boundaries(iblock_offset+iblock, offset_n, count_n, istart, iend)
            !$acc parallel loop  private(zij, zijsq, pot_ij, vi_sum)
            do i = istart, iend
               vi_sum = 0.0_wp
               !$acc loop reduction(+:vi_sum)
               do j = istart, iend
                  zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                  zijsq = zij*zij
                  pot_ij = volfactor * (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                  vi_sum = vi_sum +  q_elec(j) * pot_ij
               end do
               !$acc end loop
               V(i) = V(i) - vi_sum
            end do
            !$acc end loop
         end do

         ! Number of blocks with full interactions (below the diagonal
         num_block_full = localwork%pair_atom2atom_num_block_full_local(itype,itype)
         iblock_offset = localwork%pair_atom2atom_full_iblock_offset(itype,itype)
         do iblock = 1, num_block_full
            call update_tri_block_boundaries(iblock_offset+iblock, offset_n, count_n, &
                  istart, iend, jstart, jend)
            !$acc parallel loop  private(zij, zijsq, pot_ij, vi_sum)
            do i = istart, iend
               vi_sum = 0.0_wp
               !$acc loop reduction(+:vi_sum)
               do j = jstart, jend
                  zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                  zijsq = zij*zij
                  pot_ij = volfactor * (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                  !$acc atomic update
                  V(j) = V(j) - q_elec(i) * pot_ij
                  vi_sum = vi_sum +  q_elec(j) * pot_ij
               end do
               !$acc end loop
               V(i) = V(i) - vi_sum
            end do
            !$acc end parallel loop
         end do
      end do
      !$acc end data
   end subroutine Qelec2Qelec_potential

   !================================================================================
   ! Compute the Coulomb forces felt by each melt ions due to other melt ions
   !DEC$ ATTRIBUTES NOINLINE :: melt_forces
   subroutine melt_forces(localwork, ewald, box, ions, xyz_ions, &
         electrodes, xyz_atoms, q_elec, force, compute_force)
      implicit none

      ! Parameters in
      ! -------------
      type(MW_localwork_t),  intent(in) :: localwork       !< Local work assignment
      type(MW_ewald_t),     intent(in) :: ewald          !< Ewald summation parameters
      type(MW_box_t),       intent(in) :: box            !< Simulation box parameters
      type(MW_ion_t),       intent(in) :: ions(:)        !< ions parameters
      real(wp),             intent(in) :: xyz_ions(:,:)  !< ions xyz positions
      type(MW_electrode_t), intent(inout) :: electrodes(:)   !< electrode parameters
      real(wp),             intent(in) :: xyz_atoms(:,:)  !< electrode atoms xyz positions
      real(wp),             intent(in) :: q_elec(:)      !< Electrode atoms charge
      logical,              intent(in) :: compute_force !< compute forces on electrodes

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: force(:,:) !< Coulomb force on ions

      ! Local
      ! -----
      integer :: num_ion_types, num_elec_types
      integer :: i, j, itype, jtype
      real(wp) :: qi, qj !< ion charge
      real(wp) :: volfactor
      real(wp) :: zij, zijsq, erf_ij
      real(wp) :: fij, fiz, force_melt2ele, force_ele2ele
      real(wp) :: alpha, alphasq, alphapi, sqrpialpha
      integer :: num_block_diag, num_block_full
      integer :: iblock, iblock_offset, istart, iend, jstart, jend
      integer :: count_n, offset_n, count_m, offset_m

      num_ion_types = size(ions,1)
      num_elec_types = size(electrodes,1)

      ! Precompute some factors
      ! -----------------------
      alpha = ewald%alpha
      alphasq = alpha*alpha
      alphapi = 2.0_wp * alpha / sqrt(pi)
      sqrpialpha = sqrt(pi) / alpha

      ! Long range contribution from k=0 mode
      ! -------------------------------------
      ! Normalization
      volfactor = 2.0_wp/box%area(3)

      !$acc data present(xyz_ions(:,3), xyz_atoms(:,3), force(:,3), q_elec(:), &
      !$acc              electrodes, electrodes%force_ions)      

      do itype = 1, num_ion_types
         count_n = ions(itype)%count
         offset_n = ions(itype)%offset
         do jtype = 1, itype - 1
            count_m = ions(jtype)%count
            offset_m = ions(jtype)%offset
            qi = ions(itype)%charge
            qj = ions(jtype)%charge
            ! Number of blocks with full ion2ion interactions
            num_block_full = localwork%pair_ion2ion_num_block_full_local(itype, jtype)
            iblock_offset = localwork%pair_ion2ion_full_iblock_offset(itype, jtype)

            do iblock = 1, num_block_full
               call update_other_block_boundaries(iblock_offset+iblock, &
                     offset_n, count_n, offset_m, count_m, &
                     istart, iend, jstart, jend)
               !$acc parallel loop  private(zij, zijsq, erf_ij, fij, fiz)
               do i = istart, iend
                  fiz = 0.0_wp
                  !$acc loop reduction(+:fiz)
                  do j = jstart, jend
                     zij = xyz_ions(j,3) - xyz_ions(i,3)
                     zijsq = zij*zij
                     erf_ij = erf(alpha*zij)
                     fij = qi * qj * pi * volfactor * erf_ij
                     fiz = fiz + fij
                     !$acc atomic update
                     force(j,3) = force(j,3) + fij
                  end do
                  !$acc end loop
                  force(i,3) = force(i,3) - fiz
               end do
               !$acc end parallel loop
            end do
         end do

         ! Blocks on the diagonal compute only half of the pair interactions
         qi = ions(itype)%charge
         qj = ions(itype)%charge
         num_block_diag = localwork%pair_ion2ion_num_block_diag_local(itype)
         iblock_offset = localwork%pair_ion2ion_diag_iblock_offset(itype)

         do iblock = 1, num_block_diag
            call update_diag_block_boundaries(iblock_offset+iblock, &
                  offset_n, count_n, istart, iend)
            !$acc parallel loop private(fiz, zij, zijsq, erf_ij, fij, fiz)
            do i = istart, iend
               fiz = 0.0_wp
               !$acc loop reduction(+:fiz)
               do j = istart, iend
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  zijsq = zij*zij
                  erf_ij = erf(alpha*zij)
                  fij = qi * qj * pi * volfactor * erf_ij
                  fiz = fiz + fij
               end do
               !$acc end loop
               force(i,3) = force(i,3) - fiz
            end do
            !$acc end parallel loop
         end do

         ! Number of blocks with full interactions (below the diagonal
         qi = ions(itype)%charge
         qj = ions(itype)%charge
         num_block_full = localwork%pair_ion2ion_num_block_full_local(itype,itype)
         iblock_offset = localwork%pair_ion2ion_full_iblock_offset(itype,itype)
         do iblock = 1, num_block_full
            call update_tri_block_boundaries(iblock_offset+iblock, offset_n, count_n, &
                  istart, iend, jstart, jend)

            !$acc parallel loop private(fiz, zij, zijsq, erf_ij, fij, fiz)
            do i = istart, iend
               fiz = 0.0_wp
               !$acc loop reduction(+:fiz)
               do j = jstart, jend
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  zijsq = zij*zij
                  erf_ij = erf(alpha*zij)
                  fij = qi * qj * pi * volfactor * erf_ij
                  !$acc atomic update
                  force(j,3) = force(j,3) + fij
                  fiz = fiz + fij
               end do
               !$acc end loop
               force(i,3) = force(i,3) - fiz
            end do
            !$acc end parallel loop
         end do

         qi = ions(itype)%charge
         do jtype = 1, num_elec_types
            count_m = electrodes(jtype)%count
            offset_m = electrodes(jtype)%offset
            ! Number of blocks with full atom2ion interactions
            num_block_full = localwork%pair_atom2ion_num_block_full_local(itype, jtype)
            iblock_offset = localwork%pair_atom2ion_full_iblock_offset(itype, jtype)
            force_melt2ele = 0.0_wp
            do iblock = 1, num_block_full
               call update_other_block_boundaries(iblock_offset+iblock, &
                     offset_n, count_n, offset_m, count_m, &
                     istart, iend, jstart, jend)
               !$acc parallel loop private(zij, zijsq, erf_ij, fij, fiz)
               do i = istart, iend
                  fiz = 0.0_wp
                  !$acc loop reduction(+:fiz)
                  do j = jstart, jend
                     zij = xyz_atoms(j,3) - xyz_ions(i,3)
                     zijsq = zij*zij
                     erf_ij = erf(alpha*zij)
                     fij = qi * q_elec(j) * pi * volfactor * erf_ij
                     fiz = fiz + fij
                  end do
                  !$acc end loop
                  force(i,3) = force(i,3) - fiz
                  force_melt2ele = force_melt2ele + fiz
               end do
               !$acc end parallel loop
            end do
            electrodes(jtype)%force_ions(3,5) = electrodes(jtype)%force_ions(3,5) + force_melt2ele
         end do
      end do

      ! Elec 2 Elec interactions
      ! ------------------------
      if (compute_force) then
         ! Blocks on the diagonal compute only half of the pair interactions
         do itype = 1, num_elec_types
            count_n = electrodes(itype)%count
            offset_n = electrodes(itype)%offset
            do jtype = 1, itype-1
               count_m =  electrodes(jtype)%count
               offset_m = electrodes(jtype)%offset

               ! Number of blocks with full interactions (below the diagonal
               num_block_full = localwork%pair_atom2atom_num_block_full_local(itype,jtype)
               iblock_offset = localwork%pair_atom2atom_full_iblock_offset(itype,jtype)

               do iblock = 1, num_block_full
                  call update_other_block_boundaries(iblock_offset+iblock, &
                        offset_n, count_n, offset_m, count_m, istart, iend, jstart, jend)
                  !$acc parallel loop private(zij, zijsq, erf_ij, fij, force_ele2ele)
                  do i = istart, iend
                     force_ele2ele = 0.0_wp
                     !$acc loop reduction(+:force_ele2ele)
                     do j = jstart, jend
                        zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                        zijsq = zij*zij
                        erf_ij = erf(alpha*zij)
                        fij = q_elec(i) * q_elec(j) * pi * volfactor * erf_ij
                        !$acc atomic update
                        electrodes(jtype)%force_ions(3,4) = electrodes(jtype)%force_ions(3,4) + fij
                        force_ele2ele = force_ele2ele + fij
                     end do
                     !$acc end loop
                     electrodes(itype)%force_ions(3,4) = electrodes(itype)%force_ions(3,4) - force_ele2ele
                  end do
                  !$acc end parallel loop
               end do
            end do

            num_block_diag = localwork%pair_atom2atom_num_block_diag_local(itype)
            iblock_offset = localwork%pair_atom2atom_diag_iblock_offset(itype)

            do iblock = 1, num_block_diag
               call update_diag_block_boundaries(iblock+iblock_offset, &
                     offset_n, count_n, istart, iend)
               !$acc parallel loop private(fiz, zij, zijsq, erf_ij, fij, force_ele2ele)
               do i = istart, iend
                  force_ele2ele = 0.0_wp
                  !$acc loop reduction(+:force_ele2ele)
                  do j = istart, iend
                     zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                     zijsq = zij*zij
                     erf_ij = erf(alpha*zij)
                     fij = q_elec(i) * q_elec(j) * pi * volfactor * erf_ij
                     force_ele2ele = force_ele2ele + fij
                  end do
                  !$acc end loop
                  electrodes(itype)%force_ions(3,4) = electrodes(itype)%force_ions(3,4) - force_ele2ele
               end do
               !$acc end parallel loop
            end do

            ! Number of blocks with full interactions (below the diagonal
            num_block_full = localwork%pair_atom2atom_num_block_full_local(itype,itype)
            iblock_offset = localwork%pair_atom2atom_full_iblock_offset(itype,itype)

            do iblock = 1, num_block_full
               call update_tri_block_boundaries(iblock_offset+iblock, offset_n, count_n, &
                     istart, iend, jstart, jend)
               !$acc parallel loop private(fiz, zij, zijsq, erf_ij, fij, force_ele2ele)
               do i = istart, iend
                  force_ele2ele = 0.0_wp
                  !$acc loop reduction(+:force_ele2ele)
                  do j = jstart, jend
                     zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                     zijsq = zij*zij
                     erf_ij = erf(alpha*zij)
                     fij = q_elec(i) * q_elec(j) * pi * volfactor * erf_ij
                     !$acc atomic update
                     electrodes(itype)%force_ions(3,4) = electrodes(itype)%force_ions(3,4) + fij
                     force_ele2ele = force_ele2ele + fij
                  end do
                  !$acc end loop
                  electrodes(itype)%force_ions(3,4) = electrodes(itype)%force_ions(3,4) - force_ele2ele
               end do
               !$acc end parallel loop
            end do
         end do
      end if

      !$acc end data
   end subroutine melt_forces

   !================================================================================
   ! Compute the contibution to the energy from long-range Coulomb interaction
   !DEC$ ATTRIBUTES NOINLINE :: energy
   subroutine energy(localwork, ewald, box, ions, xyz_ions, &
         electrodes, xyz_atoms, q_elec, h)
      implicit none

      ! Parameters in
      ! -------------
      type(MW_localwork_t),  intent(in) :: localwork       !< Local work assignment
      type(MW_ewald_t),     intent(in) :: ewald          !< Ewald summation parameters
      type(MW_box_t),       intent(in) :: box            !< Simulation box parameters
      type(MW_ion_t),       intent(in) :: ions(:)        !< ions parameters
      real(wp),             intent(in) :: xyz_ions(:,:)  !< ions xyz positions
      type(MW_electrode_t), intent(in) :: electrodes(:)   !< electrode parameters
      real(wp),             intent(in) :: xyz_atoms(:,:)  !< electrode atoms xyz positions
      real(wp),             intent(in) :: q_elec(:)      !< Electrode atoms charge

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: h !< Energy

      ! Local
      ! -----
      integer :: num_ion_types, num_elec_types
      integer :: i, j, itype, jtype
      real(wp) :: qi, qj !< ion charge
      real(wp) :: volfactor
      real(wp) :: zij, zijsq
      real(wp) :: pot_ij, vsumzk0
      real(wp) :: alpha, alphasq, alphapi, sqrpialpha
      integer :: num_block_diag, num_block_full
      integer :: iblock, iblock_offset, istart, iend, jstart, jend
      integer :: count_n, offset_n, count_m, offset_m

      num_ion_types = size(ions,1)
      num_elec_types = size(electrodes,1)

      ! Precompute some factors
      ! -----------------------
      alpha = ewald%alpha
      alphasq = alpha*alpha
      alphapi = 2.0_wp * alpha / sqrt(pi)
      sqrpialpha = sqrt(pi) / alpha

      ! Zero energy
      ! ----------------
      h = 0.0_wp
      vsumzk0 = 0.0_wp

      ! Long range contribution from k=0 mode
      ! -------------------------------------
      ! Normalization
      volfactor = 2.0_wp/box%area(3)

      !$acc data present(xyz_ions(:,3), xyz_atoms(:,3), q_elec(:)) copy(vsumzk0)

      ! Melt -> Melt and Elec -> Melt interactions
      ! ------------------------------------------
      do itype = 1, num_ion_types
         count_n = ions(itype)%count
         offset_n = ions(itype)%offset
         do jtype = 1, itype - 1
            count_m = ions(jtype)%count
            offset_m = ions(jtype)%offset

            qi = ions(itype)%charge
            qj = ions(jtype)%charge
            ! Number of blocks with full ion2ion interactions
            num_block_full = localwork%pair_ion2ion_num_block_full_local(itype, jtype)
            iblock_offset = localwork%pair_ion2ion_full_iblock_offset(itype, jtype)

            do iblock = 1, num_block_full
               call update_other_block_boundaries(iblock+iblock_offset, &
                     offset_n, count_n, offset_m, count_m, &
                     istart, iend, jstart, jend)

               !$acc parallel loop collapse(2) reduction(+:vsumzk0) private(zij, zijsq, pot_ij)
               do i = istart, iend
                  do j = jstart, jend
                     zij = xyz_ions(j,3) - xyz_ions(i,3)
                     zijsq = zij*zij
                     pot_ij = (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                     vsumzk0 = vsumzk0 + qi * qj * pot_ij
                  end do
               end do
               !$acc end parallel loop
            end do
         end do

         ! Blocks on the diagonal compute only half of the pair interactions
         qi = ions(itype)%charge
         qj = ions(itype)%charge
         num_block_diag = localwork%pair_ion2ion_num_block_diag_local(itype)
         iblock_offset = localwork%pair_ion2ion_diag_iblock_offset(itype)

         do iblock = 1, num_block_diag
            call update_diag_block_boundaries(iblock_offset+iblock, offset_n, count_n, istart, iend)
            !$acc parallel loop collapse(2) reduction(+:vsumzk0) private(zij, zijsq, pot_ij)
            do i = istart, iend
               do j = istart, iend
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  zijsq = zij*zij
                  pot_ij = (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                  vsumzk0 = vsumzk0 + 0.5_wp * qi * qj * pot_ij
               end do
            end do
            !$acc end parallel loop
         end do

         ! Number of blocks with full interactions (below the diagonal
         qi = ions(itype)%charge
         qj = ions(itype)%charge
         num_block_full = localwork%pair_ion2ion_num_block_full_local(itype,itype)
         iblock_offset = localwork%pair_ion2ion_full_iblock_offset(itype,itype)

         do iblock = 1, num_block_full
            call update_tri_block_boundaries(iblock_offset+iblock, offset_n, count_n, &
                  istart, iend, jstart, jend)
            !$acc parallel loop collapse(2) reduction(+:vsumzk0) private(zij, zijsq, pot_ij)
            do i = istart, iend
               do j = jstart, jend
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  zijsq = zij*zij
                  pot_ij = (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                  vsumzk0 = vsumzk0 + qi * qj * pot_ij
               end do
            end do
            !$acc end parallel loop
         end do

         qi = ions(itype)%charge
         do jtype = 1, num_elec_types
            count_m = electrodes(jtype)%count
            offset_m = electrodes(jtype)%offset
            ! Number of blocks with full ion2ion interactions
            num_block_full = localwork%pair_atom2ion_num_block_full_local(itype, jtype)
            iblock_offset = localwork%pair_atom2ion_full_iblock_offset(itype, jtype)

            do iblock = 1, num_block_full
               call update_other_block_boundaries(iblock+iblock_offset, &
                     offset_n, count_n, offset_m, count_m, &
                     istart, iend, jstart, jend)
               !$acc parallel loop collapse(2) reduction(+:vsumzk0) private(zij, zijsq, pot_ij)
               do i = istart, iend
                  do j = jstart, jend
                     zij = xyz_atoms(j,3) - xyz_ions(i,3)
                     zijsq = zij*zij
                     pot_ij = (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                     vsumzk0 = vsumzk0 + qi * q_elec(j) * pot_ij
                  end do
               end do
               !$acc end parallel loop
            end do
         end do
      end do

      ! Elec 2 Elec interactions
      ! ------------------------
      ! Blocks on the diagonal compute only half of the pair interactions
      do itype = 1, num_elec_types
         count_n = electrodes(itype)%count
         offset_n = electrodes(itype)%offset
         do jtype = 1, itype-1
            count_m =  electrodes(jtype)%count
            offset_m = electrodes(jtype)%offset

            ! Number of blocks with full interactions (below the diagonal
            num_block_full = localwork%pair_atom2atom_num_block_full_local(itype,jtype)
            iblock_offset = localwork%pair_atom2atom_full_iblock_offset(itype,jtype)

            do iblock = 1, num_block_full
               call update_other_block_boundaries(iblock_offset+iblock, &
                     offset_n, count_n, offset_m, count_m, istart, iend, jstart, jend)
               !$acc parallel loop collapse(2) reduction(+:vsumzk0) private(zij, zijsq, pot_ij)
               do i = istart, iend
                  do j = jstart, jend
                     zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                     zijsq = zij*zij
                     pot_ij = (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                     vsumzk0 = vsumzk0 + q_elec(i) * q_elec(j) * pot_ij
                  end do
               end do
               !$acc end parallel loop
            end do
         end do

         num_block_diag = localwork%pair_atom2atom_num_block_diag_local(itype)
         iblock_offset = localwork%pair_atom2atom_diag_iblock_offset(itype)

         do iblock = 1, num_block_diag
            call update_diag_block_boundaries(iblock+iblock_offset, &
                  offset_n, count_n, istart, iend)
            !$acc parallel loop collapse(2) reduction(+:vsumzk0) private(zij, zijsq, pot_ij)
            do i = istart, iend
               do j = istart, iend
                  zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                  zijsq = zij*zij
                  pot_ij = (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                  vsumzk0 = vsumzk0 + 0.5_wp * q_elec(i) * q_elec(j) * pot_ij
               end do
            end do
            !$acc end parallel loop
         end do

         ! Number of blocks with full interactions (below the diagonal
         num_block_full = localwork%pair_atom2atom_num_block_full_local(itype,itype)
         iblock_offset = localwork%pair_atom2atom_full_iblock_offset(itype,itype)

         do iblock = 1, num_block_full

            call update_tri_block_boundaries(iblock_offset+iblock, offset_n, count_n, &
                  istart, iend, jstart, jend)
            !$acc parallel loop collapse(2) reduction(+:vsumzk0) private(zij, zijsq, pot_ij)
            do i = istart, iend
               do j = jstart, jend
                  zij = xyz_atoms(j,3) - xyz_atoms(i,3)
                  zijsq = zij*zij
                  pot_ij = (sqrpialpha*exp(-zijsq*alphasq) + pi*zij*erf(zij*alpha))
                  vsumzk0 = vsumzk0 + q_elec(i) * q_elec(j) * pot_ij
               end do
            end do
            !$acc end parallel loop
         end do
      end do
      !$acc end data

      ! scale by volume factor
      h = -vsumzk0 * volfactor

   end subroutine energy

   include 'update_diag_block_boundaries.inc'
   include 'update_tri_block_boundaries.inc'
   include 'update_other_block_boundaries.inc'

end module MW_coulomb_keq0
