! Steele Potential
! -----------------------
!
module MW_steele
#ifndef MW_SERIAL
   use MPI
#endif
   use MW_kinds, only: wp
   use MW_box, only: MW_box_t
   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_timers
   use MW_localwork, only: MW_localwork_t
   use MW_parallel, only: MW_COMM_WORLD
   use MW_molecule, only: MW_molecule_t
   implicit none
   private

   ! Public type
   ! -----------
   public :: MW_steele_t

   ! Public subroutines
   ! ------------------
   public :: define_type
   public :: void_type
   public :: print_type
   public :: set_rcut
   public :: set_parameters
   public :: melt_forces
   public :: energy

   ! Data structure to handle Steele potentials
   type MW_steele_t
      integer :: num_walls = 0                       ! number of walls
      integer :: num_species = 0                     ! number of species
      real(wp) :: rcut = 0.0_wp                      ! cut-off parameter
      real(wp) :: rcutsq = 0.0_wp                    ! cut-off parameter squared
      real(wp), dimension(:), allocatable :: delta   ! (num_walls) 
      real(wp), dimension(:,:), allocatable :: A ! (num_walls,num_species)
      real(wp), dimension(:,:), allocatable :: B ! (num_walls,num_species)
    end type MW_steele_t

   integer, parameter :: ION_NAME_LEN = 8

contains

   !================================================================================
   ! Define the data structure
   subroutine define_type(this, num_walls, num_species, rcut)
      use MW_errors, only: &
            MW_errors_allocate_error => allocate_error, &
            MW_errors_parameter_error => parameter_error
      implicit none
      ! Parameters
      ! ---------_
      type(MW_steele_t), intent(inout) :: this      !> structure to be defined
      integer, intent(in) :: num_walls, num_species                    !> number of species
      real(wp), intent(in) :: rcut

      ! Local
      ! -----
      integer :: ierr

      call set_rcut(this, rcut)

      if (num_species <= 0) then
         call MW_errors_parameter_error("define_type", "steele.f90",&
               "num_species", num_species)
      end if

      if (num_walls <= 0) then
         call MW_errors_parameter_error("define_type", "steele.f90",&
               "num_walls", num_walls)
      end if

      this%num_species = num_species
      this%num_walls = num_walls

      allocate(this%delta(num_walls), stat=ierr)
      if (ierr /= 0) &
            call MW_errors_allocate_error("define_type", "steele.f90", ierr)
      this%delta(:) = 0.0_wp

      allocate(this%A(num_walls,num_species), stat=ierr)
      if (ierr /= 0) &
            call MW_errors_allocate_error("define_type", "steele.f90", ierr)
      this%A(:,:) = 0.0_wp

      allocate(this%B(num_walls,num_species), stat=ierr)
      if (ierr /= 0) &
            call MW_errors_allocate_error("define_type", "steele.f90", ierr)
      this%B(:,:) = 0.0_wp

   end subroutine define_type

   !================================================================================
   ! Void the data structure
   subroutine void_type(this)
      use MW_errors, only: MW_errors_deallocate_error => deallocate_error
      implicit none
      type(MW_steele_t), intent(inout) :: this

      integer :: ierr

      if (allocated(this%delta)) then
         deallocate(this%delta, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "steele.f90", ierr)
         end if
      end if

      if (allocated(this%A)) then
         deallocate(this%A, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "steele.f90", ierr)
         end if
      end if

      if (allocated(this%B)) then
         deallocate(this%B, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "steele.f90", ierr)
         end if
      end if

      this%num_species = 0
      this%num_walls = 0
      this%rcut = 0.0_wp
      this%rcutsq = 0.0_wp

   end subroutine void_type

   ! ================================================================================
   ! set rcut value
   subroutine set_rcut(this, rcut)
      implicit none
      type(MW_steele_t), intent(inout) :: this
      real(wp), intent(in) :: rcut

      this%rcut = rcut
      this%rcutsq = rcut*rcut

   end subroutine set_rcut

   ! ================================================================================
   ! set rcut value
   subroutine set_parameters(this, wallA, typeB, rho_w, eps_wf, sig_wf, delta_w)
      use MW_constants, only: kJpermol2hartree, angstrom2bohr, pi
      implicit none
      type(MW_steele_t), intent(inout) :: this
      integer, intent(in) :: wallA
      integer, intent(in) :: typeB
      real(wp), intent(in) :: eps_wf ! kJ per mol
      real(wp), intent(in) :: rho_w ! angstrom**-3
      real(wp), intent(in) :: sig_wf ! angstrom
      real(wp), intent(in) :: delta_w ! angstrom

      ! Local
      real(wp) :: A, B, eps, sigma, delta, rho

      ! Convert lj_eps from kJ/mol to atomic unit
      eps = kJpermol2hartree * eps_wf
      ! convert lj_sig from Angstrom to atomic unit
      sigma = angstrom2bohr * sig_wf
      delta = angstrom2bohr * delta_w
      rho = rho_w / angstrom2bohr**3

      ! Convert (eps,sig) to (A,B)
      A = 4.0_wp / 5.0_wp * pi * rho * eps * delta * sigma**12
      B = - 2.0_wp * pi * rho * eps * delta * sigma**6

      this%A(wallA,typeB) = A

      this%B(wallA,typeB) = B

      this%delta(wallA) = delta

   end subroutine set_parameters

   !================================================================================
   ! Print the data structure to ounit
   subroutine print_type(this, species_names, ounit)
      use MW_kinds, only: PARTICLE_NAME_LEN
      implicit none
      type(MW_steele_t), intent(in) :: this               ! structure to be printed
      character(PARTICLE_NAME_LEN),   intent(in) :: species_names(:)   ! names of the species
      integer,                   intent(in) :: ounit              ! output unit

      integer :: i, j

      write(ounit, '("|steele| cut-off distance: ",es12.5)') this%rcut
      write(ounit, '("|steele| ",8x,1x,8x,1x,"     A      ",1x,"     B      ")')

      do j = 1, this%num_walls
         do i = 1, this%num_species
            write(ounit, '("|steele| ",a8,1x,es12.5,1x,es12.5,1x,es12.5)') &
                  species_names(i), this%A(i,j), this%B(i,j), this%delta(j)
         end do
      end do

   end subroutine print_type

   ! ==============================================================================
   ! Compute Lennard-Jones forces felt by melt ions
   subroutine melt_forces(num_pbc, localwork, steele, box, ions, xyz_ions, &
                    electrodes, xyz_atoms, 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_steele_t), intent(in) :: steele      !< Steele potential parameters
      type(MW_box_t),   intent(in) :: box     !< Simulation box parameters

      type(MW_ion_t), intent(in) :: ions(:)        !< ions parameters
      type(MW_electrode_t), intent(inout) :: electrodes(:)        !< electrodes parameters
      real(wp),       intent(in) :: xyz_ions(:,:)  !< ions xyz positions
      real(wp),       intent(in) :: xyz_atoms(:,:)  !< ions xyz positions
      logical,              intent(in) :: compute_force !< compute_force on electrodes

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

      ! Local
      ! -----
      integer :: ierr, count
      integer :: i, num_elec_types

      select case(num_pbc)
      case(2)
            call melt_forces_2DPBC(localwork, steele, box, ions, xyz_ions, &
                    electrodes, xyz_atoms, force, stress_tensor)
      case(3)
            call melt_forces_3DPBC(localwork, steele, box, ions, xyz_ions, &
                    electrodes, xyz_atoms, force, stress_tensor)
      end select
