! ===========================================================================
! Rattle algorithm part 1 - Update positions
subroutine constrain_positions_@PBC@(molecule, dt, box, ions, &
      xyz_now, xyz_next, velocity_next,stress_tensor)
   implicit none
   ! Passed in
   ! ---------
   type(MW_molecule_t), intent(in) :: molecule
   real(wp), intent(in) :: dt ! time step
   type(MW_box_t), intent(in) :: box ! simulation box
   type(MW_ion_t), intent(in)      :: ions(:)

   real(wp), intent(in) :: xyz_now(:,:)       ! coordinates at t

   ! Passed in/out
   ! -------------
   real(wp), intent(inout) :: xyz_next(:,:) !< coodinates at t+dt
   real(wp), intent(inout) :: velocity_next(:,:) !< velocities at t+dt
   real(wp), intent(inout) :: stress_tensor(:,:) !< contribution of constraint force on stress tensor

   ! Local
   ! -----
   integer :: imol
   integer :: num_sites
   integer :: num_constraints

   integer :: k, iter
   logical :: satisfied
   integer :: ia, ib, iasite, ibsite, iatype, ibtype
   real(wp) :: a, b, c
   real(wp) :: dr, drsq, delta, tolerance
   real(wp) :: sx, sy, sz, snorm2
   real(wp) :: rijx, rijy, rijz, rijnorm2, sdotrij
   real(wp) :: iamassrec, ibmassrec
   real(wp) :: g
   integer :: max_iteration

   call MW_timers_start(TIMER_RATTLE_POSITIONS)

   a = box%length(1)
   b = box%length(2)
   c = box%length(3)

   num_sites = molecule%num_sites
   num_constraints = molecule%num_constraints
   max_iteration = molecule%constraint_max_iteration
   tolerance = molecule%constraint_tolerance

   if (num_constraints > 0) then

      do imol = 1, molecule%n

         convergenceLoop: do iter = 1, max_iteration
            satisfied = .true.
            do k = 1, num_constraints

               iasite = molecule%constrained_sites(1,k)
               ibsite = molecule%constrained_sites(2,k)

               iatype = molecule%sites(iasite)
               ibtype = molecule%sites(ibsite)

               ia = ions(iatype)%offset + imol
               ib = ions(ibtype)%offset + imol

               call minimum_image_displacement_@PBC@(a, b, c, &
                     xyz_next(ia,1), xyz_next(ia,2), xyz_next(ia,3), &
                     xyz_next(ib,1), xyz_next(ib,2), xyz_next(ib,3), &
                     sx, sy, sz, snorm2)

               ! Test for constraint realization
               dr = molecule%constrained_dr(k)
               drsq = dr*dr
               delta = snorm2 - drsq
               if (abs(delta) > tolerance*drsq) then
                  satisfied = .false.

                  iamassrec = 1.0_wp / ions(iatype)%mass
                  ibmassrec = 1.0_wp / ions(ibtype)%mass

                  ! Compute minimum image distance between 2 points at the previous time step
                  call minimum_image_displacement_@PBC@(a, b, c, &
                        xyz_now(ia,1), xyz_now(ia,2), xyz_now(ia,3), &
                        xyz_now(ib,1), xyz_now(ib,2), xyz_now(ib,3), &
                        rijx, rijy, rijz, rijnorm2)
                  sdotrij = sx*rijx + sy*rijy + sz*rijz
                  ! this is g as defined in Andersen's paper (J. of Comp. Chem., 52, 24-34)
                  g = delta / (2.0_wp * dt* sdotrij * (iamassrec + ibmassrec))

                  xyz_next(ia,1) = xyz_next(ia,1) + dt*g*iamassrec*rijx
                  xyz_next(ia,2) = xyz_next(ia,2) + dt*g*iamassrec*rijy
                  xyz_next(ia,3) = xyz_next(ia,3) + dt*g*iamassrec*rijz

                  xyz_next(ib,1) = xyz_next(ib,1) - dt*g*ibmassrec*rijx
                  xyz_next(ib,2) = xyz_next(ib,2) - dt*g*ibmassrec*rijy
                  xyz_next(ib,3) = xyz_next(ib,3) - dt*g*ibmassrec*rijz

                  stress_tensor(1,1) = stress_tensor(1,1)-2.0_wp*g*rijx*rijx/dt
                  stress_tensor(2,2) = stress_tensor(2,2)-2.0_wp*g*rijy*rijy/dt
                  stress_tensor(3,3) = stress_tensor(3,3)-2.0_wp*g*rijz*rijz/dt
                  stress_tensor(1,2) = stress_tensor(1,2)-2.0_wp*g*rijx*rijy/dt
                  stress_tensor(1,3) = stress_tensor(1,3)-2.0_wp*g*rijx*rijz/dt
                  stress_tensor(2,3) = stress_tensor(2,3)-2.0_wp*g*rijy*rijz/dt

                  velocity_next(ia,1) = velocity_next(ia,1) + g*iamassrec*rijx
                  velocity_next(ia,2) = velocity_next(ia,2) + g*iamassrec*rijy
                  velocity_next(ia,3) = velocity_next(ia,3) + g*iamassrec*rijz

                  velocity_next(ib,1) = velocity_next(ib,1) - g*ibmassrec*rijx
                  velocity_next(ib,2) = velocity_next(ib,2) - g*ibmassrec*rijy
                  velocity_next(ib,3) = velocity_next(ib,3) - g*ibmassrec*rijz

               end if
            end do

            ! Check if all constrained are realized and finalize
            if (satisfied) then
               exit convergenceLoop
            end if

         end do convergenceLoop
      end do
   end if

   call MW_timers_stop(TIMER_RATTLE_POSITIONS)
