! Module to implement thermostatting routine for constant temperature runs
!
! Reference
! ---------
! Nose-Hoover chains: The canonical ensemble via continuous dynamics
! Glenn K. Martyna, Michael L. Klein, Marck Tuckerman
! J. Chem. Phys. 97 (4), 15 August 1992
!
module MW_thermostat
   use MW_kinds, only: wp
   use MW_ion, only: MW_ion_t
   use MW_constants, only: boltzmann
   implicit none
   private

   ! Public types
   ! ------------
   public :: MW_thermostat_t

   ! Public subroutines
   ! ------------------
   public :: define_type
   public :: allocate_velocity_work_array
   public :: void_type
   public :: print_type
   public :: update_history
   public :: initialize
   public :: ions_positions_update
   public :: ions_velocities_update
   public :: setup_masses
   public :: energy

   integer, parameter :: ZETA_NUM_STEPS = 3

   type MW_thermostat_t
      real(wp) :: temperature            !< Thermostat temperature
      real(wp) :: relaxation_time        !< Relaxation time
      integer  :: max_iteration          !< maximum number of iteration for velocity convergence
      real(wp) :: tolerance              !< tolerance on velocity convergence

      integer :: chain_length = 0        !< Number of thermostats in the chain

      real(wp), allocatable :: mass(:)        !< Thermostats mass
      real(wp), allocatable :: zeta(:,:)      !< Thermostat "position"
      real(wp), allocatable :: v_zeta(:,:)    !< Thermostat "velocity"
      real(wp), allocatable :: G_zeta(:,:)    !< Thermostat "force"

      ! history pointer
      integer :: zeta_step(ZETA_NUM_STEPS)

      real(wp), allocatable :: velocity_ions_work(:,:,:) !< work array required for velocity iterations
   end type MW_thermostat_t

