module MW_thomas_fermi
   use MW_kinds, only: wp
   use MW_box, only: MW_box_t
   use MW_electrode, only: MW_electrode_t
   use MW_ion, only: MW_ion_t
   use MW_constants, only: pi, fourpi
   implicit none
   private

   ! Public type
   ! -----------
   public :: MW_thomas_fermi_t

   ! Public subroutines
   ! ------------------
   public :: define_type
   public :: void_type
   public :: print_type
   public :: kinetic_gradQelec_potential
   public :: kinetic_elec_potential
   public :: kinetic_energy

   ! Data structure to handle Thomas-fermi model
   type MW_thomas_fermi_t
      real(wp) :: ltf = 0.0_wp                                     ! Thomas-fermi length
      real(wp) :: fourpi_ltfsquare = 0.0_wp                        ! 4pi*ltf**2
      real(wp) :: voronoi_volume = 0.0_wp                          ! Voronoi volume of the particle
   end type MW_thomas_fermi_t


contains

   !================================================================================
   ! Define the data structure
   subroutine define_type(this, ltf, voronoi_volume)
      implicit none
      ! Parameters
      ! ---------_
      type(MW_thomas_fermi_t), intent(inout) :: this      !> structure to be defined
      real(wp), intent(in) :: ltf
      real(wp), intent(in) :: voronoi_volume
      ! Local
      ! -----
      real(wp) :: ltfsquare

      ltfsquare = ltf * ltf
      this%ltf = ltf
      this%fourpi_ltfsquare = fourpi * ltfsquare
      this%voronoi_volume = voronoi_volume


   end subroutine define_type

   !================================================================================
   ! Void the data structure
   subroutine void_type(this)
      implicit none
      type(MW_thomas_fermi_t), intent(inout) :: this

      this%ltf = 0.0_wp
      this%fourpi_ltfsquare = 0.0_wp
      this%voronoi_volume = 0.0_wp

   end subroutine void_type


   !================================================================================
   ! Print the data structure to ounit
   subroutine print_type(this, ounit)
      implicit none
      type(MW_thomas_fermi_t), intent(in) :: this               ! structure to be printed
      integer,                   intent(in) :: ounit              ! output unit

      write(ounit, '("|electrode|  Thomas-Fermi length ltf = ",es12.5," Voronoi volume = ",es12.5)') &
            this%ltf, this%voronoi_volume

   end subroutine print_type

   ! ================================================================================
   ! Compute the 2nd order contribution of the kinetic energy to the  potential felt by an electrode atom
   subroutine kinetic_gradQelec_potential(thomas_fermi, electrodes, gradQ_V)
     implicit none

     ! Parameters in
     ! -------------
      type(MW_thomas_fermi_t), intent(in) :: thomas_fermi(:)      !< Lennard-Jones potential parameters
      type(MW_electrode_t), intent(in) :: electrodes(:)  !< electrode parameters
      ! Parameter out
      ! -------------
      real(wp), intent(inout) :: gradQ_V(:,:)

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: itype, ielec, i
      integer :: count_n, offset_n
      real(wp) :: fourpi_ltfsquarei, voronoii

      num_elec_types = size(electrodes,1)

      do itype = 1, num_elec_types
         if(thomas_fermi(itype)%ltf>0) then
            count_n = electrodes(itype)%count
            offset_n = electrodes(itype)%offset
            fourpi_ltfsquarei = thomas_fermi(itype)%fourpi_ltfsquare
            voronoii = thomas_fermi(itype)%voronoi_volume
            do ielec = 1, count_n
               i = ielec + offset_n
               gradQ_V(i,i) = gradQ_V(i,i) + fourpi_ltfsquarei / voronoii
            end do
         end if
      end do
   end subroutine kinetic_gradQelec_potential


   ! ================================================================================
   ! Compute the 2nd order contribution of the kinetic energy to the  potential felt by an electrode atom
   subroutine kinetic_elec_potential(this, electrodes, q_elec, V)
     implicit none

     ! Parameters in
     ! -------------
      type(MW_thomas_fermi_t), intent(in) :: this(:)      !< Lennard-Jones potential parameters
      type(MW_electrode_t), intent(in) :: electrodes(:)  !< electrode parameters
      real(wp), intent(in) :: q_elec(:)
      ! Parameter out
      ! -------------
      real(wp), intent(inout) :: V(:)

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: itype, ielec, i
      integer :: count_n, offset_n
      real(wp) :: fourpi_ltfsquarei, voronoii

      num_elec_types = size(electrodes,1)

      do itype = 1, num_elec_types
         if(this(itype)%ltf>0) then
            count_n = electrodes(itype)%count
            offset_n = electrodes(itype)%offset
            fourpi_ltfsquarei = this(itype)%fourpi_ltfsquare
            voronoii = this(itype)%voronoi_volume
            do ielec = 1, count_n
               i = ielec + offset_n
               V(i) = fourpi_ltfsquarei * q_elec(i) / voronoii
            end do
         end if
      end do
    end subroutine kinetic_elec_potential


   ! ================================================================================
   ! Compute the kinetic energy of the electron gas
   subroutine kinetic_energy(this, electrodes, q_elec, h)
     implicit none

     ! Parameters in
     ! -------------
      type(MW_thomas_fermi_t), intent(in) :: this(:)      !< Lennard-Jones potential parameters
      type(MW_electrode_t), intent(in) :: electrodes(:)  !< electrode parameters
      real(wp), intent(in) :: q_elec(:)
      ! Parameter out
      ! -------------
      real(wp), intent(inout) :: h

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: itype, ielec, i
      integer :: count_n, offset_n
      real(wp) :: fourpi_ltfsquarei, voronoii

      num_elec_types = size(electrodes,1)
      h = 0.0_wp

      do itype = 1, num_elec_types
         if(this(itype)%ltf>0) then
            count_n = electrodes(itype)%count
            offset_n = electrodes(itype)%offset
            fourpi_ltfsquarei = this(itype)%fourpi_ltfsquare
            voronoii = this(itype)%voronoi_volume
            do ielec = 1, count_n
               i = ielec + offset_n
               h = h + fourpi_ltfsquarei * 0.5 * q_elec(i)*q_elec(i) / voronoii
            end do
         end if
      end do
    end subroutine kinetic_energy


end module MW_thomas_fermi
