! Module to implement NPT runs
!
! Reference
! ---------
! Constant pressure molecular dynamics algorithms
! Glenn K. Martyna, Douglas J. Tobias, Michael L. Klein, 
! J. Chem. Phys. 101 (5), 4177, 1994
!
module MW_barostat
   use MW_kinds, only: wp
   use MW_thermostat, only: MW_thermostat_t
   use MW_ion, only: MW_ion_t
   use MW_constants, only: boltzmann
   use MW_molecule, only: MW_molecule_t
   use MW_box, only: MW_box_t
   implicit none
   private

   ! Public types
   ! ------------
   public :: MW_barostat_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_barostat_t
      real(wp) :: temperature            !< Thermostat temperature
      real(wp) :: externalpressure       !< External/applied pressure
      real(wp) :: ions_th_relaxation_time  !< Thermostat relaxation time
      real(wp) :: b_relaxation_time     !< Barostat relaxation time
      integer  :: max_iteration          !< maximum number of iteration for velocity convergence
      real(wp) :: tolerance              !< tolerance on velocity convergence

      integer :: ions_th_chain_length = 0       !< Number of thermostats in the chain
      integer :: b_th_chain_length = 0        !< Number of thermostats in the chain for the barostat

      real(wp), allocatable :: ions_th_mass(:)        !< Thermostats mass
      real(wp), allocatable :: ions_th_zeta(:,:)      !< Thermostat "position"
      real(wp), allocatable :: ions_th_v_zeta(:,:)    !< Thermostat "velocity"
      real(wp), allocatable :: ions_th_G_zeta(:,:)    !< Thermostat "force"

      real(wp) :: b_W                            !< Barostat mass
      real(wp), allocatable :: b_epsilon(:)      !< Barostat "position"
      real(wp), allocatable :: b_v_epsilon(:)    !< Barostat "velocity"
      real(wp), allocatable :: b_F_epsilon(:)    !< Barostat "force"

      real(wp), allocatable :: b_th_mass(:)        !< Barostat thermostat masses
      real(wp), allocatable :: b_th_zeta(:,:)      !< Barostat thermostat "position"
      real(wp), allocatable :: b_th_v_zeta(:,:)    !< Barostat thermostat "velocity"
      real(wp), allocatable :: b_th_G_zeta(:,:)    !< Barostat 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_barostat_t

contains

   ! ================================================================================
   ! Define a barostat structure
   subroutine define_type(this, thermostat, externalpressure, &
                   b_relaxation_time, b_th_chain_length)
      use MW_errors, only: &
            MW_errors_allocate_error => allocate_error, &
            MW_errors_parameter_error => parameter_error
      implicit none

      ! Parameters in
      ! -------------
      type(MW_thermostat_t), intent(in) :: thermostat    !< Thermostat structure
      real(wp), intent(in) :: externalpressure   !< Barostat pressure
      real(wp), intent(in) :: b_relaxation_time !< Barostat relaxation time
      integer,  intent(in) :: b_th_chain_length  !< Number of thermostats in the chain for the barostat

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

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