contains

   ! ================================================================================
   ! Define a thermostat structure
   subroutine define_type(this, temperature, relaxation_time, &
         max_iteration, tolerance, chain_length)
      use MW_errors, only: &
            MW_errors_allocate_error => allocate_error, &
            MW_errors_parameter_error => parameter_error
      implicit none

      ! Parameters in
      ! -------------
      real(wp), intent(in) :: temperature      !< Thermostat temperature
      real(wp), intent(in) :: relaxation_time  !< Relaxation time
      integer,  intent(in) :: max_iteration    !< maximum number of iteration for velocity convergence
      real(wp), intent(in) :: tolerance        !< tolerance on velocity convergence
      integer,  intent(in) :: chain_length  !< Number of thermostats in the chain

      ! Parameters out
      ! --------------
      type(MW_thermostat_t), intent(inout) :: this

      ! Local
      ! -----
      integer :: ierr, step

      if (chain_length < 1) then
         call MW_errors_parameter_error("define_type", "thermostat.f90", "chain_length", chain_length)
      end if

      this%temperature = temperature
      this%relaxation_time = relaxation_time
      this%max_iteration = max_iteration
      this%tolerance = tolerance
      this%chain_length = chain_length

      ! allocate with size chain_length+1 in order to have a more generic formulation
      ! the n+1th thermostat will always have zeta=0, v_zeta=0 and G_zeta=0
      allocate(this%zeta(chain_length+1, ZETA_NUM_STEPS), &
            this%v_zeta(chain_length+1, ZETA_NUM_STEPS), &
            this%G_zeta(chain_length+1, ZETA_NUM_STEPS), &
            this%mass(chain_length+1), &
            stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("define_type", "thermostat.f90", ierr)
      end if

      this%zeta(:,:) = 0.0_wp
      this%v_zeta(:,:) = 0.0_wp
      this%G_zeta(:,:) = 0.0_wp
      this%mass(:) = 0.0_wp

      do step = 1, ZETA_NUM_STEPS
         this%zeta_step(step) = step
      end do

   end subroutine define_type

   ! ================================================================================
   !> Allocate velocity work array
   subroutine allocate_velocity_work_array(this, num_ions)
      use MW_errors, only: &
            MW_errors_allocate_error => allocate_error, &
            MW_errors_parameter_error => parameter_error
     implicit none
     ! Parameters
     ! ----------
     type(MW_thermostat_t), intent(inout) :: this
     integer, intent(in) :: num_ions !< size of the velocity work array
     ! local
     ! -----
     integer :: ierr

      if (num_ions < 1) then
         call MW_errors_parameter_error("define_type", "thermostat.f90", "num_ions", num_ions)
      end if
      allocate(this%velocity_ions_work(num_ions,3,2), stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("define_type", "thermostat.f90", ierr)
      end if
      this%velocity_ions_work(:,:,:) = 0.0_wp
    end subroutine allocate_velocity_work_array

   ! ================================================================================
   ! Setup thermostat masses
   subroutine setup_masses(this, num_dof)
      implicit none
      ! Parameters
      ! ----------
      type(MW_thermostat_t), intent(inout) :: this
      integer, intent(in) :: num_dof

      ! Local
      ! -----
      real(wp) :: kbt
      real(wp) :: relaxation_time_sq
      integer :: i

      kbt = boltzmann * this%temperature
      relaxation_time_sq = this%relaxation_time * this%relaxation_time

      this%mass(1) = num_dof * kbt * relaxation_time_sq
      do i = 2, this%chain_length
         this%mass(i) = kbt * relaxation_time_sq
      end do
   end subroutine setup_masses

   ! ================================================================================
   ! Void a thermostat structure
   subroutine void_type(this)
      use MW_errors, only: MW_errors_deallocate_error => deallocate_error
      implicit none
      ! Parameters
      ! ----------
      type(MW_thermostat_t), intent(inout) :: this

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

      this%temperature = 0.0_wp
      this%relaxation_time = 0.0_wp
      this%chain_length = 0
      this%tolerance = 0.0_wp
      this%zeta_step(:) = 0

      if (allocated(this%velocity_ions_work)) then
         deallocate(this%velocity_ions_work, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("define_type", "thermostat.f90", ierr)
         end if
      end if
      if (allocated(this%zeta)) then
         deallocate(this%zeta, this%v_zeta, this%G_zeta, this%mass, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("define_type", "thermostat.f90", ierr)
         end if
      end if

   end subroutine void_type

   ! ================================================================================
   ! Void a thermostat structure
   subroutine print_type(this, ounit)
      implicit none
      ! Parameters
      ! ----------
      type(MW_thermostat_t), intent(in) :: this
      integer, intent(in) :: ounit

      write(ounit, '("|thermostat| temperature:                           ",es12.5)') this%temperature
      write(ounit, '("|thermostat| relaxation time:                       ",es12.5)') this%relaxation_time
      write(ounit, '("|thermostat| chain length:                          ",i12)') this%chain_length
      write(ounit, '("|thermostat| maximum number of velocity iterations: ",i12)') this%max_iteration
      write(ounit, '("|thermostat| tolerance on velocities:               ",es12.5)') this%tolerance
      if (allocated(this%mass)) then
         write(ounit, '("|thermostat| masses: ",100(es12.5,1x))') this%mass(:)
      end if

   end subroutine print_type

   ! ================================================================================
   !> Cycle history
   subroutine update_history(this)
      implicit none
      ! Parameters
      ! ----------
      type(MW_thermostat_t), intent(inout) :: this

      ! Local
      ! -----
      integer :: step_now, step

      ! current step data will replace oldest step data
      step_now = this%zeta_step(ZETA_NUM_STEPS)
      do step = ZETA_NUM_STEPS, 2, -1
         this%zeta_step(step) = this%zeta_step(step-1)
      end do
      this%zeta_step(1) = step_now

   end subroutine update_history

   ! ================================================================================
   ! Initialize thermostat position, velocity and force
   subroutine initialize(this, num_dof, ions, velocity_ions_now)
     implicit none
     ! parameters
     ! ----------
     type(MW_thermostat_t), intent(inout) :: this   !< Thermostat parameters
     integer, intent(in) :: num_dof                 !< number of degree of freedom in the system
     type(MW_ion_t), intent(in) :: ions(:)          !< ion type parameters
     real(wp), intent(in) :: velocity_ions_now(:,:) !< ions velocities

     ! Local
     ! -----
     integer :: zeta_now
     integer :: i, itype
     integer :: num_ion_types
     real(wp) :: kbt, mass, vsq, kinetic_energy
     real(wp) :: vx, vy, vz

     call setup_masses(this, num_dof)

     num_ion_types = size(ions,1)
     kbt = boltzmann*this%temperature

     kinetic_energy = 0.0_wp
     do itype = 1, num_ion_types
        mass = ions(itype)%mass
        do i = ions(itype)%offset+1, ions(itype)%offset + ions(itype)%count
           vx = velocity_ions_now(i,1)
           vy = velocity_ions_now(i,2)
           vz = velocity_ions_now(i,3)
           vsq = vx*vx + vy*vy + vz*vz
           kinetic_energy = kinetic_energy + mass*vsq
        end do
     end do

     zeta_now = this%zeta_step(1)
!    do i = 1, this%chain_length
!       this%zeta(i,zeta_now) = 0.0_wp
!       this%v_zeta(i,zeta_now) = 0.0_wp
!    end do

     ! First thermostat in the chain
     if(this%G_zeta(1,zeta_now).eq.0.0_wp) this%G_zeta(1,zeta_now) = 1.0_wp/this%mass(1) * (kinetic_energy - num_dof*kbt)

     ! Chained thermostats
     do i = 2, this%chain_length
        if(this%G_zeta(i,zeta_now).eq.0.0_wp) this%G_zeta(i,zeta_now) = - 1.0_wp/this%mass(i) * kbt
     end do

   end subroutine initialize

   ! ================================================================================
   !> Update ions position
   subroutine ions_positions_update(this, dt, ions, &
         velocity_ions_prev, xyz_ions_now, velocity_ions_now)
      implicit none
      ! Parameters
      ! ----------
      type(MW_thermostat_t), intent(inout) :: this           !< thermostat parameters
      real(wp), intent(in) :: dt                          !< time step
      type(MW_ion_t), intent(in) :: ions(:)               !< ion_types data
      real(wp), intent(in) :: velocity_ions_prev(:,:)     !< ions velocities at previous time step

      real(wp), intent(inout) :: xyz_ions_now(:,:)        !< ions positions at current time step
      real(wp), intent(inout) :: velocity_ions_now(:,:) !< ions velocities at current time step

      ! Local
      ! -----
      integer :: i, itype, ichain
      real(wp) :: halfdt, halfdtsq
      integer :: zeta_now, zeta_prev
      integer :: num_ion_types

      num_ion_types = size(ions, 1)
      zeta_now = this%zeta_step(1)
      zeta_prev = this%zeta_step(2)

      ! Update positions and velocities with thermostat "forces" computed at previous time step
      halfdt = 0.5_wp*dt
      halfdtsq = 0.5_wp*dt*dt
      do itype = 1, num_ion_types
         do i = ions(itype)%offset+1, ions(itype)%offset + ions(itype)%count

            velocity_ions_now(i,1) = velocity_ions_now(i,1) - &
                  halfdt * velocity_ions_prev(i,1)*this%v_zeta(1,zeta_prev)
            velocity_ions_now(i,2) = velocity_ions_now(i,2) - &
                  halfdt * velocity_ions_prev(i,2)*this%v_zeta(1,zeta_prev)
            velocity_ions_now(i,3) = velocity_ions_now(i,3) - &
                  halfdt * velocity_ions_prev(i,3)*this%v_zeta(1,zeta_prev)

            xyz_ions_now(i,1) = xyz_ions_now(i,1) - &
                  halfdtsq * velocity_ions_prev(i,1)*this%v_zeta(1,zeta_prev)
            xyz_ions_now(i,2) = xyz_ions_now(i,2) - &
                  halfdtsq * velocity_ions_prev(i,2)*this%v_zeta(1,zeta_prev)
            xyz_ions_now(i,3) = xyz_ions_now(i,3) - &
                  halfdtsq * velocity_ions_prev(i,3)*this%v_zeta(1,zeta_prev)

         end do
      end do

      do ichain = 1, this%chain_length
         this%zeta(ichain, zeta_now) = this%zeta(ichain, zeta_prev) + &
              dt * this%v_zeta(ichain, zeta_prev) + &
              halfdtsq * this%G_zeta(ichain, zeta_prev)
      end do

   end subroutine ions_positions_update

   ! ================================================================================
   !> Update ions position
   subroutine ions_velocities_update(this, num_dof, dt, ions, velocity_ions_now)
      use MW_errors, only: MW_errors_runtime_error => runtime_error
      implicit none
      ! Parameters
      ! ----------
      type(MW_thermostat_t), intent(inout) :: this        !< thermostat parameters
      integer, intent(in) :: num_dof                      !< number of degrees of freedom
      real(wp), intent(in) :: dt                          !< time step
      type(MW_ion_t), intent(in) :: ions(:)               !< ions parameters
      real(wp), intent(inout) :: velocity_ions_now(:,:) !< ions velocities at current time step

      ! Local
      ! -----
      integer :: num_ion_types
      integer :: zeta_now, zeta_prev, zeta_old
      integer :: i, itype, iter, icur, iprev, itmp
      real(wp) :: mass, halfdt
      real(wp) :: scalingfactor
      real(wp) :: kbt, kinetic_energy, vsq
      logical :: converge
      real(wp) :: delta
      real(wp) :: vx, vy, vz

      num_ion_types = size(ions,1)
      zeta_now = this%zeta_step(1)
      zeta_prev = this%zeta_step(2)
      zeta_old = this%zeta_step(3)

      ! precompute some factors
      kbt = boltzmann*this%temperature
      halfdt = 0.5_wp*dt

      ! Initial guess for v_zeta
      do i = 1, this%chain_length
         this%v_zeta(i,zeta_now) = this%v_zeta(i,zeta_old) + 2.0_wp * dt * this%G_zeta(i,zeta_prev)
      end do

      ! Setup work array
      icur = 1
      iprev = 2
      this%velocity_ions_work(:,:,icur) = velocity_ions_now(:,:)

      converge = .false.
      convergenceLoop: do iter = 1, this%max_iteration

         ! swap indices for work arrays
         itmp = icur
         icur = iprev
         iprev = itmp

         scalingfactor = 1.0_wp / (1.0_wp + halfdt*this%v_zeta(1,zeta_now))

         ! Scale velocities
         do itype = 1, num_ion_types
            do i = ions(itype)%offset+1, ions(itype)%offset + ions(itype)%count

               this%velocity_ions_work(i,1,icur) = scalingfactor * velocity_ions_now(i,1)
               this%velocity_ions_work(i,2,icur) = scalingfactor * velocity_ions_now(i,2)
               this%velocity_ions_work(i,3,icur) = scalingfactor * velocity_ions_now(i,3)

            end do
         end do

         ! Update thermostat "force"
         kinetic_energy = 0.0_wp
         do itype = 1, num_ion_types
            mass = ions(itype)%mass
            do i = ions(itype)%offset+1, ions(itype)%offset + ions(itype)%count
               vx = this%velocity_ions_work(i, 1, icur)
               vy = this%velocity_ions_work(i, 2, icur)
               vz = this%velocity_ions_work(i, 3, icur)
               vsq = vx*vx + vy*vy + vz*vz
               kinetic_energy = kinetic_energy + mass*vsq
            end do
         end do


         ! Chained thermostats update
         ! Implementation feature: ilast=chain_length+1 is a valid index and v_zeta(ilast,:)=0 at all time
         !
         ! First thermostat in the chain
         this%G_zeta(1,zeta_now) = 1.0_wp/this%mass(1) * (kinetic_energy - num_dof*kbt) - &
               this%v_zeta(1, zeta_now)*this%v_zeta(2, zeta_now)
         this%v_zeta(1,zeta_now) = this%v_zeta(1,zeta_prev) + halfdt * (this%G_zeta(1,zeta_prev) + this%G_zeta(1,zeta_now))
         ! Other thermostats
         do i = 2, this%chain_length
            this%G_zeta(i,zeta_now) = &
                  1.0_wp/this%mass(i) * (this%mass(i-1)*this%v_zeta(i-1,zeta_now)*this%v_zeta(i-1,zeta_now) - kbt) - &
                  this%v_zeta(i, zeta_now)*this%v_zeta(i+1, zeta_now)
            this%v_zeta(i,zeta_now) = this%v_zeta(i,zeta_prev) + halfdt * (this%G_zeta(i,zeta_prev) + this%G_zeta(i,zeta_now))
         end do

         ! check for convergence
         delta = 0.0_wp
         do itype = 1, num_ion_types
            mass = ions(itype)%mass
            do i = ions(itype)%offset+1, ions(itype)%offset + ions(itype)%count
               vx = this%velocity_ions_work(i,1,icur) - this%velocity_ions_work(i,1,iprev)
               vy = this%velocity_ions_work(i,2,icur) - this%velocity_ions_work(i,2,iprev)
               vz = this%velocity_ions_work(i,3,icur) - this%velocity_ions_work(i,3,iprev)
               vsq = vx*vx + vy*vy + vz*vz
               delta = delta + mass*vsq
            end do
         end do

         if (delta <= (num_dof*kbt)*this%tolerance) then
            converge = .true.
            exit convergenceLoop
         end if

      end do convergenceLoop
      if (.not. converge) then
         call MW_errors_runtime_error("ions_velocities_update", "thermostat.f90", &
               "Velocity thermostating failed to converge")
      end if

      ! Copy back updated velocities
      velocity_ions_now(:,:) = this%velocity_ions_work(:,:,icur)

   end subroutine ions_velocities_update

   ! ================================================================================
   !> Compute the energy of the thermostat
   subroutine energy(this, num_dof, h)
      implicit none
      ! Parameters
      ! ----------
      type(MW_thermostat_t), intent(in) :: this
      integer, intent(in) :: num_dof
      real(wp), intent(inout) :: h

      ! Locals
      ! ------
      integer :: zeta_now
      integer :: i
      real(wp) :: kinetic_energy, potential_energy
      real(wp) :: kbt

      zeta_now = this%zeta_step(1)
      kinetic_energy = 0.0_wp
      do i = 1, this%chain_length
         kinetic_energy = kinetic_energy + 0.5_wp * this%mass(i) * this%v_zeta(i, zeta_now) * this%v_zeta(i, zeta_now)
      end do

      kbt = boltzmann * this%temperature
      potential_energy = num_dof * kbt * this%zeta(1,zeta_now)
      do i = 2, this%chain_length
         potential_energy = potential_energy + kbt * this%zeta(i, zeta_now)
      end do

      h = kinetic_energy + potential_energy
   end subroutine energy

end module MW_thermostat
