! Module to compute electrostatic Coulomb potential
module MW_coulomb_keq0_pim
   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 :: gradQelec_electricfield !grad_Q [grad_mu [U^k=0_muQ]]
   public :: diag_gradmumelt_electricfield !diag[grad_mu [grad_mu [U^k=0_mumu]]]
   public :: gradmumelt_electricfield !grad_mu [grad_mu [U^k=0_mumu]]
   public :: mumelt2Qelec_potential !grad_Q [U^k=0_muQ]
   public :: qmelt2mumelt_electricfield !-grad_mu [U^k=0_qmu]
   public :: mumelt2mumelt_electricfield !-grad_mu [U^k=0_mumu]
   public :: Qelec2mumelt_electricfield !-grad_mu [U^k=0_muQ]
   public :: melt_forces !-grad_r [(U^k=0_qq + U^k=0_qmu + U^k=0_mumu) + (U^k=0_QQ) + (U^k=0_qQ + U^k=0_muQ)]
   public :: energy !(U^k=0_qq + U^k=0_qmu + U^k=0_mumu) + (U^k=0_QQ) + (U^k=0_qQ + U^k=0_muQ)

contains

   subroutine gradQelec_electricfield(localwork, ewald, box, ions, xyz_ions, &
      electrodes, xyz_atoms, gradQ_efield)
      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) :: gradQ_efield(:,:) !< Electric field

      ! Local
      ! -----
      integer :: num_ions, num_ion_types, num_elec_types
      integer :: i, j, jz, itype, jtype
      real(wp) :: volfactor
      real(wp) :: zij, erf_ij
      real(wp) :: alpha
      integer :: num_block_full
      integer :: iblock, iblock_offset, istart, iend, jstart, jend
      integer :: count_n, offset_n, count_m, offset_m

      num_ions = size(gradQ_efield,1)/3
      num_ion_types = size(ions,1)
      num_elec_types = size(electrodes,1)

      ! Precompute some factors
      ! -----------------------
      alpha = ewald%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), gradQ_efield(:,:))

      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 atom2ion interactions
            num_block_full = localwork%pair_ion2atom_num_block_full_local(itype, jtype)
            iblock_offset = localwork%pair_ion2atom_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) private(jz, zij, erf_ij)
               do i = istart, iend
                  do j = jstart, jend
                     jz = j + num_ions*2
                     zij = xyz_atoms(i,3) - xyz_ions(j,3)
                     erf_ij = erf(alpha*zij)
                     !$acc atomic update
                     gradQ_efield(jz,i) = gradQ_efield(jz,i) + volfactor*pi*erf_ij
                  end do
               end do
               !$acc end parallel loop
            end do
         end do
      end do
      !$acc end data

   end subroutine gradQelec_electricfield

   !================================================================================
   ! Compute the electric fields felt by each melt ions due to other melt ions
   ! dipoles
   subroutine diag_gradmumelt_electricfield(localwork, ewald, box, ions, &
      diag_gradmu_efield)
      implicit none

      ! Parameters Input
      ! -------------
      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

      ! Parameters Input/Output
      ! --------------
      real(wp), intent(inout) :: diag_gradmu_efield(:)      !< matrix containing the
      !< gradient of the field
      ! Local
      ! -----
      integer :: num_ion_types, num_ions
      integer :: itype, i, iion, iionz
      real(wp) :: volfactor, alpha, sqrpi, elecz_ii

      num_ion_types = size(ions,1)
      num_ions = size(diag_gradmu_efield,1)/3

      ! Precompute some factors
      ! -----------------------
      alpha = ewald%alpha
      sqrpi = sqrt(pi)
      volfactor = 2.0_wp/box%area(3)
      elecz_ii = 2.0_wp*volfactor*alpha*sqrpi

      ! Long range contribution from k=0 mode
      ! -------------------------------------
      do itype = 1, num_ion_types
         do i = 1, localwork%count_ions(itype)
            iion = localwork%offset_ions(itype) + i
            iionz = iion + num_ions*2
            diag_gradmu_efield(iionz) = diag_gradmu_efield(iionz) + elecz_ii
         end do
      end do

   end subroutine diag_gradmumelt_electricfield

   !================================================================================
   ! Compute the electric fields felt by each melt ions due to other melt ions
   ! dipoles
   subroutine gradmumelt_electricfield(localwork, ewald, box, ions, xyz_ions, &
      gradmu_efield)
      implicit none

      ! Parameters Input
      ! -------------
      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

      ! Parameters Input/Output
      ! --------------
      real(wp), intent(inout) :: gradmu_efield(:,:)      !< matrix containing the
      !< gradient of the field
      ! Local
      ! -----
      integer :: num_ion_types, num_ions
      integer :: i, j, itype, jtype, iz, jz
      real(wp) :: volfactor
      real(wp) :: zij, zijsq, exp_ij
      real(wp) :: alpha, alphasq, alphapi, sqrpialpha,sqrpi
      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_ions = size(gradmu_efield,1)/3

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

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

      !$acc data present(xyz_ions(:,3), gradmu_efield(:,:))

      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
            ! 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 collapse(2) private(iz, jz, zij, zijsq, exp_ij)
               do i = istart, iend
                  do j = jstart, jend
                     iz = i + num_ions*2
                     jz = j + num_ions*2
                     zij = xyz_ions(j,3) - xyz_ions(i,3)
                     zijsq = zij*zij
                     exp_ij = exp(-zijsq*alphasq)
                     !$acc atomic update
                     gradmu_efield(jz,iz) = gradmu_efield(jz,iz) + &
                     2.0_wp*volfactor*alpha*sqrpi*exp_ij
                  end do
               end do
               !$acc end parallel loop
            end do
         end do

         ! Blocks on the diagonal compute only half of the pair interactions
         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) private(iz, jz, zij, zijsq, exp_ij)
            do i = istart, iend
               do j = istart, iend
                  iz = i + num_ions*2
                  jz = j + num_ions*2
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  zijsq = zij*zij
                  exp_ij = exp(-zijsq*alphasq)
                  !$acc atomic update
                  gradmu_efield(jz,iz) = gradmu_efield(jz,iz) + &
                  2.0_wp*volfactor*alpha*sqrpi*exp_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_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) private(iz, jz, zij, zijsq, exp_ij)
            do i = istart, iend
               do j = jstart, jend
                  iz = i + num_ions*2
                  jz = j + num_ions*2
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  zijsq = zij*zij
                  exp_ij = exp(-zijsq*alphasq)
                  !$acc atomic update
                  gradmu_efield(jz,iz) = gradmu_efield(jz,iz) + &
                  2.0_wp*volfactor*alpha*sqrpi*exp_ij
               end do
            end do
            !$acc end parallel loop
         end do

      end do
      !$acc end data
   end subroutine gradmumelt_electricfield

   !================================================================================
   ! Compute the Coulomb potential felt by each electrode atoms due to all melt ions
   subroutine mumelt2Qelec_potential(localwork, ewald, box, &
      ions, xyz_ions, electrodes, xyz_atoms, dipolez, 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
      real(wp),             intent(in) :: dipolez(:)      !< z-component of the dipole

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

      ! Local
      ! -----
      integer :: num_ion_types, num_elec_types
      integer :: num_ions, num_elec_atoms
      integer :: i, j 
      real(wp) :: volfactor
      real(wp) :: zij, zijsq, erf_ij
      real(wp) :: alpha
      real(wp) :: vi

      integer :: 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)
      num_ions = 0
      if (num_ion_types > 0) then
         num_ions = ions(num_ion_types)%offset+ions(num_ion_types)%count
      end if
      num_elec_atoms = 0
      if (num_elec_types > 0) then
         num_elec_atoms = electrodes(num_elec_types)%offset+electrodes(num_elec_types)%count
      end if

      ! Precompute some factors
      ! -----------------------
      alpha = ewald%alpha

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

      count_m = num_ions
      offset_m = 0
   
      count_n = num_elec_atoms
      offset_n = 0

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

      ! Number of blocks with full interactions (below the diagonal
      num_block_full = localwork%pair_ion2atom_num_block_full_local_ms
      iblock_offset = localwork%pair_ion2atom_full_iblock_offset_ms
      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, vi)
         do i = istart, iend
            vi = 0.0_wp
            !$acc loop reduction(+:vi)
            do j = jstart, jend
               zij = xyz_atoms(i,3) - xyz_ions(j,3)
               zijsq = zij*zij
               erf_ij = erf(zij*alpha)
               vi = vi + volfactor*pi*erf_ij*dipolez(j)
            end do
            !$acc end loop
            V(i) = V(i) + vi
         end do
         !$acc end parallel loop
      end do
      !$acc end data
   end subroutine mumelt2Qelec_potential

   !================================================================================
   ! Compute the electric fields felt by each melt ions due to other melt ions
   ! charges
   subroutine qmelt2mumelt_electricfield(localwork, ewald, box, ions, xyz_ions, &
      efield)
      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

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: efield(:,:) !< Electric field

      ! Local
      ! -----
      integer :: num_ion_types
      integer :: i, j, itype, jtype
      real(wp) :: qi, qj !< ion charge
      real(wp) :: volfactor
      real(wp) :: zij,  erf_ij
      real(wp) :: alpha
      real(wp) :: eiz

      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)

      ! Precompute some factors
      ! -----------------------
      alpha = ewald%alpha

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

      !$acc data present(xyz_ions(:,3), efield(:,:))

      do itype = 1, num_ion_types
         count_n = ions(itype)%count
         offset_n = ions(itype)%offset
         qi = ions(itype)%charge
         do jtype = 1, itype - 1
            count_m = ions(jtype)%count
            offset_m = ions(jtype)%offset
            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, erf_ij, eiz)
               do i = istart, iend
                  eiz = 0.0_wp
                  !$acc loop reduction(+:eiz)
                  do j = jstart, jend
                     zij = xyz_ions(j,3) - xyz_ions(i,3)
                     erf_ij = erf(alpha*zij)
                     eiz = eiz + volfactor*pi*erf_ij*qj
                     !$acc atomic update
                     efield(j,3) = efield(j,3) + volfactor*pi*erf_ij*qi
                  end do
                  !$acc end loop
                  efield(i,3) = efield(i,3) - eiz
               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) private(zij, erf_ij)
            do i = istart, iend
               do j = istart, iend
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  erf_ij = erf(alpha*zij)
                  !$acc atomic update
                  efield(j,3) = efield(j,3) + volfactor*pi*erf_ij*qi
               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 private(zij, erf_ij, eiz)
            do i = istart, iend
               eiz = 0.0_wp
               !$acc loop reduction(+:eiz)
               do j = jstart, jend
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  erf_ij = erf(alpha*zij)
                  eiz = eiz + volfactor*pi*erf_ij*qj
                  !$acc atomic update
                  efield(j,3) = efield(j,3) + volfactor*pi*erf_ij*qi
               end do
               !$acc end loop
               efield(i,3) = efield(i,3) - eiz
            end do
            !$acc end parallel loop
         end do
      end do
      !$acc end data
   end subroutine qmelt2mumelt_electricfield

   !================================================================================
   ! Compute the electric fields felt by each melt ions due to other melt ions
   ! dipoles
   subroutine mumelt2mumelt_electricfield(localwork, ewald, box, ions, xyz_ions, &
      efield, dipolez)
      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
      real(wp),             intent(in) :: dipolez(:)      !< z-component of the dipole

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: efield(:,:) !< Electric field on ions

      ! Local
      ! -----
      integer :: num_ion_types
      integer :: i, j, itype
      real(wp) :: muiz,mujz
      real(wp) :: volfactor
      real(wp) :: zij, zijsq, exp_ij
      real(wp) :: eiz
      real(wp) :: alpha, alphasq, alphapi, sqrpialpha,sqrpi
      integer :: num_block_diag, num_block_full
      integer :: iblock, iblock_offset, istart, iend, jstart, jend
      integer :: count_n, offset_n
       
      integer :: num_ions

      num_ion_types = size(ions,1)
      num_ions = 0
      do itype = 1,num_ion_types
         num_ions = num_ions + ions(itype)%count
      enddo
      count_n = num_ions
      offset_n = 0

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

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

      !$acc data present(xyz_ions(:,3), dipolez(:), efield(:,:))

      ! Blocks on the diagonal compute only half of the pair interactions
      num_block_diag = localwork%pair_ion2ion_num_block_diag_local_ms
      iblock_offset = localwork%pair_ion2ion_diag_iblock_offset_ms

      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) private(muiz, mujz, zij, zijsq, exp_ij)
         do i = istart, iend
            do j = istart, iend
               muiz = dipolez(i)
               mujz = dipolez(j)
               zij = xyz_ions(j,3) - xyz_ions(i,3)
               zijsq = zij*zij
               exp_ij = exp(-zijsq*alphasq)
               !$acc atomic update
               efield(j,3) = efield(j,3) - 2.0_wp*volfactor*alpha*sqrpi*exp_ij*muiz
            end do
         end do
         !$acc end parallel loop
      end do

      ! Number of blocks with full interactions (below the diagonal
      num_block_full = localwork%pair_ion2ion_num_block_full_local_ms
      iblock_offset = localwork%pair_ion2ion_full_iblock_offset_ms
      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(muiz, mujz, zij, zijsq, exp_ij, eiz)
         do i = istart, iend
            muiz = dipolez(i)
            eiz = 0.0_wp
            !$acc loop reduction(+:eiz)
            do j = jstart, jend
               mujz = dipolez(j)
               zij = xyz_ions(j,3) - xyz_ions(i,3)
               zijsq = zij*zij
               exp_ij = exp(-zijsq*alphasq)
               eiz = eiz - 2.0_wp*volfactor*alpha*sqrpi*exp_ij*mujz
               !$acc atomic update
               efield(j,3) = efield(j,3) - 2.0_wp*volfactor*alpha*sqrpi*exp_ij*muiz
            end do
            !$acc end loop
            efield(i,3) = efield(i,3) + eiz
         end do
         !$acc end parallel loop
      end do
      !$acc end data
   end subroutine mumelt2mumelt_electricfield

   !================================================================================
   ! Compute the electric fields felt by each melt ions due to other melt ions
   ! charges
   subroutine Qelec2mumelt_electricfield(localwork, ewald, box, ions, xyz_ions, &
      electrodes, xyz_atoms, q_elec, efield)
      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) :: efield(:,:) !< Electric field

      ! Local
      ! -----
      integer :: num_ion_types, num_elec_types
      integer :: i, j
      real(wp) :: volfactor
      real(wp) :: zij, erf_ij, eiz
      real(wp) :: alpha
      integer :: num_block_full
      integer :: iblock, iblock_offset, istart, iend, jstart, jend
      integer :: count_n, offset_n, count_m, offset_m
      integer  :: num_ions, num_elec_atoms

      num_ion_types = size(ions,1)
      num_elec_types = size(electrodes,1)
      num_ions = 0
      if (num_ion_types > 0) then
         num_ions = ions(num_ion_types)%offset+ions(num_ion_types)%count
      end if
      num_elec_atoms = 0
      if (num_elec_types > 0) then
         num_elec_atoms = electrodes(num_elec_types)%offset+electrodes(num_elec_types)%count
      end if

      ! Precompute some factors
      ! -----------------------
      alpha = ewald%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), q_elec(:), efield(:,:))

      count_m = num_elec_atoms
      offset_m = 0
   
      count_n = num_ions
      offset_n = 0
   
      ! Number of blocks with full interactions (below the diagonal
      num_block_full = localwork%pair_ion2atom_num_block_full_local_ms
      iblock_offset = localwork%pair_ion2atom_full_iblock_offset_ms
      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, erf_ij, eiz)
         do i = istart, iend
            eiz = 0.0_wp
            !$acc loop reduction(+:eiz)
            do j = jstart, jend
               zij = xyz_atoms(j,3) - xyz_ions(i,3)
               erf_ij = erf(alpha*zij)
               eiz = eiz - volfactor*pi*erf_ij*q_elec(j)
            end do
            !$acc end loop
            !$acc atomic update
            efield(i,3) = efield(i,3) + eiz
         end do
         !$acc end parallel loop
      end do
      !$acc end data
   end subroutine Qelec2mumelt_electricfield

   !================================================================================
   ! Compute the Coulomb forces felt by each melt ions due to other melt ions
   subroutine melt_forces(localwork, ewald, box, ions, xyz_ions, &
      electrodes, xyz_atoms, q_elec, force, dipolez)
      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
      real(wp),             intent(in) :: dipolez(:)      !< z-component of the dipole

      ! 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) :: Q_j    !< electrode charge
      real(wp) :: muiz, mujz
      real(wp) :: volfactor
      real(wp) :: zij, zijsq, erf_ij,exp_ij
      real(wp) :: fij, fiz
      real(wp) :: alpha, alphasq, alphapi, sqrpialpha,sqrpi
      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
      sqrpi = sqrt(pi)
      alphapi = 2.0_wp * alpha / sqrpi
      sqrpialpha = sqrpi / 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), q_elec(:), dipolez(:), force(:,:))

      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(muiz, mujz, zij, zijsq, erf_ij, exp_ij, fij, fiz)
               do i = istart, iend
                  fiz = 0.0_wp
                  muiz = dipolez(i)
                  !$acc loop reduction(+:fiz)
                  do j = jstart, jend
                     mujz = dipolez(j)
                     zij = xyz_ions(j,3) - xyz_ions(i,3)
                     zijsq = zij*zij
                     erf_ij = erf(alpha*zij)
                     exp_ij = exp(-zijsq*alphasq)
                     fij = qi * qj * pi * volfactor * erf_ij   &
                     -2.0_wp*volfactor*exp_ij*sqrpi*alpha*(qj*muiz - qi*mujz)  &
                     +4.0_wp*volfactor*zij*alphasq*alpha*sqrpi*muiz*mujz*exp_ij
                     fiz = fiz + fij
                     !$acc atomic update
                     force(j,3) = force(j,3) + fij
                  end do
                  !$acc end loop
                  !$acc atomic update
                  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(muiz, mujz, zij, zijsq, erf_ij, exp_ij, fij, fiz)
            do i = istart, iend
               fiz = 0.0_wp
               muiz = dipolez(i)
               !$acc loop reduction(+:fiz)
               do j = istart, iend
                  mujz = dipolez(j)
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  zijsq = zij*zij
                  erf_ij = erf(alpha*zij)
                  exp_ij = exp(-zijsq*alphasq)
                  fij = qi * qj * pi * volfactor * erf_ij  &
                  -2.0_wp*volfactor*exp_ij*sqrpi*alpha*(qj*muiz - qi*mujz) &
                  +4.0_wp*volfactor*zij*alphasq*alpha*sqrpi*muiz*mujz*exp_ij
                  fiz = fiz + fij
               end do
               !$acc end loop
               !$acc atomic update
               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(muiz, mujz, zij, zijsq, erf_ij, exp_ij, fij, fiz)
            do i = istart, iend
               fiz = 0.0_wp
               muiz = dipolez(i)
               !$acc loop reduction(+:fiz)
               do j = jstart, jend
                  mujz = dipolez(j)
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  zijsq = zij*zij
                  erf_ij = erf(alpha*zij)
                  exp_ij = exp(-zijsq*alphasq)
                  fij = qi * qj * pi * volfactor * erf_ij &
                        -2.0_wp*volfactor*exp_ij*sqrpi*alpha*(qj*muiz - qi*mujz) &
                        +4.0_wp*volfactor*zij*alphasq*alpha*sqrpi*muiz*mujz*exp_ij
                  fiz = fiz + fij
                  !$acc atomic update
                  force(j,3) = force(j,3) + fij
               end do
               !$acc end loop
               !$acc atomic update
               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)
            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(muiz, Q_j, zij, zijsq, erf_ij, exp_ij, fij, fiz)
               do i = istart, iend
                  fiz = 0.0_wp
                  muiz = dipolez(i)
                  !$acc loop reduction(+:fiz)
                  do j = jstart, jend
                     Q_j = q_elec(j)
                     zij = xyz_atoms(j,3) - xyz_ions(i,3)
                     zijsq = zij*zij
                     erf_ij = erf(alpha*zij)
                     exp_ij = exp(-zijsq*alphasq)
                     fij = volfactor*pi*erf_ij*qi*Q_j - 2.0_wp*volfactor*exp_ij*sqrpi*alpha*muiz*Q_j
                     fiz = fiz + fij
                  end do
                  !$acc end loop
                  !$acc atomic update
                  force(i,3) = force(i,3) - fiz
               end do
               !$acc end parallel loop
            end do
         end do
      end do
      !$acc end data
   end subroutine melt_forces

   !================================================================================
   ! Compute the contibution to the energy from long-range Coulomb interaction
   subroutine energy(localwork, ewald, box, ions, xyz_ions, &
      electrodes, xyz_atoms, q_elec, h, dipolez)
      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
      real(wp),             intent(in) :: dipolez(:)      !< z-component of the dipole

      ! 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) :: Q_j !< electrode charge
      real(wp) :: muiz, mujz
      real(wp) :: volfactor
      real(wp) :: zij, zijsq
      real(wp) :: erf_ij,exp_ij
      real(wp) :: pot_ij, vsumzk0
      real(wp) :: alpha, alphasq, alphapi, sqrpialpha,sqrpi

      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
      sqrpi = sqrt(pi)
      alphapi = 2.0_wp*alpha/sqrpi
      sqrpialpha = sqrpi/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), dipolez(:), 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
         qi = ions(itype)%charge
         do jtype = 1, itype - 1
            count_m = ions(jtype)%count
            offset_m = ions(jtype)%offset
            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(muiz, mujz, zij, zijsq, erf_ij, exp_ij, pot_ij)
               do i = istart, iend
                  do j = jstart, jend
                     muiz = dipolez(i)
                     mujz = dipolez(j)
                     zij = xyz_ions(j,3) - xyz_ions(i,3)
                     zijsq = zij*zij
                     erf_ij = erf(zij*alpha)
                     exp_ij = exp(-zijsq*alphasq)
                     pot_ij = (sqrpialpha*exp_ij + pi*zij*erf_ij)
                     ! The minus sign for qmu term arises from the different definition of zij
                     ! w.r.t. Pounds PhD (erf is an impair function)
                     vsumzk0 = vsumzk0 + qi * qj * pot_ij &
                     -pi*(qj*muiz - qi*mujz)*erf_ij   &
                     -2.0_wp*muiz*mujz*alpha*sqrpi*exp_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(muiz, mujz, zij, zijsq, erf_ij, exp_ij, pot_ij)
            do i = istart, iend
               do j = istart, iend
                  muiz = dipolez(i)
                  mujz = dipolez(j)
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  zijsq = zij*zij
                  erf_ij = erf(zij*alpha)
                  exp_ij = exp(-zijsq*alphasq)
                  pot_ij = (sqrpialpha*exp_ij + pi*zij*erf_ij)
                  ! The minus sign for qmu term arises from the different definition of zij
                  ! w.r.t. Pounds PhD (erf is an impair function)
                  vsumzk0 = vsumzk0 + 0.5_wp*(qi * qj * pot_ij &
                  -pi*(qj*muiz - qi*mujz)*erf_ij   &
                  -2.0_wp*muiz*mujz*alpha*sqrpi*exp_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(muiz, mujz, zij, zijsq, erf_ij, exp_ij, pot_ij)
            do i = istart, iend
               do j = jstart, jend
                  muiz = dipolez(i)
                  mujz = dipolez(j)
                  zij = xyz_ions(j,3) - xyz_ions(i,3)
                  zijsq = zij*zij
                  erf_ij = erf(zij*alpha)
                  exp_ij = exp(-zijsq*alphasq)
                  pot_ij = (sqrpialpha*exp_ij + pi*zij*erf_ij)
                  ! The minus sign for qmu term arises from the different definition of zij
                  ! w.r.t. Pounds PhD (erf is an impair function)
                  vsumzk0 = vsumzk0 + qi * qj * pot_ij &
                        -pi*(qj*muiz - qi*mujz)*erf_ij   &
                        -2.0_wp*muiz*mujz*alpha*sqrpi*exp_ij
               end do
            end do
            !$acc end parallel loop
         end do

         !qi -> point charge (melt)
         !muiz -> z-component of dipoles (melt)
         !Q_j -> integrated gaussian charge (walls)
         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(muiz, Q_j, zij, zijsq, erf_ij, exp_ij, pot_ij)
               do i = istart, iend
                  do j = jstart, jend
                     muiz = dipolez(i)
                     Q_j = q_elec(j)
                     zij = xyz_atoms(j,3) - xyz_ions(i,3)
                     zijsq = zij*zij
                     erf_ij = erf(zij*alpha)
                     exp_ij = exp(-zijsq*alphasq)
                     pot_ij = (sqrpialpha*exp_ij + pi*zij*erf_ij)
                     vsumzk0 = vsumzk0 + qi*Q_j*pot_ij &
                           -pi*Q_j*muiz*erf_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_pim
