!> Module handling velocities transformation
module MW_velocity
   use MW_kinds, only: wp
   use MW_constants, only: boltzmann
   use MW_ion, only: MW_ion_t
   use MW_random, only: &
         MW_random_normal_random_number => normal_random_number
   use MW_molecule, only: MW_molecule_t
   use MW_errors, only: MW_errors_parameter_error => parameter_error

   implicit none
   private

   ! Public subroutines
   ! ------------------
   public :: initialize
   public :: rescale
   public :: com_velocity_compute
   public :: com_velocity_reset

   integer, parameter, public :: RESCALE_METHOD_NONE = 0
   integer, parameter, public :: RESCALE_METHOD_GLOBAL = 1
   integer, parameter, public :: RESCALE_METHOD_SPECIES = 2
   integer, parameter, public :: RESCALE_METHOD_MOLECULES_GLOBAL = 3
   integer, parameter, public :: RESCALE_METHOD_MOLECULES_INDEPENDENT = 4

contains

   ! ================================================================================
   !> Compute the center of mass velocity
   !!
   !! $v_com = \frac{\sum_i^N m_i v_i}{\sum_i^N m_i}$
   subroutine com_velocity_compute(ions, ions_vx, ions_vy, ions_vz, &
         com_vx, com_vy, com_vz, com_ke)
      implicit none
      ! Parameters
      ! ----------
      type(MW_ion_t), intent(in) :: ions(:)   !< ions parameters
      real(wp), intent(in) :: ions_vx(:)      !< ions velocity in x
      real(wp), intent(in) :: ions_vy(:)      !< ions velocity in y
      real(wp), intent(in) :: ions_vz(:)      !< ions velocity in z
      real(wp), intent(out) :: com_vx         !< com velocity in x
      real(wp), intent(out) :: com_vy         !< com velocity in y
      real(wp), intent(out) :: com_vz         !< com velocity in z
      real(wp), intent(out) :: com_ke         !< kinetic energy of the com

      ! Local
      ! -----
      integer :: num_ion_types, itype, iion, i
      real(wp) :: mi, mass_total
      real(wp) :: com_px, com_py, com_pz

      com_px = 0.0_wp
      com_py = 0.0_wp
      com_pz = 0.0_wp
      mass_total = 0.0_wp

      num_ion_types = size(ions,1)
      do itype = 1, num_ion_types
         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then
            mi = ions(itype)%mass
            do i = 1, ions(itype)%count
               iion = i + ions(itype)%offset
               com_px = com_px + mi * ions_vx(iion)
               com_py = com_py + mi * ions_vy(iion)
               com_pz = com_pz + mi * ions_vz(iion)
               mass_total = mass_total + mi
            end do
         end if
      end do

      com_ke = 0.5_wp * (com_px*com_px + com_py*com_py + com_pz*com_pz) / mass_total

      com_vx = com_px / mass_total
      com_vy = com_py / mass_total
      com_vz = com_pz / mass_total

   end subroutine com_velocity_compute

   ! ================================================================================
   ! Set the center of mass velocity to 0
   subroutine com_velocity_reset(com_vx, com_vy, com_vz, &
         ions, ions_vx, ions_vy, ions_vz)
      implicit none
      ! Parameters
      ! ----------
      real(wp), intent(in)        :: com_vx    ! com velocity in x
      real(wp), intent(in)        :: com_vy    ! com velocity in y
      real(wp), intent(in)        :: com_vz    ! com velocity in z
      type(MW_ion_t), intent(in) :: ions(:)   !< ions parameters
      real(wp), intent(inout) :: ions_vx(:)      !< ions velocity in x
      real(wp), intent(inout) :: ions_vy(:)      !< ions velocity in y
      real(wp), intent(inout) :: ions_vz(:)      !< ions velocity in z

      ! Local
      ! -----
      integer :: num_ion_types, itype, iion, i

      num_ion_types = size(ions,1)
      do itype = 1, num_ion_types
         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then
            do i = 1, ions(itype)%count
               iion = i + ions(itype)%offset
               ions_vx(iion) = ions_vx(iion) - com_vx
               ions_vy(iion) = ions_vy(iion) - com_vy
               ions_vz(iion) = ions_vz(iion) - com_vz
            end do
         end if
      end do
   end subroutine com_velocity_reset

   ! ================================================================================
   ! Initialize particule velocities with random normal distribution
   subroutine initialize(temperature, num_dof, ions, velocity_ions)
      implicit none
      ! Parameters
      ! ----------
      real(wp), intent(in) :: temperature
      integer, intent(in) :: num_dof
      type(MW_ion_t), intent(in) :: ions(:)
      real(wp), intent(inout) :: velocity_ions(:,:)

      ! Locals
      ! ------
      integer :: itype, iion, ixyz
      integer :: num_ion_types
      real(wp) :: harvest
      real(wp) :: sigma
      real(wp) :: kt

      num_ion_types = size(ions, 1)
      kt = boltzmann * temperature
      do itype = 1, num_ion_types
         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then
            sigma = sqrt(kt / ions(itype)%mass)
            do iion = ions(itype)%offset+1, ions(itype)%offset + ions(itype)%count
               do ixyz = 1, 3
                  call MW_random_normal_random_number(harvest)
                  velocity_ions(iion,ixyz) = harvest * sigma
               end do
            end do
         end if
      end do
      call global_scaling(temperature, num_dof, ions, velocity_ions)
   end subroutine initialize

   ! ================================================================================
   ! Rescale particule velocities
   subroutine rescale(method, temperature, num_dof, molecules, ions, velocity_ions)
      implicit none
      ! Parameters
      ! ----------
      integer, intent(in) :: method                 !< Rescale method to apply
      real(wp), intent(in) :: temperature           !< Target temperature
      integer, intent(in) :: num_dof                !< Number of degrees of freedom
      type(MW_molecule_t), intent(in) :: molecules(:)         !< molecules parameters
      type(MW_ion_t), intent(in) :: ions(:)         !< ions parameters
      real(wp), intent(inout) :: velocity_ions(:,:) !< ions velocities

      select case(method)
      case(RESCALE_METHOD_NONE)
         continue
      case(RESCALE_METHOD_GLOBAL)
         call global_scaling(temperature, num_dof, ions,  velocity_ions)
      case(RESCALE_METHOD_SPECIES)
         call species_scaling(temperature, num_dof, ions, velocity_ions)
      case(RESCALE_METHOD_MOLECULES_GLOBAL)
         call molecules_global_scaling(temperature, num_dof, molecules, ions, velocity_ions)
      case(RESCALE_METHOD_MOLECULES_INDEPENDENT)
         call molecules_independent_scaling(temperature, num_dof, molecules, ions, velocity_ions)
      case default
         call MW_errors_parameter_error("rescale", "velocity.f90", "method", method)
      end select
   end subroutine rescale

   ! ================================================================================
   ! Rescale all particules velocities by the same amount
   !
   ! The center of mass for all particles has its velocity set to 0
   ! The same scaling factor is applied to all particle velocities
   subroutine global_scaling(temperature, num_dof, ions, velocity_ions)
      implicit none
      ! Parameters
      ! ----------
      real(wp), intent(in) :: temperature           !< Target temperature
      integer, intent(in) :: num_dof                !< Number of degrees of freedom
      type(MW_ion_t), intent(in) :: ions(:)         !< ions parameters
      real(wp), intent(inout) :: velocity_ions(:,:) !< ions velocities
      ! Local
      ! -----
      integer ::iion, ixyz
      integer :: num_ion_types, itype
      real(wp) :: mi, visq, ke_total
      real(wp) :: temperature_system, scale
      real(wp) :: com_vx, com_vy, com_vz, com_ke

      call com_velocity_compute(ions, &
            velocity_ions(:,1), velocity_ions(:,2), velocity_ions(:,3), &
            com_vx, com_vy, com_vz, com_ke)

      call com_velocity_reset(com_vx, com_vy, com_vz, &
            ions, velocity_ions(:,1), velocity_ions(:,2), velocity_ions(:,3))

      ! Compute kinetic energy of the system and actual temperature of the system
      ke_total = 0.0_wp
      num_ion_types = size(ions,1)
      do itype = 1, num_ion_types
         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then
            mi = ions(itype)%mass
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               visq = 0.0_wp
               do ixyz = 1, 3
                  visq = visq + velocity_ions(iion,ixyz) * velocity_ions(iion,ixyz)
               end do
               ke_total = ke_total + 0.5_wp * mi * visq
            end do
         end if
      end do

      ! ke = num_dof*k*T/2
      temperature_system = 2.0_wp * ke_total / (real(num_dof,wp) * boltzmann)
      scale = sqrt(temperature/temperature_system)

      ! Rescale velocities
      do itype = 1, num_ion_types
         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_ions(iion,ixyz) = scale * velocity_ions(iion, ixyz)
               end do
            end do
         end if
      end do

   end subroutine global_scaling

   ! ================================================================================
   ! Rescale particules velocities by species
   !
   ! The centers of mass for each species group have their velocities set to 0.
   ! The scaling factor is applied on per species group basis.
   subroutine species_scaling(temperature, num_dof, ions, velocity_ions)
      use MW_kinds, only: wp
      use MW_constants, only: boltzmann
      use MW_ion, only: MW_ion_t
      implicit none
      ! Parameters
      ! ----------
      real(wp), intent(in) :: temperature           !< Target temperature
      integer, intent(in) :: num_dof                !< Number of degrees of freedom in the system
      type(MW_ion_t), intent(in) :: ions(:)         !< ions parameters
      real(wp), intent(inout) :: velocity_ions(:,:) !< ions velocities
      ! Local
      ! -----
      integer :: iion, ixyz
      integer :: num_ion_types, itype
      real(wp) :: mi, visq, vsq_total, ke_total
      real(wp) :: velocity_com(3)
      real(wp) :: temperature_species, scale
      integer :: num_mobile_ions, num_ions_species, num_dof_species
      real(wp) :: dof_scaled

      num_ion_types = size(ions,1)

      num_mobile_ions = 0
      do itype = 1, num_ion_types
         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then
            num_mobile_ions = num_mobile_ions + ions(itype)%count
         end if
      end do

      do itype = 1, num_ion_types

         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then

            ! Compute center of mass velocity for each ion species
            velocity_com(:) = 0.0_wp
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_com(ixyz) = velocity_com(ixyz) + velocity_ions(iion,ixyz)
               end do
            end do

            ! Subtract velocity_com from all ions
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_ions(iion,ixyz) = velocity_ions(iion,ixyz) - velocity_com(ixyz)
               end do
            end do

            ! Compute kinetic energy of the system and actual temperature of the system
            vsq_total = 0.0_wp
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               visq = 0.0_wp
               do ixyz = 1, 3
                  visq = visq + velocity_ions(iion,ixyz) * velocity_ions(iion,ixyz)
               end do
               vsq_total = vsq_total + visq
            end do
            mi = ions(itype)%mass
            ke_total = 0.5_wp * mi * vsq_total

            ! ke = num_dof*k*T/2
            num_ions_species = ions(itype)%count
            num_dof_species = 3 * num_ions_species
            dof_scaled = real(num_dof,wp) * (real(num_dof_species, wp) / real(3*num_mobile_ions, wp))
            temperature_species = 2.0_wp * ke_total / (dof_scaled * boltzmann)
            scale = sqrt(temperature/temperature_species)

            ! Rescale velocities
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_ions(iion,ixyz) = scale * velocity_ions(iion, ixyz)
               end do
            end do
         end if

      end do

   end subroutine species_scaling

   ! ================================================================================
   ! Rescale particules velocities by species
   ! The centers of mass for each molecule group have their velocities set to 0.
   ! The centers of mass for each species (not in a molecule) group have their velocities set to 0.
   ! The same scaling factor is applied to all particles
   subroutine molecules_global_scaling(temperature, num_dof, molecules, ions, velocity_ions)
      use MW_kinds, only: wp
      use MW_constants, only: boltzmann
      use MW_errors, only: &
            MW_errors_allocate_error => allocate_error, &
            MW_errors_deallocate_error => deallocate_error
      use MW_ion, only: MW_ion_t
      use MW_molecule, only: MW_molecule_t
      implicit none
      ! Parameters
      ! ----------
      real(wp), intent(in) :: temperature             !< Target temperature
      integer, intent(in) :: num_dof                  !< Number of degrees of freedom in the system
      type(MW_molecule_t), intent(in) :: molecules(:) !< molecules parameters
      type(MW_ion_t), intent(in) :: ions(:)           !< ions parameters
      real(wp), intent(inout) :: velocity_ions(:,:)   !< ions velocities
      ! Local
      ! -----
      integer :: iion, ixyz
      integer :: num_ion_types, itype
      integer :: num_molecules, imol, isite
      real(wp) :: mi, visq, ke_total
      real(wp) :: velocity_com(3)
      real(wp) :: mass_total, temperature_species, scale
      logical :: in_molecule

      num_ion_types = size(ions,1)
      num_molecules = size(molecules,1)

      ! Loop through molecules
      ! Zero center of mass velocity for each molecule types
      do imol = 1, num_molecules

         mass_total = 0.0_wp
         velocity_com(:) = 0.0_wp
         do isite = 1, molecules(imol)%num_sites
            itype = molecules(imol)%sites(isite)
            mi = ions(itype)%mass
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_com(ixyz) = velocity_com(ixyz) + mi * velocity_ions(iion,ixyz)
               end do
               mass_total = mass_total + mi
            end do
         end do

         do ixyz = 1, 3
            velocity_com(ixyz) = velocity_com(ixyz) / mass_total
         end do

         ! 0 COM velocity for this molecule
         do isite = 1, molecules(imol)%num_sites
            itype = molecules(imol)%sites(isite)
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_ions(iion,ixyz) = velocity_ions(iion,ixyz) - velocity_com(ixyz)
               end do
            end do
         end do
      end do

      ! Zero center of mass velocity for each ion species which are not in a molecule
      do itype = 1, num_ion_types
         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then
            in_molecule = .false.
            InMoleculeSearchLoop: do imol = 1, num_molecules
               do isite = 1, molecules(imol)%num_sites
                  if (itype == molecules(imol)%sites(isite)) then
                     in_molecule = .true.
                     exit InMoleculeSearchLoop
                  end if
               end do
            end do InMoleculeSearchLoop
            if (.not. in_molecule) then
               velocity_com(:) = 0.0_wp
               do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
                  do ixyz = 1, 3
                     velocity_com(ixyz) = velocity_com(ixyz) + velocity_ions(iion,ixyz)
                  end do
               end do

               ! Subtract velocity_com from all ions
               do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
                  do ixyz = 1, 3
                     velocity_ions(iion,ixyz) = velocity_ions(iion,ixyz) - velocity_com(ixyz)
                  end do
               end do
            end if
         end if
      end do

      ! Compute kinetic energy of the system and actual temperature of the system
      ke_total = 0.0_wp
      do itype = 1, num_ion_types
         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then
            mi = ions(itype)%mass
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               visq = 0.0_wp
               do ixyz = 1, 3
                  visq = visq + velocity_ions(iion,ixyz) * velocity_ions(iion,ixyz)
               end do
               ke_total = ke_total + 0.5_wp * mi * visq
            end do
         end if
      end do

      ! ke = num_dof*k*T/2
      temperature_species = 2.0_wp * ke_total / (real(num_dof,wp) * boltzmann)
      scale = sqrt(temperature/temperature_species)

      ! Rescale velocities
      do itype = 1, num_ion_types
         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_ions(iion,ixyz) = scale * velocity_ions(iion, ixyz)
               end do
            end do
         end if
      end do

   end subroutine molecules_global_scaling

   ! ================================================================================
   ! Rescale particules velocities by species
   ! The centers of mass for each molecule group have their velocities set to 0.
   ! The centers of mass for each species (not in a molecule) group have their velocities set to 0.
   ! Independent scaling factor is applied to each molecule type and independent particule type
   subroutine molecules_independent_scaling(temperature, num_dof, molecules, ions, velocity_ions)
      use MW_kinds, only: wp
      use MW_constants, only: boltzmann
      use MW_errors, only: &
            MW_errors_allocate_error => allocate_error, &
            MW_errors_deallocate_error => deallocate_error
      use MW_ion, only: MW_ion_t
      use MW_molecule, only: MW_molecule_t
      implicit none
      ! Parameters
      ! ----------
      real(wp), intent(in) :: temperature             !< Target temperature
      integer, intent(in) :: num_dof                  !< Number of degrees of freedom in the system
      type(MW_molecule_t), intent(in) :: molecules(:) !< molecules parameters
      type(MW_ion_t), intent(in) :: ions(:)           !< ions parameters
      real(wp), intent(inout) :: velocity_ions(:,:)   !< ions velocities
      ! Local
      ! -----
      integer :: iion, ixyz
      integer :: num_ion_types, itype
      integer :: num_molecules, imol, isite
      real(wp) :: mi, visq, ke_total
      real(wp) :: velocity_com(3)
      real(wp) :: mass_total, temperature_species, scale
      logical :: in_molecule
      integer :: num_mobile_ions, num_ions_species, num_dof_species
      real(wp) :: dof_scaled


      num_ion_types = size(ions,1)
      num_molecules = size(molecules,1)

      num_mobile_ions = 0
      do itype = 1, num_ion_types
         if (ions(itype)%mobile .and. (.not. ions(itype)%is_fourth_site)) then
            num_mobile_ions = num_mobile_ions + ions(itype)%count
         end if
      end do

      ! Loop through molecules
      ! Zero center of mass velocity for each molecule types
      do imol = 1, num_molecules

         mass_total = 0.0_wp
         velocity_com(:) = 0.0_wp
         do isite = 1, molecules(imol)%num_sites
            itype = molecules(imol)%sites(isite)
            mi = ions(itype)%mass
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_com(ixyz) = velocity_com(ixyz) + mi * velocity_ions(iion,ixyz)
               end do
               mass_total = mass_total + mi
            end do
         end do

         do ixyz = 1, 3
            velocity_com(ixyz) = velocity_com(ixyz) / mass_total
         end do

         ! 0 COM velocity for this molecule
         do isite = 1, molecules(imol)%num_sites
            itype = molecules(imol)%sites(isite)
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_ions(iion,ixyz) = velocity_ions(iion,ixyz) - velocity_com(ixyz)
               end do
            end do
         end do

         ! Compute kinetic energy of the sub-system for this molecule
         ! and actual temperature of the system
         ke_total = 0.0_wp
         do isite = 1, molecules(imol)%num_sites
            itype = molecules(imol)%sites(isite)
            mi = ions(itype)%mass
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               visq = 0.0_wp
               do ixyz = 1, 3
                  visq = visq + velocity_ions(iion,ixyz) * velocity_ions(iion,ixyz)
               end do
               ke_total = ke_total + 0.5_wp * mi * visq
            end do
         end do

         num_ions_species = molecules(imol)%num_sites * molecules(imol)%n
         num_dof_species = 3 * num_ions_species - molecules(imol)%num_constraints
         dof_scaled = real(num_dof,wp) * (real(num_dof_species,wp) / real(3*num_mobile_ions,wp))
         temperature_species = 2.0_wp * ke_total / (dof_scaled * boltzmann)
         scale = sqrt(temperature/temperature_species)

         do isite = 1, molecules(imol)%num_sites
            itype = molecules(imol)%sites(isite)
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_ions(iion,ixyz) = scale * velocity_ions(iion, ixyz)
               end do
            end do
         end do
      end do


      ! Zero center of mass velocity and scale temperature for each
      ! ion species which are not in a molecule
      do itype = 1, num_ion_types
         in_molecule = .false.
         InMoleculeSearchLoop: do imol = 1, num_molecules
            do isite = 1, molecules(imol)%num_sites
               if (itype == molecules(imol)%sites(isite)) then
                  in_molecule = .true.
                  exit InMoleculeSearchLoop
               end if
            end do
         end do InMoleculeSearchLoop
         if (.not. in_molecule) then
            velocity_com(:) = 0.0_wp
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_com(ixyz) = velocity_com(ixyz) + velocity_ions(iion,ixyz)
               end do
            end do

            ! Subtract velocity_com from all ions
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_ions(iion,ixyz) = velocity_ions(iion,ixyz) - velocity_com(ixyz)
               end do
            end do

            ! Compute kinetic energy for this species
            mi = ions(itype)%mass
            ke_total = 0.0_wp
            do iion = ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               visq = 0.0_wp
               do ixyz = 1, 3
                  visq = visq + velocity_ions(iion,ixyz) * velocity_ions(iion,ixyz)
               end do
               ke_total = ke_total + 0.5_wp * mi * visq
            end do

            num_ions_species = ions(itype)%count
            num_dof_species = 3 * num_ions_species
            dof_scaled = real(num_dof,wp) * (real(num_dof_species,wp) / real(3*num_mobile_ions,wp))
            temperature_species = 2.0_wp * ke_total / (dof_scaled * boltzmann)
            scale = sqrt(temperature/temperature_species)

            ! Rescale velocities
            do iion = 1, ions(itype)%offset+1, ions(itype)%offset+ions(itype)%count
               do ixyz = 1, 3
                  velocity_ions(iion,ixyz) = scale * velocity_ions(iion, ixyz)
               end do
            end do

         end if
      end do

   end subroutine molecules_independent_scaling


end module MW_velocity
