! Rattle
! ------
!
! Rattle: A "Velocity" Version of the Shake Algorithm for Molecular Dynamics Calculation
! Hans C. Andersen
! Journal of Computationanl Physics 52, 24-34, (1983)
!
! RATTLE is an algorithm for integrating the equations of motion in
! molecular dynamics calculations for molecular models with internal
! constraints. RATTLE calculates the positions and velocities at the
! next time step from the positions and celocities at the present time
! step, without requiring information about earlier history. It is
! based on the Verlet algorithm and retains the simplicity of using
! Cartesian coordinates for each of the atoms to describe the
! configuration of a molecule with internal constraints.
!
! Algorithm
!
! r(t+h) = r(t) + h*v(t) + h^2 * [ F(r(t)) + g_rr(t) ] / 2
! v(t+h) = v(t) + h* [ F(r(t)) + g_rr(t) + F(r(t+h)) + g_rv(t) ] / 2
!
! where g_rr and g_rv represent two separate approximations for the
! forces associated with the constraints such that both the positions
! and the velocities satisfy the constraints.
module MW_rattle
   use MW_timers
   use MW_kinds, only: wp
   use MW_box, only: MW_box_t, &
         MW_box_fix_outliers => fix_outliers
   use MW_ion, only: MW_ion_t
   use MW_molecule, only: MW_molecule_t
   implicit none
   private


   ! Public type
   ! -----------
   public :: MW_rattle_t

   ! Public subroutines
   ! ------------------
   public :: define_type
   public :: void_type
   public :: print_type
   public :: constrain_positions
   public :: constrain_velocities

   type MW_rattle_t
      integer :: max_iterations !< maximum iteration in rattle
      real(wp) :: tolerance     !< tolerance on constraints

      ! statistics
      integer :: max_positions_iterations_count = 0
      integer :: max_velocities_iterations_count = 0
   end type MW_rattle_t

contains

   ! ========================================================================
   ! Define the data structure
   subroutine define_type(this, max_iterations, tolerance)
      implicit none
      ! Passed in
      ! ---------
      type(MW_rattle_t), intent(inout) :: this
      integer, intent(in) :: max_iterations
      real(wp) ,intent(in) :: tolerance

      this%max_iterations = max_iterations
      this%tolerance = tolerance

      this%max_positions_iterations_count = 0
      this%max_velocities_iterations_count = 0

   end subroutine define_type

   ! ========================================================================
   ! Void the data structure
   subroutine void_type(this)
      implicit none
      ! Passed in
      ! ---------
      type(MW_rattle_t), intent(inout) :: this

      this%max_iterations = 0
      this%tolerance = 0

      this%max_positions_iterations_count = 0
      this%max_velocities_iterations_count = 0

   end subroutine void_type

   ! ===========================================================================
   ! Print info about the data structure
   subroutine print_type(this, ounit)
      implicit none
      type(MW_rattle_t), intent(in) :: this
      integer, intent(in) :: ounit

      write(ounit, '("|rattle| maximum number of iterations:         ", i12)') this%max_iterations
      write(ounit, '("|rattle| tolerance on constraints realisation: ", es12.5)') this%tolerance
   end subroutine print_type

   ! ===========================================================================
   ! Rattle algorithm part 1 - Update positions
   subroutine constrain_positions(num_pbc, molecule, dt, box, ions, &
         xyz_now, xyz_next, velocity_next,stress_tensor)
      implicit none
      ! Passed in
      ! ---------
      integer, intent(in) :: num_pbc
      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 to stress


      select case(num_pbc)
      case(2)
         call constrain_positions_2DPBC(molecule, dt, box, ions, &
         xyz_now, xyz_next, velocity_next,stress_tensor)
      case(3)
         call constrain_positions_3DPBC(molecule, dt, box, ions, &
         xyz_now, xyz_next, velocity_next,stress_tensor)
      end select

   end subroutine constrain_positions

   ! ===========================================================================
   ! Rattle algorithm part 2 - Update velocities
   subroutine constrain_velocities(num_pbc, molecule, box, ions, &
         xyz_next, velocity_next)
      implicit none
      ! Passed in
      ! ---------
      integer, intent(in) :: num_pbc
      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

      select case(num_pbc)
      case(2)
         call constrain_velocities_2DPBC(molecule, box, ions, &
         xyz_next, velocity_next)
      case(3)
         call constrain_velocities_3DPBC(molecule, box, ions, &
         xyz_next, velocity_next)
      end select

   end subroutine constrain_velocities

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

end module MW_rattle