!     if (ions_th_chain_length < 1) then
!        call MW_errors_parameter_error("define_type", "barostat.f90", "ions_th_chain_length", ions_th_chain_length)
!     end if

      if (b_th_chain_length < 1) then
         call MW_errors_parameter_error("define_type", "barostat.f90", "b_th_chain_length", b_th_chain_length)
      end if

      this%temperature = thermostat%temperature
      this%externalpressure = externalpressure
      this%ions_th_relaxation_time = thermostat%relaxation_time
      this%b_relaxation_time = b_relaxation_time
      this%max_iteration = thermostat%max_iteration
      this%tolerance = thermostat%tolerance
      this%ions_th_chain_length = thermostat%chain_length
      this%b_th_chain_length = b_th_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%ions_th_zeta(this%ions_th_chain_length+1, ZETA_NUM_STEPS), &
            this%ions_th_v_zeta(this%ions_th_chain_length+1, ZETA_NUM_STEPS), &
            this%ions_th_G_zeta(this%ions_th_chain_length+1, ZETA_NUM_STEPS), &
            this%ions_th_mass(this%ions_th_chain_length+1), &
            stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("define_type", "barostat.f90", ierr)
      end if

      allocate(this%b_epsilon(ZETA_NUM_STEPS), &
            this%b_v_epsilon(ZETA_NUM_STEPS), &
            this%b_F_epsilon(ZETA_NUM_STEPS), &
            stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("define_type", "barostat.f90", ierr)
      end if

      allocate(this%b_th_zeta(b_th_chain_length+1, ZETA_NUM_STEPS), &
            this%b_th_v_zeta(b_th_chain_length+1, ZETA_NUM_STEPS), &
            this%b_th_G_zeta(b_th_chain_length+1, ZETA_NUM_STEPS), &
            this%b_th_mass(b_th_chain_length+1), &
            stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("define_type", "barostat.f90", ierr)
      end if

      this%ions_th_zeta(:,:) = 0.0_wp
      this%ions_th_v_zeta(:,:) = 0.0_wp
      this%ions_th_G_zeta(:,:) = 0.0_wp
      this%ions_th_mass(:) = 0.0_wp

      this%b_epsilon(:) = 0.0_wp
      this%b_v_epsilon(:) = 0.0_wp
      this%b_F_epsilon(:) = 0.0_wp
      this%b_W = 0.0_wp

      this%b_th_zeta(:,:) = 0.0_wp
      this%b_th_v_zeta(:,:) = 0.0_wp
      this%b_th_G_zeta(:,:) = 0.0_wp
      this%b_th_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_barostat_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", "barostat.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", "barostat.f90", ierr)
      end if
      this%velocity_ions_work(:,:,:) = 0.0_wp
    end subroutine allocate_velocity_work_array

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

      ! Local
      ! -----
      real(wp) :: kbt
      real(wp) :: ions_th_relaxation_time_sq
      real(wp) :: b_relaxation_time_sq
      integer :: i

      kbt = boltzmann * this%temperature
      ions_th_relaxation_time_sq = this%ions_th_relaxation_time * this%ions_th_relaxation_time
      b_relaxation_time_sq = this%b_relaxation_time * this%b_relaxation_time

      this%ions_th_mass(1) = num_dof * kbt * ions_th_relaxation_time_sq
      do i = 2, this%ions_th_chain_length
         this%ions_th_mass(i) = kbt * ions_th_relaxation_time_sq
      end do

      this%b_W = (num_dof + 3) * kbt * b_relaxation_time_sq

      do i = 1, this%b_th_chain_length
         this%b_th_mass(i) = kbt * ions_th_relaxation_time_sq     ! Le ions ici n'est pas une erreur 
      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_barostat_t), intent(inout) :: this

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

      this%temperature = 0.0_wp
      this%externalpressure = 0.0_wp
      this%ions_th_relaxation_time = 0.0_wp
      this%b_relaxation_time = 0.0_wp
      this%b_W = 0.0_wp
      this%ions_th_chain_length = 0
      this%b_th_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", "barostat.f90", ierr)
         end if
      end if
      if (allocated(this%ions_th_zeta)) then
         deallocate(this%ions_th_zeta, this%ions_th_v_zeta, &
                    this%ions_th_G_zeta, this%ions_th_mass, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("define_type", "barostat.f90", ierr)
         end if
      end if
      if (allocated(this%b_th_zeta)) then
         deallocate(this%b_th_zeta, this%b_th_v_zeta, &
                    this%b_th_G_zeta, this%b_th_mass, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("define_type", "barostat.f90", ierr)
         end if
      end if
      if (allocated(this%b_epsilon)) then
         deallocate(this%b_epsilon, this%b_v_epsilon, this%b_F_epsilon, stat=ierr)
         if (ierr /= 0) then
            call MW_errors_deallocate_error("define_type", "barostat.f90", ierr)
         end if
      end if

   end subroutine void_type

   ! ================================================================================
   ! Print the thermostat structure
   subroutine print_type(this, ounit)
      implicit none
      ! Parameters
      ! ----------
      type(MW_barostat_t), intent(in) :: this
      integer, intent(in) :: ounit

      write(ounit, '("|thermostat| temperature:                           ",es12.5)') this%temperature
      write(ounit, '("|thermostat| relaxation time:                       ",es12.5)') &
                                           this%ions_th_relaxation_time
      write(ounit, '("|thermostat| chain length:                          ",i12)') &
                                           this%ions_th_chain_length
      if (allocated(this%ions_th_mass)) then
         write(ounit, '("|thermostat| masses: ",100(es12.5,1x))') this%ions_th_mass(:)
      end if
      write(ounit, '("|barostat| external pressure:                     ",es12.5)') this%externalpressure
      write(ounit, '("|barostat| relaxation time:                       ",es12.5)') this%b_relaxation_time
      write(ounit, '("|barostat| barostat mass: ",(es12.5,1x))') this%b_W
      if (allocated(this%b_th_mass)) then
         write(ounit, '("|barostat| barostat thermostat masses: ",100(es12.5,1x))') this%b_th_mass(:)
      endif
      write(ounit, '("|thermostat/barostat| maximum number of velocity iterations: ",i12)') this%max_iteration
      write(ounit, '("|thermostat/barostat| tolerance on velocities:               ",es12.5)') this%tolerance

   end subroutine print_type

   ! ================================================================================
   !> Cycle history
   subroutine update_history(this)
      implicit none
      ! Parameters
      ! ----------
      type(MW_barostat_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, volume_now, configurational_pressure)
     implicit none
     ! parameters
     ! ----------
     type(MW_barostat_t), intent(inout) :: this   !< Barostat/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
     real(wp), intent(in) :: volume_now             !< volume of the box
     real(wp), intent(in) :: configurational_pressure  !< pressure of the box

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

     call setup_masses(this, num_dof)

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

     volfactor_now = 1.0_wp / volume_now

     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
     pressure_now = configurational_pressure + kinetic_energy * volfactor_now / 3.0_wp

     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 for ions
     if(this%ions_th_G_zeta(1,zeta_now).eq.0.0_wp) this%ions_th_G_zeta(1,zeta_now) = 1.0_wp/this%ions_th_mass(1) *  &
                               (kinetic_energy - num_dof * kbt)

     ! Chained thermostats for ions
     do i = 2, this%ions_th_chain_length
        if(this%ions_th_G_zeta(i,zeta_now).eq.0.0_wp) this%ions_th_G_zeta(i,zeta_now) = - 1.0_wp/this%ions_th_mass(i) * kbt
     end do

     ! Chained thermostats forthe barostat 
     do i = 1, this%b_th_chain_length
        if(this%b_th_G_zeta(i,zeta_now).eq.0.0_wp) this%b_th_G_zeta(i,zeta_now) = - 1.0_wp/this%b_th_mass(i) * kbt
     end do

     ! Barostat
     if(this%b_F_epsilon(zeta_now).eq.0.0_wp) &
             this%b_F_epsilon(zeta_now) = 3.0_wp * (volume_now * (pressure_now - this%externalpressure) + kinetic_energy / num_dof) 

   end subroutine initialize

   ! ================================================================================
   !> Update ions position
   subroutine ions_positions_update(this, num_dof, dt, ions, &
         velocity_ions_prev, xyz_ions_now, velocity_ions_now, volume_now)
      implicit none
      ! Parameters
      ! ----------
      type(MW_barostat_t), intent(inout) :: this           !< barostat/thermostat parameters
      integer, intent(in) :: num_dof                       !< number of degree of freedom in the system
      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
      real(wp), intent(inout) :: volume_now             !< volume at current time step

      ! Local
      ! -----
      integer :: i, itype, ichain
      real(wp) :: halfdt, halfdtsq
      real(wp) :: exponentialprefactor
      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)

      halfdt = 0.5_wp*dt
      halfdtsq = 0.5_wp*dt*dt


      ! Update positions and velocities with thermostat and barostat "forces" computed at previous time step
      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%ions_th_v_zeta(1,zeta_prev) - &
                  halfdt * (2.0_wp+3.0_wp/num_dof) * velocity_ions_prev(i,1)*this%b_v_epsilon(zeta_prev)
            velocity_ions_now(i,2) = velocity_ions_now(i,2) - &
                  halfdt * velocity_ions_prev(i,2)*this%ions_th_v_zeta(1,zeta_prev) - &
                  halfdt * (2.0_wp+3.0_wp/num_dof) * velocity_ions_prev(i,2)*this%b_v_epsilon(zeta_prev)
            velocity_ions_now(i,3) = velocity_ions_now(i,3) - &
                  halfdt * velocity_ions_prev(i,3)*this%ions_th_v_zeta(1,zeta_prev) - &
                  halfdt * (2.0_wp+3.0_wp/num_dof) * velocity_ions_prev(i,3)*this%b_v_epsilon(zeta_prev)

            xyz_ions_now(i,1) = xyz_ions_now(i,1) - &
                  halfdtsq * velocity_ions_prev(i,1)*this%ions_th_v_zeta(1,zeta_prev) - &
                  halfdtsq * (2.0_wp+3.0_wp/num_dof) * velocity_ions_prev(i,1)*this%b_v_epsilon(zeta_prev)
            xyz_ions_now(i,2) = xyz_ions_now(i,2) - &
                  halfdtsq * velocity_ions_prev(i,2)*this%ions_th_v_zeta(1,zeta_prev) - &
                  halfdtsq * (2.0_wp+3.0_wp/num_dof) * velocity_ions_prev(i,2)*this%b_v_epsilon(zeta_prev)
            xyz_ions_now(i,3) = xyz_ions_now(i,3) - &
                  halfdtsq * velocity_ions_prev(i,3)*this%ions_th_v_zeta(1,zeta_prev) - &
                  halfdtsq * (2.0_wp+3.0_wp/num_dof) * velocity_ions_prev(i,3)*this%b_v_epsilon(zeta_prev)

         end do
      end do

      ! Update ions zetas 
      do ichain = 1, this%ions_th_chain_length
         this%ions_th_zeta(ichain, zeta_now) = this%ions_th_zeta(ichain, zeta_prev) + &
              dt * this%ions_th_v_zeta(ichain, zeta_prev) + &
              halfdtsq * this%ions_th_G_zeta(ichain, zeta_prev)
      end do

      ! Update barostat zetas 
      do ichain = 1, this%b_th_chain_length
         this%b_th_zeta(ichain, zeta_now) = this%b_th_zeta(ichain, zeta_prev) + &
              dt * this%b_th_v_zeta(ichain, zeta_prev) + &
              halfdtsq * this%b_th_G_zeta(ichain, zeta_prev)
      end do
      ! Update epsilon 
      this%b_epsilon(zeta_now) = this%b_epsilon(zeta_prev) + this%b_v_epsilon(zeta_prev) * dt &
                  + (this%b_F_epsilon(zeta_prev) / this%b_W  &
                     - this%b_v_epsilon(zeta_prev) * this%b_th_v_zeta(1,zeta_prev)) * halfdtsq
      ! Update positions and velocities with the new volume
      exponentialprefactor = exp(this%b_epsilon(zeta_now) - this%b_epsilon(zeta_prev))

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

            velocity_ions_now(i,1) = exponentialprefactor * velocity_ions_now(i,1) 
            velocity_ions_now(i,2) = exponentialprefactor * velocity_ions_now(i,2) 
            velocity_ions_now(i,3) = exponentialprefactor * velocity_ions_now(i,3) 

            xyz_ions_now(i,1) = exponentialprefactor * xyz_ions_now(i,1) 
            xyz_ions_now(i,2) = exponentialprefactor * xyz_ions_now(i,2) 
            xyz_ions_now(i,3) = exponentialprefactor * xyz_ions_now(i,3) 

         end do
      end do      

      ! Update volume
      volume_now =  volume_now * exponentialprefactor**3

   end subroutine ions_positions_update

   ! ================================================================================
   !> Update ions position
   subroutine ions_velocities_update(this, num_dof, dt, ions, velocity_ions_now,configurational_pressure,volume_now) 
!             num_pbc, molecule, box, xyz_next,num_molecule_types)
      use MW_errors, only: MW_errors_runtime_error => runtime_error
!     use MW_rattle, only: MW_rattle_constrain_velocities => constrain_velocities
!     use MW_molecule, only: CONSTRAINT_ALGORITHM_RATTLE
      implicit none
      ! Parameters
      ! ----------
      type(MW_barostat_t), intent(inout) :: this        !< barostat/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
      real(wp), intent(in) :: configurational_pressure    !< configurational contribution to the pressure
      real(wp), intent(in) :: volume_now                 !< volume of the box at current time step
!     integer, intent(in) :: num_pbc
!     type(MW_molecule_t), intent(in) :: molecule(:)
!     type(MW_box_t), intent(in) :: box
!     real(wp), intent(in) :: xyz_next(:,:)
!     integer, intent(in) :: num_molecule_types

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

      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
      volfactor_now = 1.0_wp / volume_now

      ! Initial guess for ions_th_v_zeta and b_th_v_zeta
      do i = 1, this%ions_th_chain_length
         this%ions_th_v_zeta(i,zeta_now) = this%ions_th_v_zeta(i,zeta_old) + 2.0_wp * dt * this%ions_th_G_zeta(i,zeta_prev)
      end do

      do i = 1, this%b_th_chain_length
         this%b_th_v_zeta(i,zeta_now) = this%b_th_v_zeta(i,zeta_old) + 2.0_wp * dt * this%b_th_G_zeta(i,zeta_prev)
      end do

      ! Initial guess for b_v_epsilon
      this%b_v_epsilon(zeta_now) = this%b_v_epsilon(zeta_old) + 2.0_wp * dt * this%b_F_epsilon(zeta_prev)/this%b_W

      ! 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

         th_scalingfactor = 1.0_wp / (1.0_wp + halfdt*(this%ions_th_v_zeta(1,zeta_now) &
                                      + (2.0_wp + 3.0_wp / num_dof)*this%b_v_epsilon(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) = th_scalingfactor * velocity_ions_now(i,1)
               this%velocity_ions_work(i,2,icur) = th_scalingfactor * velocity_ions_now(i,2)
               this%velocity_ions_work(i,3,icur) = th_scalingfactor * velocity_ions_now(i,3)

            end do
         end do

!        do itype = 1, num_molecule_types
!           if(molecule(itype)%constraint_algorithm == CONSTRAINT_ALGORITHM_RATTLE) then
!               max_iteration = molecule(itype)%constraint_max_iteration
!               tolerance = molecule(itype)%constraint_tolerance    
!               call MW_rattle_constrain_velocities(num_PBC,molecule(itype), &
!                   box, ions, &
!                    xyz_next(:,:), this%velocity_ions_work(:,:,icur))
!           endif
!        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

         pressure_now = configurational_pressure + kinetic_energy * volfactor_now / 3.0_wp

         ! Chained thermostats for ions 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%ions_th_G_zeta(1,zeta_now) = 1.0_wp/this%ions_th_mass(1) * (kinetic_energy - num_dof * kbt) - &
               this%ions_th_v_zeta(1, zeta_now)*this%ions_th_v_zeta(2, zeta_now)
         this%ions_th_v_zeta(1,zeta_now) = this%ions_th_v_zeta(1,zeta_prev) + halfdt * (this%ions_th_G_zeta(1,zeta_prev) + &
                                                                                        this%ions_th_G_zeta(1,zeta_now))
         ! Other thermostats
         do i = 2, this%ions_th_chain_length
            this%ions_th_G_zeta(i,zeta_now) = &
                  1.0_wp/this%ions_th_mass(i) * &
                  (this%ions_th_mass(i-1)*this%ions_th_v_zeta(i-1,zeta_now)*this%ions_th_v_zeta(i-1,zeta_now) - kbt) - &
                  this%ions_th_v_zeta(i, zeta_now)*this%ions_th_v_zeta(i+1, zeta_now)
            this%ions_th_v_zeta(i,zeta_now) = this%ions_th_v_zeta(i,zeta_prev) + &
                                              halfdt * (this%ions_th_G_zeta(i,zeta_prev) + this%ions_th_G_zeta(i,zeta_now))
         end do

         ! barostat force and velocity update
!        this%b_F_epsilon(zeta_now) = 3.0_wp * (volume_now * (pressure_now - this%externalpressure) + kinetic_energy / num_dof)
!        this%b_v_epsilon(zeta_now) = this%b_v_epsilon(zeta_prev) + &
!                       halfdt * (this%b_F_epsilon(zeta_prev) / this%b_W - &
!                                 this%b_v_epsilon(zeta_prev)*this%b_th_v_zeta(1,zeta_prev))+ &
!                       halfdt * (this%b_F_epsilon(zeta_now) / this%b_W  &
!                                - this%b_v_epsilon(zeta_now)*this%b_th_v_zeta(1,zeta_now)) 
         this%b_F_epsilon(zeta_now) = 3.0_wp * (volume_now * (pressure_now - this%externalpressure) + kinetic_energy / num_dof) &
                                 - this%b_W * this%b_v_epsilon(zeta_now)*this%b_th_v_zeta(1,zeta_now)
         this%b_v_epsilon(zeta_now) = this%b_v_epsilon(zeta_prev) + &
                        halfdt * (this%b_F_epsilon(zeta_prev) + this%b_F_epsilon(zeta_now)) / this%b_W  



         ! Chained thermostats for barostat 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%b_th_G_zeta(1,zeta_now) = 1.0_wp/this%b_th_mass(1) * (this%b_W * this%b_v_epsilon(zeta_now)**2 - kbt) - &
               this%b_th_v_zeta(1, zeta_now)*this%b_th_v_zeta(2, zeta_now)
         this%b_th_v_zeta(1,zeta_now) = this%b_th_v_zeta(1,zeta_prev) + halfdt * (this%b_th_G_zeta(1,zeta_prev) + &
                                                                                        this%b_th_G_zeta(1,zeta_now))
         ! Other thermostats
         do i = 2, this%b_th_chain_length
            this%b_th_G_zeta(i,zeta_now) = &
                  1.0_wp/this%b_th_mass(i) * &
                  (this%b_th_mass(i-1)*this%b_th_v_zeta(i-1,zeta_now)*this%b_th_v_zeta(i-1,zeta_now) - kbt) - &
                  this%b_th_v_zeta(i, zeta_now)*this%b_th_v_zeta(i+1, zeta_now)
            this%b_th_v_zeta(i,zeta_now) = this%b_th_v_zeta(i,zeta_prev) + &
                                              halfdt * (this%b_th_G_zeta(i,zeta_prev) + this%b_th_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", "barostat.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 barostat
   subroutine energy(this, num_dof, volume, h)
      implicit none
      ! Parameters
      ! ----------
      type(MW_barostat_t), intent(in) :: this
      integer, intent(in) :: num_dof
      real(wp), intent(in) :: volume
      real(wp), intent(inout) :: h

      ! Locals
      ! ------
      integer :: zeta_now
      integer :: i
      real(wp) :: ions_th_kinetic_energy, ions_th_potential_energy
      real(wp) :: b_th_kinetic_energy, b_th_potential_energy
      real(wp) :: b_kinetic_energy, b_potential_energy
      real(wp) :: PV
      real(wp) :: kbt

      zeta_now = this%zeta_step(1)
      kbt = boltzmann * this%temperature

      ions_th_kinetic_energy = 0.0_wp
      do i = 1, this%ions_th_chain_length
         ions_th_kinetic_energy = ions_th_kinetic_energy + &
                  0.5_wp * this%ions_th_mass(i) * this%ions_th_v_zeta(i, zeta_now) * this%ions_th_v_zeta(i, zeta_now)
      end do

      ions_th_potential_energy = num_dof * kbt * this%ions_th_zeta(1,zeta_now)
      do i = 2, this%ions_th_chain_length
         ions_th_potential_energy = ions_th_potential_energy + kbt * this%ions_th_zeta(i, zeta_now)
      end do

      b_kinetic_energy = 0.5_wp * this%b_W * this%b_v_epsilon(zeta_now) * this%b_v_epsilon(zeta_now)

      b_potential_energy = kbt * this%b_epsilon(zeta_now)

      b_th_kinetic_energy = 0.0_wp
      do i = 1, this%b_th_chain_length
         b_th_kinetic_energy = b_th_kinetic_energy + &
                  0.5_wp * this%b_th_mass(i) * this%b_th_v_zeta(i, zeta_now) * this%b_th_v_zeta(i, zeta_now)
      end do

      b_th_potential_energy = 0.0_wp
      do i = 1, this%b_th_chain_length
         b_th_potential_energy = b_th_potential_energy + kbt * this%b_th_zeta(i, zeta_now)
      end do

      PV =  this%externalpressure * volume 

      h = ions_th_kinetic_energy + ions_th_potential_energy + &
          b_kinetic_energy + b_potential_energy    + &
          b_th_kinetic_energy + b_th_potential_energy + &
          PV   
   end subroutine energy

end module MW_barostat
