! Module to compute electrostatic Coulomb potential
module MW_coulomb_self
   implicit none
   private

   ! Public subroutines
   ! ------------------
   !NO PIM
   public :: diag_gradQelec_potential
   public :: gradQelec_potential
   public :: Qelec2Qelec_potential
   public :: melt_efg ! correction to the efg
   public :: energy
   !PIM
   public :: diag_gradmumelt_electricfield
   public :: gradmumelt_electricfield
   public :: mumelt2mumelt_electricfield
   public :: pim_energy

contains

   !================================================================================
   ! Compute the Coulomb potential felt by each electrode atoms due to all electrode atoms
   subroutine diag_gradQelec_potential(localwork, ewald, electrodes, diag_gradQ_V)
      use MW_kinds, only: wp
      use MW_constants, only: pi
      use MW_ewald, only: MW_ewald_t
      use MW_box, only: MW_box_t
      use MW_electrode, only: MW_electrode_t
      use MW_ion, only: MW_ion_t
      use MW_localwork, only: MW_localwork_t
      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_electrode_t), intent(in)    :: electrodes(:)  !< electrode parameters

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

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: itype, i, iatom
      real(wp) :: selffactor
      real(wp) :: etai
      real(wp) :: alpha

      num_elec_types = size(electrodes,1)

      !$acc data &
      !$acc present(electrodes(:), localwork, localwork%count_atoms(:), localwork%offset_atoms(:), diag_gradQ_V(:))

      ! Self interaction correction
      ! ---------------------------
      alpha = ewald%alpha
      !$acc parallel loop gang private(itype, etai, selffactor, i, iatom)
      do itype = 1, num_elec_types
         etai = electrodes(itype)%eta
         selffactor = (2.0_wp * alpha - sqrt(2.0_wp) * etai ) / sqrt(pi)
         !$acc loop vector
         do i = 1, localwork%count_atoms(itype)
            iatom = localwork%offset_atoms(itype) + i
            diag_gradQ_V(iatom) = diag_gradQ_V(iatom) - selffactor
         end do
         !$acc end loop
      end do
      !$acc end parallel loop
      !$acc end data
   end subroutine diag_gradQelec_potential

   !================================================================================
   ! Compute the Coulomb potential felt by each electrode atoms due to all electrode atoms
   subroutine gradQelec_potential(localwork, ewald, electrodes, gradQ_V)
      use MW_kinds, only: wp
      use MW_constants, only: pi
      use MW_ewald, only: MW_ewald_t
      use MW_box, only: MW_box_t
      use MW_electrode, only: MW_electrode_t
      use MW_ion, only: MW_ion_t
      use MW_localwork, only: MW_localwork_t
      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_electrode_t), intent(in)    :: electrodes(:)  !< electrode parameters

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

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: itype, i, iatom
      real(wp) :: selffactor
      real(wp) :: etai
      real(wp) :: alpha

      num_elec_types = size(electrodes,1)


      ! Self interaction correction
      ! ---------------------------
      alpha = ewald%alpha
      do itype = 1, num_elec_types
         etai = electrodes(itype)%eta
         selffactor = (2.0_wp * alpha - sqrt(2.0_wp) * etai ) / sqrt(pi)
         do i = 1, localwork%count_atoms(itype)
            iatom = localwork%offset_atoms(itype) + i
            gradQ_V(iatom,iatom) = gradQ_V(iatom,iatom) - selffactor
         end do
      end do
   end subroutine gradQelec_potential

   !================================================================================
   ! Compute the Coulomb potential felt by each electrode atoms due to all electrode atoms
   subroutine Qelec2Qelec_potential(localwork, ewald, electrodes, q_elec, V)
      use MW_kinds, only: wp
      use MW_constants, only: pi
      use MW_ewald, only: MW_ewald_t
      use MW_box, only: MW_box_t
      use MW_electrode, only: MW_electrode_t
      use MW_ion, only: MW_ion_t
      use MW_localwork, only: MW_localwork_t
      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_electrode_t), intent(in)    :: electrodes(:)  !< electrode parameters
      real(wp),             intent(in)    :: q_elec(:)      !< Electrode atoms charge

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

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: itype, i, iatom
      real(wp) :: selffactor
      real(wp) :: etai
      real(wp) :: alpha

      num_elec_types = size(electrodes,1)


      ! Self interaction correction
      ! ---------------------------
      alpha = ewald%alpha
      do itype = 1, num_elec_types
         etai = electrodes(itype)%eta
         selffactor = (2.0_wp*alpha - sqrt(2.0_wp)*etai )/sqrt(pi)
         do i = 1, localwork%count_atoms(itype)
            iatom = localwork%offset_atoms(itype) + i
            V(iatom) = V(iatom) - selffactor*q_elec(iatom)
         end do
      end do
   end subroutine Qelec2Qelec_potential

   !================================================================================
   ! Compute the contibution to the energy from long-range Coulomb interaction
   subroutine energy(localwork, ewald, ions, electrodes, q_elec, h)
      use MW_kinds, only: wp
      use MW_constants, only: sqrtpi, sqrt2
      use MW_ewald, only: MW_ewald_t
      use MW_box, only: MW_box_t
      use MW_electrode, only: MW_electrode_t
      use MW_ion, only: MW_ion_t
      use MW_localwork, only: MW_localwork_t
      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_ion_t),       intent(in)    :: ions(:) !< ion types parameters
      type(MW_electrode_t), intent(in)    :: electrodes(:)  !< electrode parameters
      real(wp),             intent(in)    :: q_elec(:)      !< Electrode atoms charge

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

      ! Local
      ! -----
      integer :: num_ion_types, num_elec_types
      integer :: itype, i, iatom
      real(wp) :: selffactor
      real(wp) :: etai, qi, vi
      real(wp) :: alpha
      real(wp) :: h_ions, h_atoms

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

      ! Self interaction correction
      ! ---------------------------
      h_ions = 0.0_wp
      alpha = ewald%alpha
      selffactor = alpha / sqrtpi
      do itype = 1, num_ion_types
         qi = ions(itype)%charge
         h_ions = h_ions + localwork%count_ions(itype)*selffactor*qi*qi
      end do

      h_atoms = 0.0_wp
      do itype = 1, num_elec_types
         etai = electrodes(itype)%eta
         vi = electrodes(itype)%V
         selffactor = (alpha - etai/sqrt2)/sqrtpi
         do i = 1, localwork%count_atoms(itype)
            iatom = localwork%offset_atoms(itype) + i
            h_atoms = h_atoms + selffactor*q_elec(iatom)*q_elec(iatom) + q_elec(iatom)*vi
         end do
      end do

      h = -(h_ions + h_atoms)

   end subroutine energy

   !================================================================================
   ! Compute the contibution to the gradient of the electric field from self term
   subroutine diag_gradmumelt_electricfield(localwork, ewald, ions, diag_gradmu_efield)
      use MW_kinds, only: wp
      use MW_constants, only: pi
      use MW_ewald, only: MW_ewald_t
      use MW_box, only: MW_box_t
      use MW_ion, only: MW_ion_t
      use MW_localwork, only: MW_localwork_t
      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_ion_t),       intent(in)    :: ions(:)   !< ion types parameters

      ! Parameters out
      ! --------------
      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, iionx, iiony, iionz
      real(wp) :: selffactor
      real(wp) :: alpha

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

      ! Self interaction correction
      ! ---------------------------
      alpha = ewald%alpha
      selffactor = 4.0_wp*alpha**3.0_wp/(3.0_wp*sqrt(pi))
      do itype = 1, num_ion_types
         do i = 1, localwork%count_ions(itype)
            iion = localwork%offset_ions(itype) + i
            iionx = iion
            iiony = iion + num_ions
            iionz = iion + num_ions*2
            diag_gradmu_efield(iionx) = diag_gradmu_efield(iionx) - selffactor
            diag_gradmu_efield(iiony) = diag_gradmu_efield(iiony) - selffactor
            diag_gradmu_efield(iionz) = diag_gradmu_efield(iionz) - selffactor
         end do
      end do

   end subroutine diag_gradmumelt_electricfield

   !================================================================================
   ! Compute the contibution to the gradient of the electric field from self term
   subroutine gradmumelt_electricfield(localwork, ewald, ions, gradmu_efield)
      use MW_kinds, only: wp
      use MW_constants, only: pi
      use MW_ewald, only: MW_ewald_t
      use MW_box, only: MW_box_t
      use MW_ion, only: MW_ion_t
      use MW_localwork, only: MW_localwork_t
      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_ion_t),       intent(in)    :: ions(:)   !< ion types parameters

      ! Parameters out
      ! --------------
      real(wp),             intent(inout) :: gradmu_efield(:,:)  !< matrix containing the
      !< gradient of the field

      ! Local
      ! -----
      integer :: num_ion_types, num_ions
      integer :: itype, i, iion, iionx, iiony, iionz
      real(wp) :: selffactor
      real(wp) :: alpha

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

      ! Self interaction correction
      ! ---------------------------
      alpha = ewald%alpha
      selffactor = 4.0_wp*alpha**3.0_wp/(3.0_wp*sqrt(pi))
      do itype = 1, num_ion_types
         do i = 1, localwork%count_ions(itype)
            iion = localwork%offset_ions(itype) + i
            iionx = iion
            iiony = iion + num_ions
            iionz = iion + num_ions*2
            gradmu_efield(iionx,iionx) = gradmu_efield(iionx,iionx) - selffactor
            gradmu_efield(iiony,iiony) = gradmu_efield(iiony,iiony) - selffactor
            gradmu_efield(iionz,iionz) = gradmu_efield(iionz,iionz) - selffactor
         end do
      end do

   end subroutine gradmumelt_electricfield

   !================================================================================
   ! Compute the contibution to the electric field from self term
   subroutine mumelt2mumelt_electricfield(localwork, ewald, ions, efield, dipoles)
      use MW_kinds, only: wp
      use MW_constants, only: pi
      use MW_ewald, only: MW_ewald_t
      use MW_box, only: MW_box_t
      use MW_ion, only: MW_ion_t
      use MW_localwork, only: MW_localwork_t
      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_ion_t),       intent(in)    :: ions(:) !< ion types parameters
      real(wp),             intent(in)    :: dipoles(:,:)      !< ions dipoles

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

      ! Local
      ! -----
      integer :: num_ion_types
      integer :: itype, i, iion
      real(wp) :: selffactor
      real(wp) :: alpha

      num_ion_types = size(ions,1)

      ! Self interaction correction
      ! ---------------------------
      alpha = ewald%alpha
      selffactor = 4.0_wp*alpha**3.0_wp/(3.0_wp*sqrt(pi))
      do itype = 1, num_ion_types
         do i= 1, localwork%count_ions(itype)
            iion = localwork%offset_ions(itype) + i
            efield(iion,1) = efield(iion,1) + selffactor*dipoles(iion,1)
            efield(iion,2) = efield(iion,2) + selffactor*dipoles(iion,2)
            efield(iion,3) = efield(iion,3) + selffactor*dipoles(iion,3)
         enddo
      end do

   end subroutine mumelt2mumelt_electricfield

   !================================================================================
   ! Compute the contibution to the electric field gradient from self term
   subroutine melt_efg(localwork, ewald, ions, efg)
      use MW_kinds, only: wp
      use MW_constants, only: pi
      use MW_ewald, only: MW_ewald_t
      use MW_box, only: MW_box_t
      use MW_ion, only: MW_ion_t
      use MW_localwork, only: MW_localwork_t
      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_ion_t),       intent(in)    :: ions(:)    !< ion types parameters
      ! Parameters out
      ! --------------
      real(wp),             intent(inout) :: efg(:,:)   !< The electric field gradient

      ! Local
      ! -----
      integer :: num_ion_types
      integer :: itype, i, iion
      real(wp) :: selffactor, qi
      real(wp) :: alpha

      num_ion_types = size(ions,1)

      ! Self interaction correction
      ! ---------------------------
      alpha = ewald%alpha
      selffactor = 4.0_wp*alpha**3.0_wp/(3.0_wp*sqrt(pi))
      do itype = 1, num_ion_types
         qi = ions(itype)%charge
         if (ions(itype)%dump_efg) then 
            do i = 1, localwork%count_ions(itype)
               iion = localwork%offset_ions(itype) + i
               ! Only the diagonal components have the correction
               ! efg: 1-xx, 2-xy, 3-xz, 4: yy, 5: yz, 6: zz
               efg(iion,1) = efg(iion,1) - selffactor*qi
               efg(iion,4) = efg(iion,4) - selffactor*qi
               efg(iion,6) = efg(iion,6) - selffactor*qi
            enddo
         end if 
      end do

   end subroutine melt_efg

   !================================================================================
   ! Compute the contibution to the energy from long-range Coulomb interaction
   ! charge-charge and dipole-dipole terms
   subroutine pim_energy(localwork, ewald, ions, electrodes, q_elec, dipoles, h)
      use MW_kinds, only: wp
      use MW_constants, only: sqrtpi, sqrt2
      use MW_ewald, only: MW_ewald_t
      use MW_box, only: MW_box_t
      use MW_electrode, only: MW_electrode_t
      use MW_ion, only: MW_ion_t
      use MW_localwork, only: MW_localwork_t
      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_ion_t),       intent(in)    :: ions(:) !< ion types parameters
      type(MW_electrode_t), intent(in)    :: electrodes(:)  !< electrode parameters
      real(wp),             intent(in)    :: q_elec(:)      !< Electrode atoms charge
      real(wp),             intent(in)    :: dipoles(:,:)      !< ions dipoles

      ! Parameters out
      ! --------------
      real(wp),             intent(inout) :: h       !< Self-component of the energy

      ! Local
      ! -----
      integer :: num_ion_types, num_elec_types
      integer :: itype, i, iatom, iion
      real(wp) :: selffactor
      real(wp) :: etai, qi, vi
      real(wp) :: alpha
      real(wp) :: h_ions, h_atoms, h_dipoles

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

      ! Self interaction correction
      ! ---------------------------
      h_ions = 0.0_wp
      alpha = ewald%alpha
      selffactor = alpha / sqrtpi
      do itype = 1, num_ion_types
         qi = ions(itype)%charge
         h_ions = h_ions + localwork%count_ions(itype)*selffactor*qi*qi
      end do

      h_atoms = 0.0_wp
      do itype = 1, num_elec_types
         etai = electrodes(itype)%eta
         vi = electrodes(itype)%V
         selffactor = (alpha - etai/sqrt2)/sqrtpi
         do i = 1, localwork%count_atoms(itype)
            iatom = localwork%offset_atoms(itype) + i
            h_atoms = h_atoms + selffactor*q_elec(iatom)*q_elec(iatom) + q_elec(iatom)*vi
         end do
      end do

      h_dipoles = 0.0_wp
      selffactor = 2.0_wp*alpha**3.0_wp/(3.0_wp*sqrtpi)
      do itype = 1, num_ion_types
         do i = 1, localwork%count_ions(itype)
            iion = localwork%offset_ions(itype) + i
            h_dipoles = h_dipoles + selffactor*dipoles(iion,1)*dipoles(iion,1)
            h_dipoles = h_dipoles + selffactor*dipoles(iion,2)*dipoles(iion,2)
            h_dipoles = h_dipoles + selffactor*dipoles(iion,3)*dipoles(iion,3)
         enddo
      end do

      h = -(h_ions + h_atoms + h_dipoles)

   end subroutine pim_energy

end module MW_coulomb_self