end subroutine constrain_positions_@PBC@

! ===========================================================================
! Rattle algorithm part 2 - Update velocities
subroutine constrain_velocities_@PBC@(molecule, box, ions, &
      xyz_next, velocity_next)
   implicit none
   ! Passed in
   ! ---------
   type(MW_box_t), intent(in) :: box

   type(MW_molecule_t), intent(in) :: molecule
   type(MW_ion_t), intent(in) :: ions(:)
   real(wp), intent(in) :: xyz_next(:,:) ! coordinates at t + dt

   ! Passed out
   ! ----------
   real(wp), intent(inout) :: velocity_next(:,:)   ! coordinates at t+dt

   ! Local
   ! -----
   integer :: imol
   integer :: num_sites, num_constraints

   integer :: k, iter
   logical :: satisfied
   integer :: ia, ib, iasite, ibsite, iatype, ibtype
   real(wp) :: a, b, c
   real(wp) :: rijx, rijy, rijz, rijnorm2
   real(wp) :: vijx, vijy, vijz
   real(wp) :: dr, rdotvij, g
   real(wp) :: iamassrec, ibmassrec
   integer :: max_iteration
   real(wp) :: tolerance

   call MW_timers_start(TIMER_RATTLE_VELOCITIES)
   a = box%length(1)
   b = box%length(2)
   c = box%length(3)

   num_sites = molecule%num_sites
   num_constraints = molecule%num_constraints
   max_iteration = molecule%constraint_max_iteration
   tolerance = molecule%constraint_tolerance

   if (num_constraints > 0) then

      do imol = 1, molecule%n

         convergenceLoop: do iter = 1, max_iteration
            satisfied = .true.
            do k = 1, num_constraints

               iasite = molecule%constrained_sites(1,k)
               ibsite = molecule%constrained_sites(2,k)

               iatype = molecule%sites(iasite)
               ibtype = molecule%sites(ibsite)

               ia = ions(iatype)%offset + imol
               ib = ions(ibtype)%offset + imol

               call minimum_image_displacement_@PBC@(a, b, c, &
                     xyz_next(ia,1), xyz_next(ia,2), xyz_next(ia,3), &
                     xyz_next(ib,1), xyz_next(ib,2), xyz_next(ib,3), &
                     rijx, rijy, rijz, rijnorm2)

               vijx = velocity_next(ib,1) - velocity_next(ia,1)
               vijy = velocity_next(ib,2) - velocity_next(ia,2)
               vijz = velocity_next(ib,3) - velocity_next(ia,3)

               ! Test for constraint realization
               rdotvij = rijx*vijx + rijy*vijy + rijz*vijz
               if (abs(rdotvij) >= tolerance) then
                  satisfied = .false.

                  dr = molecule%constrained_dr(k)
                  iamassrec = 1.0_wp / ions(iatype)%mass
                  ibmassrec = 1.0_wp / ions(ibtype)%mass

                  ! this is k as defined in Andersen's paper (J. of Comp. Chem., 52, 24-34)
                  g = rdotvij / ((dr*dr) * (iamassrec + ibmassrec))

                  velocity_next(ia,1) = velocity_next(ia,1) + g*iamassrec*rijx
                  velocity_next(ia,2) = velocity_next(ia,2) + g*iamassrec*rijy
                  velocity_next(ia,3) = velocity_next(ia,3) + g*iamassrec*rijz

                  velocity_next(ib,1) = velocity_next(ib,1) - g*ibmassrec*rijx
                  velocity_next(ib,2) = velocity_next(ib,2) - g*ibmassrec*rijy
                  velocity_next(ib,3) = velocity_next(ib,3) - g*ibmassrec*rijz

               end if
            end do

            ! Check if all constrained are realized and finalize
            if (satisfied) then
               exit convergenceLoop
            end if

         end do convergenceLoop
      end do
   end if
   call MW_timers_stop(TIMER_RATTLE_VELOCITIES)
end subroutine constrain_velocities_@PBC@
