!--------------------------------------------------------------------------------------------------!
!   CP2K: A general program to perform molecular dynamics simulations                              !
!   Copyright (C) 2000 - 2018  CP2K developers group                                               !
!--------------------------------------------------------------------------------------------------!

! **************************************************************************************************
!> \brief Defines functions to perform rmsd in 3D
!> \author Teodoro Laino 09.2006
! **************************************************************************************************
MODULE rmsd

   USE kinds,                           ONLY: dp
   USE mathlib,                         ONLY: diamat_all,&
                                              matvec_3x3
   USE particle_types,                  ONLY: particle_type
#include "./base/base_uses.f90"

   IMPLICIT NONE
   PRIVATE

   CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'rmsd'
   REAL(KIND=dp), PARAMETER, PRIVATE    :: Epsi = EPSILON(0.0_dp)*1.0E4_dp

   PUBLIC :: rmsd3

CONTAINS

! **************************************************************************************************
!> \brief Computes the RMSD in 3D. Provides also derivatives.
!> \param particle_set ...
!> \param r ...
!> \param r0 ...
!> \param output_unit ...
!> \param weights ...
!> \param my_val ...
!> \param rotate ...
!> \param transl ...
!> \param rot ...
!> \param drmsd3 ...
!> \author Teodoro Laino 08.2006
!> \note
!>      Optional arguments:
!>          my_val -> gives back the value of the RMSD
!>          transl -> provides the translational vector
!>          rotate -> if .true. gives back in r the coordinates rotated
!>                    in order to minimize the RMSD3
!>          rot    -> provides the rotational matrix
!>          drmsd3 -> derivatives of RMSD3 w.r.t. atomic positions
! **************************************************************************************************
   SUBROUTINE rmsd3(particle_set, r, r0, output_unit, weights, my_val, &
                    rotate, transl, rot, drmsd3)
      TYPE(particle_type), DIMENSION(:), POINTER         :: particle_set
      REAL(KIND=dp), DIMENSION(:), INTENT(INOUT)         :: r, r0
      INTEGER, INTENT(IN)                                :: output_unit
      REAL(KIND=dp), DIMENSION(:), INTENT(IN), OPTIONAL  :: weights
      REAL(KIND=dp), INTENT(OUT), OPTIONAL               :: my_val
      LOGICAL, INTENT(IN), OPTIONAL                      :: rotate
      REAL(KIND=dp), DIMENSION(:), INTENT(OUT), OPTIONAL :: transl
      REAL(KIND=dp), DIMENSION(:, :), INTENT(INOUT), &
         OPTIONAL                                        :: rot, drmsd3

      CHARACTER(len=*), PARAMETER :: routineN = 'rmsd3', routineP = moduleN//':'//routineN

      INTEGER                                            :: i, ix, j, k, natom
      LOGICAL                                            :: my_rotate
      REAL(KIND=dp)                                      :: dm_r(4, 4, 3), lambda(4), loc_tr(3), &
                                                            M(4, 4), mtot, my_rot(3, 3), Q(0:3), &
                                                            rr(3), rr0(3), rrsq, s, xx, yy, &
                                                            Z(4, 4), zz
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:)           :: w
      REAL(KIND=dp), ALLOCATABLE, DIMENSION(:, :)        :: r0p, rp

      CPASSERT(SIZE(r) == SIZE(r0))
      natom = SIZE(particle_set)
      my_rotate = .FALSE.
      IF (PRESENT(rotate)) my_rotate = rotate
      ALLOCATE (w(natom))
      IF (PRESENT(weights)) THEN
         w(:) = weights
      ELSE
         ! All atoms have a weight proportional to their mass in the RMSD unless
         ! differently requested
         DO i = 1, natom
            w(i) = particle_set(i)%atomic_kind%mass
         END DO
         mtot = MINVAL(w)
         IF (mtot /= 0.0_dp) w(:) = w(:)/mtot
      END IF
      ALLOCATE (rp(3, natom))
      ALLOCATE (r0p(3, natom))
      ! Molecule given by coordinates R
      ! Find COM and center molecule in COM
      xx = 0.0_dp
      yy = 0.0_dp
      zz = 0.0_dp
      mtot = 0.0_dp
      DO i = 1, natom
         mtot = mtot+particle_set(i)%atomic_kind%mass
         xx = xx+r((i-1)*3+1)*particle_set(i)%atomic_kind%mass
         yy = yy+r((i-1)*3+2)*particle_set(i)%atomic_kind%mass
         zz = zz+r((i-1)*3+3)*particle_set(i)%atomic_kind%mass
      END DO
      xx = xx/mtot
      yy = yy/mtot
      zz = zz/mtot
      DO i = 1, natom
         rp(1, i) = r((i-1)*3+1)-xx
         rp(2, i) = r((i-1)*3+2)-yy
         rp(3, i) = r((i-1)*3+3)-zz
      END DO
      IF (PRESENT(transl)) THEN
         transl(1) = xx
         transl(2) = yy
         transl(3) = zz
      END IF
      ! Molecule given by coordinates R0
      ! Find COM and center molecule in COM
      xx = 0.0_dp
      yy = 0.0_dp
      zz = 0.0_dp
      DO i = 1, natom
         xx = xx+r0((i-1)*3+1)*particle_set(i)%atomic_kind%mass
         yy = yy+r0((i-1)*3+2)*particle_set(i)%atomic_kind%mass
         zz = zz+r0((i-1)*3+3)*particle_set(i)%atomic_kind%mass
      END DO
      xx = xx/mtot
      yy = yy/mtot
      zz = zz/mtot
      DO i = 1, natom
         r0p(1, i) = r0((i-1)*3+1)-xx
         r0p(2, i) = r0((i-1)*3+2)-yy
         r0p(3, i) = r0((i-1)*3+3)-zz
      END DO
      loc_tr(1) = xx
      loc_tr(2) = yy
      loc_tr(3) = zz
      ! Give back the translational vector
      IF (PRESENT(transl)) THEN
         transl(1) = transl(1)-xx
         transl(2) = transl(2)-yy
         transl(3) = transl(3)-zz
      END IF
      M = 0.0_dp
      !
      DO i = 1, natom
         IF (w(i) == 0.0_dp) CYCLE
         rr(1) = rp(1, I)
         rr(2) = rp(2, I)
         rr(3) = rp(3, I)
         rr0(1) = r0p(1, I)
         rr0(2) = r0p(2, I)
         rr0(3) = r0p(3, I)
         rrsq = w(I)*(rr0(1)**2+rr0(2)**2+rr0(3)**2+rr(1)**2+rr(2)**2+rr(3)**2)
         rr0(1) = w(I)*rr0(1)
         rr0(2) = w(I)*rr0(2)
         rr0(3) = w(I)*rr0(3)
         M(1, 1) = M(1, 1)+rrsq+2.0_dp*(-rr0(1)*rr(1)-rr0(2)*rr(2)-rr0(3)*rr(3))
         M(2, 2) = M(2, 2)+rrsq+2.0_dp*(-rr0(1)*rr(1)+rr0(2)*rr(2)+rr0(3)*rr(3))
         M(3, 3) = M(3, 3)+rrsq+2.0_dp*(rr0(1)*rr(1)-rr0(2)*rr(2)+rr0(3)*rr(3))
         M(4, 4) = M(4, 4)+rrsq+2.0_dp*(rr0(1)*rr(1)+rr0(2)*rr(2)-rr0(3)*rr(3))
         M(1, 2) = M(1, 2)+2.0_dp*(-rr0(2)*rr(3)+rr0(3)*rr(2))
         M(1, 3) = M(1, 3)+2.0_dp*(rr0(1)*rr(3)-rr0(3)*rr(1))
         M(1, 4) = M(1, 4)+2.0_dp*(-rr0(1)*rr(2)+rr0(2)*rr(1))
         M(2, 3) = M(2, 3)-2.0_dp*(rr0(1)*rr(2)+rr0(2)*rr(1))
         M(2, 4) = M(2, 4)-2.0_dp*(rr0(1)*rr(3)+rr0(3)*rr(1))
         M(3, 4) = M(3, 4)-2.0_dp*(rr0(2)*rr(3)+rr0(3)*rr(2))
      END DO
      ! Symmetrize
      M(2, 1) = M(1, 2)
      M(3, 1) = M(1, 3)
      M(3, 2) = M(2, 3)
      M(4, 1) = M(1, 4)
      M(4, 2) = M(2, 4)
      M(4, 3) = M(3, 4)
      ! Solve the eigenvalue problem for M
      Z = M
      CALL diamat_all(Z, lambda)
      ! Pick the correct eigenvectors
      S = 1.0_dp
      IF (Z(1, 1) .LT. 0.0_dp) S = -1.0_dp
      Q(0) = S*Z(1, 1)
      Q(1) = S*Z(2, 1)
      Q(2) = S*Z(3, 1)
      Q(3) = S*Z(4, 1)
      IF (PRESENT(my_val)) THEN
         IF (ABS(lambda(1)) < epsi) THEN
            my_val = 0.0_dp
         ELSE
            my_val = lambda(1)/REAL(natom, KIND=dp)
         ENDIF
      END IF
      IF (ABS(lambda(1)-lambda(2)) < epsi) THEN
         IF (output_unit > 0) WRITE (output_unit, FMT='(T2,"RMSD3|",A)') &
            'NORMAL EXECUTION, NON-UNIQUE SOLUTION'
      END IF
      ! Computes derivatives w.r.t. the positions if requested
      IF (PRESENT(drmsd3)) THEN
         DO I = 1, natom
            IF (W(I) == 0.0_dp) CYCLE
            rr(1) = W(I)*2.0_dp*rp(1, I)
            rr(2) = W(I)*2.0_dp*rp(2, I)
            rr(3) = W(I)*2.0_dp*rp(3, I)
            rr0(1) = W(I)*2.0_dp*r0p(1, I)
            rr0(2) = W(I)*2.0_dp*r0p(2, I)
            rr0(3) = W(I)*2.0_dp*r0p(3, I)

            dm_r(1, 1, 1) = (rr(1)-rr0(1))
            dm_r(1, 1, 2) = (rr(2)-rr0(2))
            dm_r(1, 1, 3) = (rr(3)-rr0(3))

            dm_r(1, 2, 1) = 0.0_dp
            dm_r(1, 2, 2) = rr0(3)
            dm_r(1, 2, 3) = -rr0(2)

            dm_r(1, 3, 1) = -rr0(3)
            dm_r(1, 3, 2) = 0.0_dp
            dm_r(1, 3, 3) = rr0(1)

            dm_r(1, 4, 1) = rr0(2)
            dm_r(1, 4, 2) = -rr0(1)
            dm_r(1, 4, 3) = 0.0_dp

            dm_r(2, 2, 1) = (rr(1)-rr0(1))
            dm_r(2, 2, 2) = (rr(2)+rr0(2))
            dm_r(2, 2, 3) = (rr(3)+rr0(3))

            dm_r(2, 3, 1) = -rr0(2)
            dm_r(2, 3, 2) = -rr0(1)
            dm_r(2, 3, 3) = 0.0_dp

            dm_r(2, 4, 1) = -rr0(3)
            dm_r(2, 4, 2) = 0.0_dp
            dm_r(2, 4, 3) = -rr0(1)

            dm_r(3, 3, 1) = (rr(1)+rr0(1))
            dm_r(3, 3, 2) = (rr(2)-rr0(2))
            dm_r(3, 3, 3) = (rr(3)+rr0(3))

            dm_r(3, 4, 1) = 0.0_dp
            dm_r(3, 4, 2) = -rr0(3)
            dm_r(3, 4, 3) = -rr0(2)

            dm_r(4, 4, 1) = (rr(1)+rr0(1))
            dm_r(4, 4, 2) = (rr(2)+rr0(2))
            dm_r(4, 4, 3) = (rr(3)-rr0(3))

            DO ix = 1, 3
               dm_r(2, 1, ix) = dm_r(1, 2, ix)
               dm_r(3, 1, ix) = dm_r(1, 3, ix)
               dm_r(4, 1, ix) = dm_r(1, 4, ix)
               dm_r(3, 2, ix) = dm_r(2, 3, ix)
               dm_r(4, 2, ix) = dm_r(2, 4, ix)
               dm_r(4, 3, ix) = dm_r(3, 4, ix)
            END DO
            !
            DO ix = 1, 3
               drmsd3(ix, I) = 0.0_dp
               DO k = 1, 4
                  DO j = 1, 4
                     drmsd3(ix, i) = drmsd3(ix, i)+Q(K-1)*Q(j-1)*dm_r(j, k, ix)
                  END DO
               END DO
               drmsd3(ix, I) = drmsd3(ix, I)/REAL(natom, KIND=dp)
            END DO
         END DO
      END IF
      ! Computes the rotation matrix in terms of quaternions
      my_rot(1, 1) = -2.0_dp*Q(2)**2-2.0_dp*Q(3)**2+1.0_dp
      my_rot(1, 2) = 2.0_dp*(-Q(0)*Q(3)+Q(1)*Q(2))
      my_rot(1, 3) = 2.0_dp*(Q(0)*Q(2)+Q(1)*Q(3))
      my_rot(2, 1) = 2.0_dp*(Q(0)*Q(3)+Q(1)*Q(2))
      my_rot(2, 2) = -2.0_dp*Q(1)**2-2.0_dp*Q(3)**2+1.0_dp
      my_rot(2, 3) = 2.0_dp*(-Q(0)*Q(1)+Q(2)*Q(3))
      my_rot(3, 1) = 2.0_dp*(-Q(0)*Q(2)+Q(1)*Q(3))
      my_rot(3, 2) = 2.0_dp*(Q(0)*Q(1)+Q(2)*Q(3))
      my_rot(3, 3) = -2.0_dp*Q(1)**2-2.0_dp*Q(2)**2+1.0_dp
      IF (PRESENT(rot)) rot = my_rot
      ! Give back coordinates rotated in order to minimize the RMSD
      IF (my_rotate) THEN
         DO i = 1, natom
            CALL matvec_3x3(r((i-1)*3+1:(i-1)*3+3), TRANSPOSE(my_rot), rp(:, i))
            r((i-1)*3+1:(i-1)*3+3) = r((i-1)*3+1:(i-1)*3+3)+loc_tr
         END DO
      END IF
      DEALLOCATE (w)
      DEALLOCATE (rp)
      DEALLOCATE (r0p)
   END SUBROUTINE rmsd3

END MODULE rmsd
