! Module to compute short range electrostatic potential
module MW_coulomb_sr
   use MW_kinds, only: wp
   use MW_errors, only: MW_errors_allocate_error => allocate_error
   use MW_ewald, only: MW_ewald_t
   use MW_box, only: MW_box_t
   use MW_damping, only: MW_damping_t
   use MW_electrode, only: MW_electrode_t
   use MW_ion, only: MW_ion_t
   use MW_localwork, only: MW_localwork_t
   use MW_constants, only: sqrtpi
   use MW_multipole_tensors, only: MW_multipole_tensors_tij_ab => tij_ab 
   implicit none
   private

   ! Public subroutines
   ! ------------------
   public :: gradQelec_potential !grad_Q [grad_Q [U^sr_QQ]]
   public :: qmelt2Qelec_potential !grad_Q [U^sr_qQ]
   public :: Qelec2Qelec_potential !grad_Q [U^sr_QQ]
   public :: melt_efg !-grad_r grad_r [(U^sr_qq)] 
   public :: melt_forces !-grad_r [(U^sr_qq) + (U^sr_QQ) + (U^sr_qQ)]
   public :: energy !(U^sr_qq) + (U^sr_QQ) + (U^sr_qQ)
   public :: setup_sr_interaction

contains

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

      ! Parameters in
      ! -------------
      integer,              intent(in)    :: num_PBC !< number of periodic boundary conditions
      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

      select case (num_PBC)
      case(2)
         call gradQelec_potential_2DPBC(localwork, ewald, box, &
         electrodes, xyz_atoms, gradQ_V)
      case(3)
         call gradQelec_potential_3DPBC(localwork, ewald, box, &
         electrodes, xyz_atoms, gradQ_V)
      end select
   end subroutine gradQelec_potential

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

      ! Parameters in
      ! -------------
      integer,              intent(in)    :: num_PBC !< Number of periodic boundary conditions
      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

      select case (num_PBC)
      case(2)
         call qmelt2Qelec_potential_2DPBC(localwork, ewald, box, &
         ions, xyz_ions, electrodes, xyz_atoms, V)
      case(3)
         call qmelt2Qelec_potential_3DPBC(localwork, ewald, box, &
         ions, xyz_ions, electrodes, xyz_atoms, V)
      end select
   end subroutine qmelt2Qelec_potential

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

      ! Parameters in
      ! -------------
      integer,              intent(in)    :: num_PBC !< number of periodic boundary conditions
      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

      select case (num_PBC)
      case(2)
         call Qelec2Qelec_potential_2DPBC(localwork, ewald, box, &
         electrodes, xyz_atoms, q_elec, V)
      case(3)
         call Qelec2Qelec_potential_3DPBC(localwork, ewald, box, &
         electrodes, xyz_atoms, q_elec, V)
      end select
   end subroutine Qelec2Qelec_potential

   !================================================================================
   ! Compute the Coulomb forces felt by each melt ions due to other melt ions
   subroutine melt_forces(num_PBC, localwork, ewald, box, ions, xyz_ions, type_ions, &
                   electrodes, xyz_atoms, type_electrode_atoms, q_elec, force,stress_tensor, compute_force)
      implicit none

      ! Parameters in
      ! -------------
      integer,              intent(in)    :: num_PBC !< number of periodic boundary conditions
      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
      integer,              intent(in)    :: type_ions(:)    !< Ion types
      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
      integer,              intent(in)    :: type_electrode_atoms(:) !< Electrode atoms types      
      real(wp),             intent(in)    :: q_elec(:)      !< Electrode atoms charge
      logical,              intent(in)    :: compute_force

      ! Parameters out
      ! --------------
      real(wp),             intent(inout) :: force(:,:) !< Coulomb force on ions
      real(wp),             intent(inout) :: stress_tensor(:,:) !< Coulomb force contribution to the stress tensor

      select case (num_PBC)
      case(2)
              call melt_forces_2DPBC(localwork, ewald, box, ions, xyz_ions, type_ions, &
                      electrodes, xyz_atoms, type_electrode_atoms, q_elec, force,stress_tensor,compute_force)
      case(3)
              call melt_forces_3DPBC(localwork, ewald, box, ions, xyz_ions, type_ions, &
                      electrodes, xyz_atoms, type_electrode_atoms, q_elec, force,stress_tensor,compute_force)
      end select

   end subroutine melt_forces

   !==========================================================================================
   ! Compute the electric field gradient felt by melt ions due to point charges of melt ions.
   subroutine melt_efg(num_PBC, localwork, ewald, box, ions, xyz_ions, efg)
      implicit none

      ! Parameters in
      ! -------------
      integer,              intent(in)    :: num_PBC !< number of periodic boundary conditions
      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) :: efg(:,:) !< the EFG on ions

      select case (num_PBC)
      case(2)
         call melt_efg_2DPBC(localwork, ewald, box, ions, xyz_ions, efg)
      case(3)
         call melt_efg_3DPBC(localwork, ewald, box, ions, xyz_ions, efg)
      end select

   end subroutine melt_efg

   !================================================================================
   ! Compute the Energy due to Coulomb interactions
   subroutine energy(num_PBC, localwork, ewald, box, ions, xyz_ions, &
      electrodes, xyz_atoms, q_elec, h)
      implicit none

      ! Parameters in
      ! -------------
      integer,              intent(in)    :: num_PBC !< number of periodic boundary conditions
      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 !< Coulomb energy

      select case (num_PBC)
      case(2)
         call energy_2DPBC(localwork, ewald, box, ions, xyz_ions, &
         electrodes, xyz_atoms, q_elec, h)
      case(3)
         call energy_3DPBC(localwork, ewald, box, ions, xyz_ions, &
         electrodes, xyz_atoms, q_elec, h)
      end select
   end subroutine energy

   ! ==========================================================================================
   ! Determine if short-range coulomb interaction should be taken into account between 2 particle types
   !
   ! By default short-range coulomb interaction is taken into account
   ! For fixed particles, short-range coulomb interaction is not taken
   ! into account between 2 particle types if the minimum distance
   ! between 2 particles of the distinct sets is greater than the
   ! specified cut-off (rcut)
   subroutine setup_sr_interaction(num_PBC, ewald, ions, electrodes, &
      xyz_ions, xyz_atoms, box)
      implicit none
      ! Parameters
      ! ----------
      integer,              intent(in)    :: num_PBC !< number of periodic boundary conditions
      type(MW_ion_t),       intent(in)    :: ions(:)             !< ions parameters
      type(MW_electrode_t), intent(in)    :: electrodes(:) !< electrode parameters
      real(wp),             intent(in)    :: xyz_ions(:,:)             !< ions xyz coordinates
      real(wp),             intent(in)    :: xyz_atoms(:,:)            !< atoms xyz coordinates
      type(MW_box_t),       intent(in)    :: box                 !< box parameters
      type(MW_ewald_t),     intent(inout) :: ewald           !< Coulomb ewald structure

      select case(num_PBC)
      case(2)
         call setup_sr_interaction_2DPBC(ewald, ions, electrodes, &
         xyz_ions, xyz_atoms, box)
      case(3)
         call setup_sr_interaction_3DPBC(ewald, ions, electrodes, &
         xyz_ions, xyz_atoms, box)
      end select
   end subroutine setup_sr_interaction

   !Functions=============================

   function force_elec2melt_ij(drnorm2, alpha, alphasq, &
      alphapi, eta, etasq, etapi, qi, Q_j) &
      result(fijqq)
      implicit none
      !$acc routine seq      

      real(wp), intent(in) :: drnorm2
      real(wp), intent(in) :: alpha
      real(wp), intent(in) :: alphasq
      real(wp), intent(in) :: alphapi
      real(wp), intent(in) :: eta
      real(wp), intent(in) :: etasq
      real(wp), intent(in) :: etapi
      real(wp), intent(in) :: qi
      real(wp), intent(in) :: Q_j

      ! Result
      ! ------
      real(wp) :: fijqq

      ! Local
      ! -----
      real(wp) :: drnorm, drnormrec, drnorm2rec
      real(wp) :: erfc_alpha,exp_alpha
      real(wp) :: erfc_eta,exp_eta

      drnorm = sqrt(drnorm2)
      drnormrec = 1.0_wp/drnorm
      drnorm2rec = drnormrec*drnormrec

      ! Change erfc to 1-erf because of underflows
      erfc_alpha = erfc(alpha*drnorm)*drnormrec
      exp_alpha = exp(-alphasq*drnorm2)*alphapi
      erfc_eta = erfc(eta*drnorm)*drnormrec
      exp_eta = exp(-etasq*drnorm2)*etapi

      fijqq = qi*Q_j*drnorm2rec*((erfc_alpha + exp_alpha) - (erfc_eta + exp_eta))

   end function force_elec2melt_ij

   ! ================================================================================
   include 'coulomb_sr_2DPBC.inc'
   include 'coulomb_sr_3DPBC.inc'
   ! ================================================================================
   include 'minimum_image_displacement_2DPBC.inc'
   include 'minimum_image_distance_2DPBC.inc'
   include 'minimum_image_displacement_3DPBC.inc'
   include 'minimum_image_distance_3DPBC.inc'
   ! ================================================================================
   include 'update_diag_block_boundaries.inc'
   include 'update_tri_block_boundaries.inc'
   include 'update_other_block_boundaries.inc'

end module MW_coulomb_sr