#ifndef MW_SERIAL
      count = size(force,1)
      call MPI_Allreduce(MPI_IN_PLACE, force, count, MPI_DOUBLE_PRECISION, MPI_SUM, MW_COMM_WORLD, ierr)

      count = size(stress_tensor,1)*size(stress_tensor,2)      ! this should always be 9
      call MPI_Allreduce(MPI_IN_PLACE, stress_tensor, count, MPI_DOUBLE_PRECISION, MPI_SUM, MW_COMM_WORLD, ierr)      

      if (compute_force) then
         num_elec_types = size(electrodes, 1)
         do i = 1, num_elec_types
            call MPI_Allreduce(MPI_IN_PLACE, electrodes(i)%force_ions(:,9), 3, &
                  MPI_DOUBLE_PRECISION, MPI_SUM, MW_COMM_WORLD, ierr)
         end do
      end if
#endif
   end subroutine melt_forces

   ! ==============================================================================
   ! Compute Lennard-Jones potential felt by melt ions
   subroutine energy(num_pbc, localwork, steele, box, ions, xyz_ions, electrodes, xyz_atoms, 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_steele_t), intent(in) :: steele      !< Steele potential 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(:)        !< electrodes parameters
      real(wp),       intent(in) :: xyz_atoms(:,:)  !< atoms xyz positatoms

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: h       !< Potential energy due to steele interactions

      ! Local
      ! -----
      integer :: ierr

      select case(num_pbc)
      case(2)
            call energy_2DPBC(localwork, steele, box, ions, xyz_ions, electrodes, xyz_atoms, h)
      case(3)
            call energy_3DPBC(localwork, steele, box, ions, xyz_ions, electrodes, xyz_atoms, h)
      end select

#ifndef MW_SERIAL
      call MPI_Allreduce(MPI_IN_PLACE, h, 1, MPI_DOUBLE_PRECISION, MPI_SUM, MW_COMM_WORLD, ierr)
#endif
   end subroutine energy

   ! ================================================================================
   include 'steele_2DPBC.inc'
   include 'steele_3DPBC.inc'
   ! ================================================================================
   include 'minimum_image_displacement_2DPBC.inc'
   include 'minimum_image_displacement_3DPBC.inc'
   include 'minimum_image_distance_2DPBC.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_steele
