module MW_external_field
   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_localwork, only: MW_localwork_t
   use MW_constants, only: fourpi, eightpi_inv
   implicit none

   private

   ! Public type
   ! -----------
   public :: MW_external_field_t

   ! Public subroutines
   ! ------------------
   public :: define_type
   public :: void_type
   public :: print_type
   public :: melt_forces
   public :: energy
   public :: Efield_elec_potential
   public :: Dfield_gradQelec_potential
   public :: Dfield_elec_potentialb
   public :: Dfield_elec_potentialA
   public :: field2mumelt_electricfieldb
   public :: field2mumelt_electricfieldA
   public :: elec_polarization
   public :: melt_polarization
   public :: dipole_polarization
   public :: initialize_polarization

   ! Data structure to handle external electric field
   type MW_external_field_t
      integer :: field_type                               ! field type (electric field or displacement)
      real(wp), allocatable :: electric_field(:)           ! electric field
      real(wp), allocatable :: displacement_field(:)       ! field
      logical,  allocatable :: apply_direction(:)          ! apply field in x, y or z?
   end type MW_external_field_t

   integer, parameter, public :: FIELD_TYPE_ELECTRIC = 2
   integer, parameter, public :: FIELD_TYPE_DISPLACEMENT = 3

contains

   !================================================================================
   ! Define the data structure
   subroutine define_type(this, field_type, amplitude, direction)
      use MW_errors, only: &
            MW_errors_allocate_error => allocate_error, &
            MW_errors_parameter_error => parameter_error

      implicit none
      ! Parameters
      ! ---------_
      type(MW_external_field_t), intent(inout) :: this      !> structure to be defined
      character, intent(in) :: field_type !> field type
      real(wp), intent(in) :: amplitude !> field amplitude
      real(wp), intent(in) :: direction(:) !> field direction

      ! Local
      ! -----
      integer :: ierr, i
      real(wp) :: fvalue

      select case (field_type)
      case("E")
        this%field_type = FIELD_TYPE_ELECTRIC
      case("D")
        this%field_type = FIELD_TYPE_DISPLACEMENT
      case default
         call MW_errors_parameter_error("define_type", "external_field.f90",&
               "type", field_type)
      end select

      allocate(this%electric_field(3), this%displacement_field(3), &
           this%apply_direction(3), stat=ierr)
      if (ierr /= 0) &
            call MW_errors_allocate_error("define_type", "external_field.f90", ierr)

      this%electric_field(:) = 0.0_wp
      this%displacement_field(:) = 0.0_wp
      this%apply_direction(:) = .false.

      select case(this%field_type)
      case (FIELD_TYPE_ELECTRIC)
         do i = 1, 3
            fvalue = amplitude * direction(i)
            this%electric_field(i) = fvalue
            this%displacement_field(i) = 0.0_wp
            if (abs(direction(i)) > 0.0_wp) then
               this%apply_direction(i) = .true.
            else
               this%apply_direction(i) = .false.
            end if
         end do

      case(FIELD_TYPE_DISPLACEMENT)
         do i = 1, 3
            fvalue = amplitude * direction(i)
            this%electric_field(i) = 0.0_wp
            this%displacement_field(i) = fvalue
            if (abs(direction(i)) > 0.0_wp) then
               this%apply_direction(i) = .true.
            else
               this%apply_direction(i) = .false.
            end if
         end do
      end select

   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_external_field_t), intent(inout) :: this

      integer :: ierr

      if (allocated(this%electric_field)) then
         deallocate(this%electric_field, this%displacement_field, &
              this%apply_direction, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("void_type", "external_field.f90", ierr)
         end if
      end if

      this%field_type = 0
   end subroutine void_type

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

      write(ounit, '("|external field| Type:                      ",i12)') this%field_type
      write(ounit, '("|external field| Electric Field:            ",3(1x,es12.5))') this%electric_field(:)
      write(ounit, '("|external field| Displacement Field:        ",3(1x,es12.5))') this%displacement_field(:)
      write(ounit, '("|external field| Apply in direction (x,y,z):",3(1x,L12))') this%apply_direction(:)

   end subroutine print_type

   ! ==============================================================================
   ! Compute external field forces felt by melt ions
   subroutine melt_forces(field, box, electrode, ions, xyz_atoms, q_elec, polarization_field, force)
      implicit none
      
      ! Parameters in
      ! -------------
      type(MW_external_field_t), intent(inout) :: field      !< External field parameters
      type(MW_box_t), intent(in) :: box
      type(MW_electrode_t), intent(in) :: electrode(:)
      type(MW_ion_t), intent(in) :: ions(:)        !< ions parameters
      real(wp),       intent(in) :: xyz_atoms(:,:)  !< atoms xyz positions
      real(wp),       intent(in) :: q_elec(:)
      real(wp),       intent(inout) :: polarization_field(:,:)
      
      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: force(:,:) !< force on ions due to external field
      
      ! Locals
      ! ------
      integer :: num_ion_types
      integer :: itype, i, ixyz, iion
      integer :: count_n, offset_n
      real(wp) :: qi
      real(wp) :: Efield(3)
      
      num_ion_types = size(ions,1)
      force(:,:) = 0.0_wp
      Efield(:) = 0.0_wp
      
      select case(field%field_type)
      
      case(FIELD_TYPE_ELECTRIC)
      
      do itype = 1, num_ion_types
         count_n = ions(itype)%count
         offset_n = ions(itype)%offset
         qi = ions(itype)%charge
         do iion = 1, count_n
            i = iion + offset_n
            do ixyz = 1, 3
               force(i, ixyz) = force(i, ixyz) + qi * field%electric_field(ixyz)
            end do
         end do
      end do
      
      case(FIELD_TYPE_DISPLACEMENT)
      
      call elec_polarization(box, electrode, xyz_atoms, q_elec, polarization_field)
      
      do ixyz = 1,3
         if(field%apply_direction(ixyz)) then
            Efield(ixyz) = field%displacement_field(ixyz) - fourpi * polarization_field(ixyz,1)
         endif
      enddo
      
      do itype = 1, num_ion_types
         count_n = ions(itype)%count
         offset_n = ions(itype)%offset
         qi = ions(itype)%charge
         do iion = 1, count_n
            i = iion + offset_n
            do ixyz = 1, 3
               force(i, ixyz) = force(i, ixyz) + qi * Efield(ixyz)
            end do
         end do
      end do
      
      end select

   end subroutine melt_forces

   ! ==============================================================================
   ! Compute the interaction energy with the external field
   subroutine energy(field, box, polarization_field, h)
      implicit none

      ! Parameters in
      ! -------------
      type(MW_external_field_t), intent(inout) :: field      !< External field parameters
      type(MW_box_t),   intent(in) :: box     !< Simulation box parameters
      real(wp), intent(in) :: polarization_field(:,:)

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: h       !< Potential energy due to interaction with the external field

      ! Local
      ! -----
      integer :: ixyz

      h = 0.0_wp

      select case(field%field_type)

      case(FIELD_TYPE_ELECTRIC)
      do ixyz = 1, 3
         h = h - (field%electric_field(ixyz) * polarization_field(ixyz,1))
      end do
      h = box%volume * h

      case(FIELD_TYPE_DISPLACEMENT)
      do ixyz = 1, 3
         if (field%apply_direction(ixyz)) then
            h = h - (field%displacement_field(ixyz) - fourpi * polarization_field(ixyz,1))**2
         endif
      end do
      h = box%volume * eightpi_inv * h

      end select

   end subroutine energy

   ! ================================================================================
   ! Compute the potential felt by an electrode atom due to an external electric field
   subroutine Efield_elec_potential(field, electrode, xyz_atoms, V)
      implicit none
     
      ! Parameters in
      ! -------------
      type(MW_external_field_t), intent(in) :: field      !< Lennard-Jones potential parameters
      type(MW_electrode_t), intent(in) :: electrode(:)  !< electrode parameters
      real(wp),       intent(in) :: xyz_atoms(:,:)  !< atoms xyz positions
      ! Parameter out
      ! -------------
      real(wp), intent(inout) :: V(:)
     
      ! Local
      ! -----
      integer :: num_elec_types
      integer :: itype, ielec, i, ixyz
      integer :: count_n, offset_n
     
      num_elec_types = size(electrode,1)
      V(:) = 0.0_wp
     
      do itype = 1, num_elec_types
         count_n = electrode(itype)%count
         offset_n = electrode(itype)%offset
         do ielec = 1, count_n
            i = ielec + offset_n
            V(i) = 0.0_wp
            do ixyz = 1, 3
               V(i) = V(i) - xyz_atoms(i,ixyz) * field%electric_field(ixyz)
            end do
         end do
      end do

   end subroutine Efield_elec_potential

   ! ================================================================================
   ! Compute the potential felt by an electrode atom due to an external electric displacement
   subroutine Dfield_gradQelec_potential(box, field, electrode, xyz_atoms, q_elec, polarization_field, gradQ_V)
      implicit none
      ! Parameters in
      ! -------------
      type(MW_external_field_t), intent(inout) :: field      !< Lennard-Jones potential parameters
      type(MW_electrode_t), intent(in) :: electrode(:)  !< electrode parameters
      type(MW_box_t), intent(in) :: box  !< box parameters
      real(wp), intent(in) :: xyz_atoms(:,:)
      real(wp), intent(in) :: q_elec(:)
     
      ! Parameter out
      ! -------------
      real(wp), intent(inout) :: polarization_field(:,:)
      real(wp), intent(inout) :: gradQ_V(:,:)
     
      ! Local
      ! -----
      integer :: num_elec_types
      integer :: ixyz
      integer :: itype, ielec, i, jtype, jelec, j
      integer :: count_n, offset_n, count_m, offset_m
     
      num_elec_types = size(electrode,1)
     
      !call elec_polarization(box, electrode, xyz_atoms, q_elec, polarization_field)
     
      do itype = 1, num_elec_types
         count_n = electrode(itype)%count
         offset_n = electrode(itype)%offset
         do ielec = 1, count_n
            i = ielec + offset_n
            do jtype = 1, num_elec_types
               count_m = electrode(jtype)%count
               offset_m = electrode(jtype)%offset
               do jelec = 1, count_m
                  j = jelec + offset_m
                  do ixyz = 1, 3
                     if(field%apply_direction(ixyz)) then
                        gradQ_V(j,i) = gradQ_V(j,i) + fourpi / box%volume * xyz_atoms(i, ixyz) * xyz_atoms(j, ixyz)
                     end if
                  end do
               end do
            end do
         end do
      end do

   end subroutine Dfield_gradQelec_potential


   ! ================================================================================
   ! Compute the potential felt by an electrode atom due to an external electric displacement
   subroutine Dfield_elec_potentialA(box, field, electrode, xyz_atoms, q_elec, polarization_field, V)
      implicit none
      ! Parameters in
      ! -------------
      type(MW_external_field_t), intent(inout) :: field      !< Lennard-Jones potential parameters
      type(MW_electrode_t), intent(in) :: electrode(:)  !< electrode parameters
      type(MW_box_t), intent(in) :: box  !< box parameters
      real(wp), intent(in) :: xyz_atoms(:,:)
      real(wp), intent(in) :: q_elec(:)

      ! Parameter out
      ! -------------
      real(wp), intent(inout) :: polarization_field(:,:)
      real(wp), intent(inout) :: V(:)

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: itype, ielec, i, ixyz
      integer :: count_n, offset_n
      real(wp) :: Efield(3)

      num_elec_types = size(electrode,1)
      V(:) = 0.0_wp
      Efield(:) = 0.0_wp

      call elec_polarization(box, electrode, xyz_atoms, q_elec, polarization_field)

      do ixyz = 1,3
         if(field%apply_direction(ixyz)) then
            Efield(ixyz) = - fourpi * (polarization_field(ixyz,3) + polarization_field(ixyz,4))
         endif
      enddo

      do itype = 1, num_elec_types
         count_n = electrode(itype)%count
         offset_n = electrode(itype)%offset
         do ielec = 1, count_n
            i = ielec + offset_n
            do ixyz = 1, 3
               V(i) = V(i) - xyz_atoms(i, ixyz) * Efield(ixyz)
            end do
         end do
      end do

   end subroutine Dfield_elec_potentialA


   ! ================================================================================
   ! Compute the potential felt by an electrode atom due to an external electric displacement
   subroutine Dfield_elec_potentialb(field, electrode, xyz_atoms, polarization_field, V)
      implicit none
      ! Parameters in
      ! -------------
      type(MW_external_field_t), intent(inout) :: field      !< Lennard-Jones potential parameters
      type(MW_electrode_t), intent(in) :: electrode(:)  !< electrode parameters
      real(wp), intent(in) :: xyz_atoms(:,:)

      ! Parameter out
      ! -------------
      real(wp), intent(inout) :: polarization_field(:,:)
      real(wp), intent(inout) :: V(:)

      ! Local
      ! -----
      integer :: num_elec_types
      integer :: itype, ielec, i, ixyz
      integer :: count_n, offset_n
      real(wp) :: Efield(3)

      num_elec_types = size(electrode,1)
      V(:) = 0.0_wp
      Efield(:) = 0.0_wp

      do ixyz = 1,3
         if(field%apply_direction(ixyz)) then
            Efield(ixyz) = field%displacement_field(ixyz) - fourpi * polarization_field(ixyz,2)
         endif
      enddo

      do itype = 1, num_elec_types
         count_n = electrode(itype)%count
         offset_n = electrode(itype)%offset
         do ielec = 1, count_n
            i = ielec + offset_n
            do ixyz = 1, 3
               V(i) = V(i) - xyz_atoms(i, ixyz) * Efield(ixyz)
            end do
         end do
      end do
   end subroutine Dfield_elec_potentialb


   ! ================================================================================
   ! Compute the polarization field for electrode atoms
   ! P(t) = 1/V * sum (q_i(t)*r_i)
   subroutine elec_polarization(box, electrode, xyz_atoms, q_elec, polarization_field)
      implicit none
     
      ! Parameters in
      ! -------------
      type(MW_box_t), intent(in) :: box !< box parameters
      type(MW_electrode_t), intent(in) :: electrode(:)  !< electrode parameters
      real(wp), intent(in) :: xyz_atoms(:,:)
      real(wp), intent(in) :: q_elec(:)
      ! Parameters out
      ! --------------
       real(wp), intent(inout) :: polarization_field(:,:)
      
      ! Local
      ! -----
      integer :: num_elec_types
      integer :: itype, ielec, i, ixyz
      integer :: count_n, offset_n

      num_elec_types = size(electrode,1)
      polarization_field(:,3) = 0.0_wp

      do itype = 1, num_elec_types
         count_n = electrode(itype)%count
         offset_n = electrode(itype)%offset
         do ielec = 1, count_n
            i = ielec + offset_n
            do ixyz = 1, 3
               polarization_field(ixyz,3) = polarization_field(ixyz,3) + q_elec(i) * xyz_atoms(i,ixyz)
            end do
         end do
      end do

      polarization_field(:,3) = polarization_field(:,3) / box%volume
      polarization_field(:,1) = polarization_field(:,2) + polarization_field(:,3) + polarization_field(:,4)

   end subroutine elec_polarization

   ! ================================================================================
   ! Compute the polarization field for melt ions
   ! P(t+dt) = P(t) + 1/V * sum_(qi * delta_i(t+dt))
   ! with delta_i(t+dt) = (r_i(t+dt) - r_i(t)) // without PBC folding
   !
   subroutine melt_polarization(num_pbc, box, ions, xyz_ions_prev, xyz_ions_now, polarization_field)
      implicit none
    
      ! Parameters in
      ! -------------
      integer, intent(in)        :: num_pbc
      type(MW_box_t), intent(in) :: box
      type(MW_ion_t), intent(in) :: ions(:)  !< electrode parameters
      real(wp), intent(in) :: xyz_ions_prev(:,:) !< ions position at previous time step
      real(wp), intent(in) :: xyz_ions_now(:,:)  !< ions position at current time step

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: polarization_field(:,:)

      ! Local
      ! -----
      integer :: num_ion_types
      integer :: itype, iion, i, ixyz
      real(wp) :: dr(3), drnorm2
      real(wp) :: volfactor, qi
      real(wp) :: tmpfield(3)
      integer :: count_n, offset_n

      num_ion_types = size(ions,1)
      volfactor = 1.0_wp / box%volume
      tmpfield(:) = 0.0_wp

      select case(num_pbc)

      case(2)
      do itype = 1, num_ion_types
         count_n = ions(itype)%count
         offset_n = ions(itype)%offset
         qi = ions(itype)%charge
         do iion = 1, count_n
            i = iion + offset_n

            call minimum_image_displacement_2DPBC(box%length(1), box%length(2), box%length(3), &
                 xyz_ions_prev(i,1), xyz_ions_prev(i,2), xyz_ions_prev(i,3), &
                 xyz_ions_now(i,1), xyz_ions_now(i,2), xyz_ions_now(i,3), &
                 dr(1),dr(2),dr(3), drnorm2)
            do ixyz = 1, 3
               tmpfield(ixyz) = tmpfield(ixyz) + qi*dr(ixyz)
            end do
         end do
      end do

      case(3)
      do itype = 1, num_ion_types
         count_n = ions(itype)%count
         offset_n = ions(itype)%offset
         qi = ions(itype)%charge
         do iion = 1, count_n
            i = iion + offset_n

            call minimum_image_displacement_3DPBC(box%length(1), box%length(2), box%length(3), &
                 xyz_ions_prev(i,1), xyz_ions_prev(i,2), xyz_ions_prev(i,3), &
                 xyz_ions_now(i,1), xyz_ions_now(i,2), xyz_ions_now(i,3), &
                 dr(1),dr(2),dr(3), drnorm2)
            do ixyz = 1, 3
               tmpfield(ixyz) = tmpfield(ixyz) + qi*dr(ixyz)
            end do
         end do
      end do

      end select

      polarization_field(:,2) = polarization_field(:,2) + volfactor*tmpfield(:)
      polarization_field(:,1) = polarization_field(:,2) + polarization_field(:,3) + polarization_field(:,4)
   end subroutine melt_polarization

   ! ================================================================================
   include 'minimum_image_displacement_3DPBC.inc'
   include 'minimum_image_displacement_2DPBC.inc'

   ! ================================================================================
   ! Compute the polarization field for the dipoles
   ! P(t)=1/V * sum mui(t)
   !
   subroutine dipole_polarization(box, ions, dipole_ions, polarization_field)
      implicit none
      ! Parameters in
      type(MW_box_t), intent(in) :: box
      type(MW_ion_t), intent(in) :: ions(:)  !< electrode parameters
      real(wp), intent(in) :: dipole_ions(:,:)

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: polarization_field(:,:)

      ! Local
      ! -----
      integer :: num_ion_types
      integer :: itype, iion, i, ixyz
      real(wp) :: volfactor
      integer :: count_n, offset_n

      num_ion_types = size(ions,1)
      volfactor = 1.0_wp / box%volume
      polarization_field(:,4) = 0.0_wp

      do itype = 1, num_ion_types
         count_n = ions(itype)%count
         offset_n = ions(itype)%offset
         do iion = 1, count_n
            i = iion + offset_n
            do ixyz = 1, 3
               polarization_field(ixyz,4) = polarization_field(ixyz,4) + dipole_ions(i,ixyz)
            end do
         end do
      end do

      polarization_field(:,4) = polarization_field(:,4) / box%volume
      polarization_field(:,1) = polarization_field(:,2) + polarization_field(:,3) + polarization_field(:,4)
   end subroutine dipole_polarization

   ! ================================================================================
   ! Initialze the polarization of the melt ions
   ! P(t=0)=sum qi * ri
   !
   subroutine initialize_polarization(box, xyz_ions, ions, polarization_field)
      implicit none
      ! Parameters in
      type(MW_box_t), intent(in) :: box
      real(wp), intent(in) :: xyz_ions(:,:)
      type(MW_ion_t), intent(in) :: ions(:)  !< electrode parameters

      ! Parameters out
      ! --------------
      real(wp), intent(inout) :: polarization_field(:,:)

      ! Local
      ! -----
      integer :: num_ion_types
      integer :: itype, iion, i, ixyz
      real(wp) :: volfactor, qi
      integer :: count_n, offset_n

      num_ion_types = size(ions,1)
      volfactor = 1.0_wp / box%volume
      polarization_field(:,2) = 0.0_wp

      do itype = 1, num_ion_types
         count_n = ions(itype)%count
         offset_n = ions(itype)%offset
         qi = ions(itype)%charge
         do iion = 1, count_n
            i = iion + offset_n
            do ixyz = 1, 3
               polarization_field(ixyz,2) = polarization_field(ixyz,2) + qi*xyz_ions(i,ixyz)
            end do
         end do
      end do

      polarization_field(:,2) = polarization_field(:,2) * volfactor
   end subroutine initialize_polarization

   subroutine field2mumelt_electricfieldb(field, ions, polarization_field, electric_field)
      implicit none
      ! Parameters in
      ! -------------
      type(MW_external_field_t), intent(in) :: field      !< Lennard-Jones potential parameters
      type(MW_ion_t), intent(in) :: ions(:)  !< electrode parameters
      real(wp), intent(in) :: polarization_field(:,:)

      ! Parameter out
      ! -------------
      real(wp), intent(out) :: electric_field(:,:)

      ! Local
      ! -----
      integer :: i, itype, ixyz
      electric_field = 0.0_wp

      select case(field%field_type)

      case(FIELD_TYPE_ELECTRIC)
      do itype = 1,size(ions,1)
         do i=ions(itype)%offset+1,ions(itype)%offset+ions(itype)%count
            do ixyz = 1, 3
               electric_field(i,ixyz) = -field%electric_field(ixyz)
            end do
         end do
      end do

      case(FIELD_TYPE_DISPLACEMENT)
      do itype = 1,size(ions,1)
         do i=ions(itype)%offset+1,ions(itype)%offset+ions(itype)%count
            do ixyz = 1, 3
               if (field%apply_direction(ixyz)) then
                  electric_field(i,ixyz) =  -(field%displacement_field(ixyz) - fourpi * polarization_field(ixyz,2))
               end if
            end do
         end do
      end do

      end select
   end subroutine field2mumelt_electricfieldb

   subroutine field2mumelt_electricfieldA(box, field, electrode, xyz_atoms, q_elec, ions, &
          dipole_ions, polarization_field, electric_field)
      implicit none
      ! Parameters in
      ! -------------
      type(MW_external_field_t), intent(inout) :: field      !< Lennard-Jones potential parameters
      type(MW_electrode_t), intent(in) :: electrode(:)  !< electrode parameters
      type(MW_box_t), intent(in) :: box  !< box parameters
      type(MW_ion_t), intent(in) :: ions(:)  !< electrode parameters
      real(wp), intent(in) :: xyz_atoms(:,:)
      real(wp), intent(in) :: q_elec(:)
      real(wp), intent(in) :: dipole_ions(:,:)

      ! Parameter out
      ! -------------
      real(wp), intent(inout) :: polarization_field(:,:)
      real(wp), intent(out) :: electric_field(:,:)

      ! Local
      ! -----
      integer :: i, itype, ixyz
      electric_field = 0.0_wp

      call elec_polarization(box, electrode, xyz_atoms, q_elec, polarization_field)
      call dipole_polarization(box, ions, dipole_ions, polarization_field)

      do itype = 1,size(ions,1)
         do i=ions(itype)%offset+1,ions(itype)%offset+ions(itype)%count
            do ixyz = 1,3
               if(field%apply_direction(ixyz)) then
                  electric_field(i,ixyz) =  fourpi * (polarization_field(ixyz,3) +polarization_field(ixyz,4))
               end if
            end do
         end do
      end do

    end subroutine field2mumelt_electricfieldA

end module MW_external_field
