!
! Copyright (C) 2002-2005 FPMD-CPV groups
! This file is distributed under the terms of the
! GNU General Public License. See the file `License'
! in the root directory of the present distribution,
! or http://www.gnu.org/copyleft/gpl.txt .
!
!------------------------------------------------------------------------------!
  MODULE ions_nose
      !------------------------------------------------------------------------------!
      !! The present code allows one to use "massive" Nose-Hoover chains:  
      !! TOBIAS DJ, MARTYNA GJ, KLEIN ML  
      !! JOURNAL OF PHYSICAL CHEMISTRY 97 (49): 12959-12966 DEC 9 1993.
      !
      USE kinds, ONLY: DP
      !
      IMPLICIT NONE
      !
      ! Some comments are in order on how Nose-Hoover chains work here (K.N. Kudin)
      !
      INTEGER :: nhpdim=1
      !! the total number of the resulting NH chains
      INTEGER :: nhpend=0
      !! it is 1 if there is a chain above all chains, otherwise it is 0
      INTEGER :: nhpbeg=0
      !! it is usually 0, however, if using groups (nhptyp = 3), it might
      !! be desirable to have atoms with uncontrolled temperature, so then
      !! nhpbeg becomes 1, and all the "uncontrolled" atoms are assigned to the
      !! 1st thermostat that is always zero in velocity and so it does not
      !! affect ionic motion.
      INTEGER :: nhptyp=0
      !! one chain for the whole system is specified by nhptyp=0 (or nothing)
      !! currently input options allow one chain per atomic type (nhptyp=1),
      !! one chain per atom (nhptyp=2), and fancy stuff with nhptyp=3 (& nhgrp).
      !
      INTEGER :: nhpcl=1, ndega 
      !
      ! see subroutine ions_nose_allocate on what are the dimensions of these
      ! variables
      !
      INTEGER, ALLOCATABLE   :: atm2nhp(:)
      !! gives the chain number from the atom list (which is sorted by type)
      INTEGER, ALLOCATABLE   :: anum2nhp(:)
      !! the number of degrees of freedom per chain (now just 3*nat_i)
      REAL(DP), ALLOCATABLE :: ekin2nhp(:)
      !! the kinetic energy of the present chain
      REAL(DP), ALLOCATABLE :: gkbt2nhp(:)
      !! the NH chain parameters
      REAL(DP), ALLOCATABLE :: qnp(:)
      !! the chain masses
      REAL(DP), ALLOCATABLE :: qnp_(:)
      !! a temporary array
      !
      REAL(DP), ALLOCATABLE :: vnhp(:), xnhp0(:), xnhpm(:), xnhpp(:), &
       scal2nhp(:), fnosep(:)

      REAL(DP) :: gkbt = 0.0_DP
      REAL(DP) :: kbt = 0.0_DP
      REAL(DP) :: tempw = 0.0_DP

!------------------------------------------------------------------------------!
  CONTAINS 
!------------------------------------------------------------------------------!

  subroutine ions_nose_init( tempw_ , fnosep_ , nhpcl_ , nhptyp_ , ndega_ , nhgrp_ , fnhscl_)
    use constants,      only: k_boltzmann_au, pi, au_terahertz
    use control_flags,  only: tnosep
    use ions_base,      only: ndofp, tions_base_init, nsp, nat, na, amass, ityp
    real(DP), intent(in)  :: tempw_ , fnosep_(:), fnhscl_(:) 
    integer, intent(in) :: nhpcl_ , nhptyp_ , ndega_ , nhgrp_(:)
    integer :: i, j, is, ia
    REAL(DP) :: amass_mean

    IF( .NOT. tions_base_init ) &
      CALL errore(' ions_nose_init ', ' you should call ions_base_init first ', 1 )
    !
    tempw     = tempw_
    !
    IF( ALLOCATED( atm2nhp ) ) DEALLOCATE( atm2nhp )
    ALLOCATE( atm2nhp( nat ) )
    !
    atm2nhp(1:nat) = 1
    !
    if (tnosep) then
       nhpcl = MAX( nhpcl_ , 1 )
       if (abs(nhptyp_).eq.1) then
          nhptyp = 1
          if (nhptyp_.gt.0) nhpend = 1
          nhpdim = nsp
          do ia=1,nat
             atm2nhp(ia) = ityp(ia)
          enddo
       elseif (abs(nhptyp_).eq.2) then
          nhptyp = 2
          if (nhptyp_.gt.0) nhpend = 1
          nhpdim = nat
          do i=1,nat
             atm2nhp(i) = i
          enddo
       elseif (abs(nhptyp_).eq.3) then
          nhptyp = 3
          if (nhptyp_.gt.0) nhpend = 1
          call set_atmnhp(nhgrp_,atm2nhp,nhpdim,nhpbeg)
       endif
       ! Add one more chain on top if needed
       nhpdim = nhpdim + nhpend

    endif
    !
    CALL ions_nose_allocate()
    !
    !  Setup Nose-Hoover chain masses
    !
    if ( ndega_ > 0 ) then
       ndega = ndega_
    else if ( ndega_ < 0 ) then
       ndega = ndofp - ( - ndega_ )
    else
       ndega = ndofp
    endif

    !  there is no need to initialize any Nose variables except for nhpcl
    !  and ndega if the thermostat is not used
    !

    IF( tnosep ) THEN

      IF( nhpcl > SIZE( fnosep_ ) ) &
        CALL errore(' ions_nose_init ', ' fnosep size too small ', nhpcl )

      ! count the number of atoms per thermostat and set the value
      anum2nhp = 0
      ! Here we shall check if the scaling factors are provided
      If (maxval(fnhscl_(1:nsp)).lt.0.0d0) then
         scal2nhp = DBLE(ndega)/DBLE(3*nat)
      else
         scal2nhp = -1.0_DP
      endif
      !
      do ia=1,nat
         is = ityp(ia)
         anum2nhp(atm2nhp(ia)) = anum2nhp(atm2nhp(ia)) + 3
         if (scal2nhp(atm2nhp(ia)).lt.0.0_DP) &
              scal2nhp(atm2nhp(ia)) = fnhscl_(is)
      enddo
      if (nhpend.eq.1) anum2nhp(nhpdim) = nhpdim - 1 - nhpbeg
      ! set gkbt2nhp for each thermostat
      do is=1,nhpdim
         gkbt2nhp(is) = DBLE(anum2nhp(is)) * tempw * k_boltzmann_au
      enddo
      ! scale the target energy by some factor convering 3*nat to ndega
      if (nhpdim.gt.1) then
         do is=1,(nhpdim-nhpend)
            if (scal2nhp(is).lt.0.0_DP) scal2nhp(is) = 1.0_DP
            gkbt2nhp(is) = gkbt2nhp(is)*scal2nhp(is)
         enddo
      endif
      !
      gkbt = DBLE( ndega ) * tempw * k_boltzmann_au
      if (nhpdim.eq.1) gkbt2nhp(1) = gkbt
      kbt  = tempw * k_boltzmann_au

      fnosep(1) = fnosep_ (1)
      if( fnosep(1) > 0.0_DP ) then
        qnp_(1) = 2.0_DP * gkbt / ( fnosep(1) * ( 2.0_DP * pi ) * au_terahertz )**2
      end if

      if ( nhpcl > 1 ) then
        do i = 2, nhpcl
          fnosep(i) = fnosep_ (i)
          if( fnosep(i) > 0.0_DP ) then
            qnp_(i) = 2.0_DP * tempw  * k_boltzmann_au / &
                 ( fnosep(i) * ( 2.0_DP * pi ) * au_terahertz )**2
          else
            qnp_(i) = qnp_(1) / DBLE(ndega)
          endif
        enddo
      endif
      !
      ! set the NH masses for all the chains
      !
      IF((nhptyp.EQ.1).OR.(nhptyp.EQ.2)) THEN 
        !===============================================================
        ! for nhptyp = 1 or 2, the default NH masses are multiplied by 
        ! the atomic masses so that the relative rates of thermostat 
        ! fluctuations are inversely proportional to the masses of the
        ! of the particles to which they are coupled (JOURNAL OF PHYSICAL
        ! CHEMISTRY 97 (49): 12959-12966 DEC 9 1993). Helps in faster
        ! equipartitioning of thermal energy. When nhpend = 1, the mass
        ! of the common thermostat is multiplied by mean atomic mass of 
        ! the system, amass _mean, = total atomic mass / number of atoms  
        !
        DO j=1,(nhpdim-nhpend)
          ! 
          IF(nhptyp.EQ.2) qnp((j-1)*nhpcl+1)=amass(ityp(j))*qnp_(1)*gkbt2nhp(j)/gkbt  
          IF(nhptyp.EQ.1) qnp((j-1)*nhpcl+1)=amass(j)*qnp_(1)*gkbt2nhp(j)/gkbt  
          !
          IF(nhpcl.GT.1) THEN
            !
            DO i=2,nhpcl
              !
              IF(nhptyp.EQ.2) qnp((j-1)*nhpcl+i)=amass(ityp(j))*qnp_(i)
              IF(nhptyp.EQ.1) qnp((j-1)*nhpcl+i)=amass(j)*qnp_(i)
              !
            END DO
            !
          END IF
          ! 
        END DO !j
        !
        IF(nhpend.EQ.1) THEN
          !
          amass_mean=0.0_DP
          !
          DO is=1,nat
            ! 
            amass_mean=amass_mean+amass(ityp(is))
            !
          END DO
          !
          amass_mean=amass_mean/DBLE(nat)
          !
          qnp((nhpdim-1)*nhpcl+1)=amass_mean*qnp_(1)*gkbt2nhp(nhpdim)/gkbt
          !
          IF(nhpcl.GT.1) THEN
            !
            DO i=2,nhpcl
              !
              qnp((nhpdim-1)*nhpcl+i)=amass_mean*qnp_(i)
              !
            END DO
            !
          END IF
          !
        END IF
        !
        !===============================================================
        !
      ELSE
        !
        !===============================================================
        ! Default code : for nhptyp = 1 or 3 
        DO j=1,nhpdim
           qnp((j-1)*nhpcl+1) = qnp_(1)*gkbt2nhp(j)/gkbt
           IF (nhpcl > 1) THEN
              DO i=2,nhpcl
                 qnp((j-1)*nhpcl+i) = qnp_(i)
              END DO
           END IF
        END DO
        ! Default code end 
        !
      END IF !nhptyp
      !
      !
   END IF !tnosep


    !    WRITE( stdout,100)
    !    WRITE( stdout,110) QNOSEP,TEMPW
    !    WRITE( stdout,120) GLIB
    !    WRITE( stdout,130) NSVAR
! 100  FORMAT(//' * Temperature control of ions with nose thermostat'/)
! 110  FORMAT(3X,'nose mass:',F12.4,' temperature (K):',F12.4)
! 120  FORMAT(3X,'ionic degrees of freedom:        ',F5.0)
! 130  FORMAT(3X,'time steps per nose oscillation: ',I5,//)

    return
  end subroutine ions_nose_init

  subroutine set_atmnhp(nhgrp,atm2nhp,nhpdim,nhpbeg)
    !
    use ions_base,      only: nsp, nat, ityp
    IMPLICIT NONE
    integer, intent(in) :: nhgrp(:)
    integer, intent(out) :: nhpdim, nhpbeg, atm2nhp(:)
    !
    integer :: i,is,ia,igrpmax,ith
    INTEGER, ALLOCATABLE   :: indtmp(:)
    !
    ! find maximum group
    igrpmax = max(maxval(nhgrp(1:nsp)),1)
    ! find out which groups are assigned (assuming gaps)
    allocate(indtmp(igrpmax))
    indtmp=0
    do is=1,nsp
       if (nhgrp(is).gt.0) indtmp(nhgrp(is)) = 1 
    enddo
    ! assign thermostat index to requested groups
    ith = 0
    ! make the 1st thermostat idle if there are negative groups
    if (minval(nhgrp(1:nsp)).lt.0) ith = 1
    nhpbeg = ith
    !
    do i=1,igrpmax
       if (indtmp(i).gt.0) then
          ith = ith + 1
          indtmp(i) = ith
       endif
    enddo
    ! assign thermostats to atoms depending on what is requested
    do ia=1,nat
       is=ityp(ia)
       if (nhgrp(is).gt.0) then
          atm2nhp(ia) = indtmp(nhgrp(is))
       elseif (nhgrp(is).eq.0) then
          ith = ith + 1
          atm2nhp(ia) = ith
       else
          atm2nhp(ia) = 1
       endif
    enddo
    nhpdim = ith
    deallocate(indtmp)
    return
    !
  end subroutine set_atmnhp

  SUBROUTINE ions_nose_allocate()
    !
    IMPLICIT NONE
    !
    IF ( .NOT. ALLOCATED( vnhp ) )     ALLOCATE( vnhp(  nhpcl*nhpdim ) )
    IF ( .NOT. ALLOCATED( xnhp0 ) )    ALLOCATE( xnhp0( nhpcl*nhpdim ) )
    IF ( .NOT. ALLOCATED( xnhpm ) )    ALLOCATE( xnhpm( nhpcl*nhpdim ) )
    IF ( .NOT. ALLOCATED( xnhpp ) )    ALLOCATE( xnhpp( nhpcl*nhpdim ) )
    IF ( .NOT. ALLOCATED( ekin2nhp ) ) ALLOCATE( ekin2nhp( nhpdim ) )
    IF ( .NOT. ALLOCATED( gkbt2nhp ) ) ALLOCATE( gkbt2nhp( nhpdim ) )
    IF ( .NOT. ALLOCATED( scal2nhp ) ) ALLOCATE( scal2nhp( nhpdim ) )
    IF ( .NOT. ALLOCATED( anum2nhp ) ) ALLOCATE( anum2nhp( nhpdim ) )
    IF ( .NOT. ALLOCATED( qnp ) )      ALLOCATE( qnp( nhpcl*nhpdim ) )
    IF ( .NOT. ALLOCATED( qnp_ ) )     ALLOCATE( qnp_( nhpcl ) )
    IF ( .NOT. ALLOCATED( fnosep ) )   ALLOCATE( fnosep( nhpcl ) )
    !
    vnhp  = 0.0_DP
    xnhp0 = 0.0_DP
    xnhpm = 0.0_DP
    xnhpp = 0.0_DP
    qnp   = 0.0_DP
    qnp_  = 0.0_DP
    !
    RETURN
    !
  END SUBROUTINE ions_nose_allocate

  SUBROUTINE ions_nose_deallocate()
    !
    IMPLICIT NONE
    !
    IF ( ALLOCATED( vnhp ) )     DEALLOCATE( vnhp )
    IF ( ALLOCATED( xnhp0 ) )    DEALLOCATE( xnhp0 )
    IF ( ALLOCATED( xnhpm ) )    DEALLOCATE( xnhpm )
    IF ( ALLOCATED( xnhpp ) )    DEALLOCATE( xnhpp )
    IF ( ALLOCATED( ekin2nhp ) ) DEALLOCATE( ekin2nhp )
    IF ( ALLOCATED( gkbt2nhp ) ) DEALLOCATE( gkbt2nhp )
    IF ( ALLOCATED( scal2nhp ) ) DEALLOCATE( scal2nhp )
    IF ( ALLOCATED( anum2nhp ) ) DEALLOCATE( anum2nhp )
    IF ( ALLOCATED( qnp ) )      DEALLOCATE( qnp )
    IF ( ALLOCATED( qnp_ ) )     DEALLOCATE( qnp_ )
    IF ( ALLOCATED( fnosep ) )   DEALLOCATE( fnosep )
    !
    IF( ALLOCATED( atm2nhp ) ) DEALLOCATE( atm2nhp )
    !
    RETURN
    !
  END SUBROUTINE ions_nose_deallocate

  SUBROUTINE ions_nose_info(delt)

      use constants,     only: au_terahertz, pi
      USE io_global,     ONLY: stdout
      USE control_flags, ONLY: tnosep
      use ions_base,  only: nat

      IMPLICIT NONE
      REAL(DP), INTENT(IN) :: delt

      INTEGER   :: nsvar, i, j
      REAL(DP) :: wnosep

      IF( tnosep ) THEN
        !
        IF( fnosep(1) <= 0.0_DP) &
          CALL errore(' ions_nose_info ', ' fnosep less than zero ', 1)
        IF( delt <= 0.0_DP) &
          CALL errore(' ions_nose_info ', ' delt less than zero ', 1)

        wnosep = fnosep(1) * ( 2.0_DP * pi ) * au_terahertz
        nsvar  = ( 2.0_DP * pi ) / ( wnosep * delt )

        WRITE( stdout,563) tempw, nhpcl, ndega, nsvar
        WRITE( stdout,564) (fnosep(i),i=1,nhpcl)
        WRITE( stdout,565) nhptyp, (nhpdim-nhpend), nhpend , nhpbeg, &
             (anum2nhp(j),j=1,nhpdim)
        IF(nhptyp.EQ.1.OR.nhptyp.EQ.2)WRITE( stdout,'(//, &
        & "*** default NH masses are multiplied by atomic masses ***")')
        do j=1,nhpdim
           WRITE( stdout,566) j,(qnp((j-1)*nhpcl+i),i=1,nhpcl)
        enddo
        WRITE( stdout,567)
        do j=1,nat,20
           WRITE( stdout,568) atm2nhp(j:min(nat,j+19))
        enddo
     END IF

 563  format( //, &
     &       3X,'ion dynamics with nose` temperature control:', /, &
     &       3X,'temperature required      = ', f10.5, ' (kelvin) ', /, &
     &       3X,'NH chain length           = ', i3, /, &
     &       3X,'active degrees of freedom = ', i6, /, &
     &       3X,'time steps per nose osc.  = ', i6 )
 564  format( //, &
     &       3X,'nose` frequency(es)       = ', 20(1X,f10.3) ) 
! 565  format( //, &
!     &       3X,'nose` mass(es)            = ', 20(1X,f10.3), // ) 
 565  FORMAT( //, &
     &       3X,'the requested type of NH chains is ',I5, /, &
     &       3X,'total number of thermostats used ',I5,1X,I1,1X,I1, /, &
     &       3X,'ionic degrees of freedom for each chain ',20(1X,I3)) 
 566  format( //, &
     &       3X,'nose` mass(es) for chain ',i4,' = ', 20(1X,f10.3)) 
567  format( //, &
     &       3X,'atom i (in sorted order) is assigned to this thermostat :')
568  format(20(1X,I3))
    RETURN
  END SUBROUTINE ions_nose_info

  

  subroutine ions_nosevel( vnhp, xnhp0, xnhpm, delt, nhpcl, nhpdim )
    implicit none
    integer, intent(in) :: nhpcl, nhpdim
    real(DP), intent(inout) :: vnhp(nhpcl,nhpdim)
    real(DP), intent(in) :: xnhp0(nhpcl,nhpdim), xnhpm(nhpcl,nhpdim), delt
    integer :: i,j
    do j=1,nhpdim
       do i=1,nhpcl
          vnhp(i,j)=2.0_DP * (xnhp0(i,j)-xnhpm(i,j)) / delt-vnhp(i,j)
       end do
    end do
        !
        !  this is equivalent to:
        !  velocity = ( 3.0_DP * xnos0(1) - 4.0_DP * xnosm(1) + xnos2m(1) ) / ( 2.0_DP * delt )
        !  but we do not need variables at time t-2dt ( xnos2m )
        !
    return
  end subroutine ions_nosevel


 subroutine ions_noseupd( xnhpp, xnhp0, xnhpm, delt, qnp, ekin2nhp, gkbt2nhp, vnhp, kbt, nhpcl, nhpdim, nhpbeg, nhpend )
    implicit none
    integer, intent(in) :: nhpcl, nhpdim, nhpbeg, nhpend
    real(DP), intent(out) :: xnhpp(nhpcl,nhpdim)
    real(DP), intent(in) :: xnhp0(nhpcl,nhpdim), xnhpm(nhpcl,nhpdim), delt, qnp(nhpcl,nhpdim), gkbt2nhp(:), kbt
    real(DP), intent(inout) :: ekin2nhp(:), vnhp(nhpcl,nhpdim)
    integer :: i, j
    real(DP) :: dt2, zetfrc, vp1dlt, ekinend, vp1dend


    ekinend = 0.0_DP
    vp1dend = 0.0_DP
    if ( nhpend == 1 ) vp1dend = 0.5_DP * delt * vnhp(1,nhpdim)
    dt2 = delt**2
    if (nhpbeg.gt.0) then
       xnhpp(:,1:nhpbeg) = 0.0_DP
       vnhp (:,1:nhpbeg) = 0.0_DP
    endif
    do j=(1+nhpbeg),nhpdim
    zetfrc = dt2 * ( 2.0_DP * ekin2nhp(j) - gkbt2nhp(j) )
    if ( nhpcl > 1 ) then
       do i=1,(nhpcl-1)
          vp1dlt = 0.5_DP * delt * vnhp(i+1,j)
          xnhpp(i,j)=(2.0_DP * xnhp0(i,j)-(1.0_DP-vp1dlt)*xnhpm(i,j)+zetfrc/qnp(i,j))&
               &   /(1.0_DP+vp1dlt)
!           xnhpp(i,j)=(4.d0*xnhp0(i,j)-(2.d0-delt*vnhp(i+1,j))*xnhpm(i,j)+2.0d0*dt2*zetfrc/qnp(i,j))&
!                &   /(2.d0+delt*vnhp(i+1,j))
          vnhp(i,j) =(xnhpp(i,j)-xnhpm(i,j))/( 2.0_DP * delt )
          zetfrc = dt2*(qnp(i,j)*vnhp(i,j)**2-kbt)
       end do
    end if
    ! Last variable
    i = nhpcl
    if ( nhpend == 0 ) then
       xnhpp(i,j)=2.0_DP * xnhp0(i,j)-xnhpm(i,j) + zetfrc / qnp(i,j)
       vnhp(i,j) =(xnhpp(i,j)-xnhpm(i,j))/( 2.0_DP * delt )
    elseif (nhpend == 1) then
       xnhpp(i,j)=(2.0_DP*xnhp0(i,j)-(1.0_DP-vp1dend)*xnhpm(i,j)+zetfrc/qnp(i,j))&
            &   /(1.0_DP+vp1dend)       
       vnhp(i,j) =(xnhpp(i,j)-xnhpm(i,j))/( 2.0_DP * delt )
       ekinend = ekinend + (qnp(i,j)*vnhp(i,j)**2)
       if (j.eq.(nhpdim-nhpend)) then
          ekin2nhp(nhpdim) = 0.5_DP*ekinend       
          vp1dend = 0.0_DP
       endif
    endif
    enddo
    ! Update velocities
!    do i=1,nhpcl
!       vnhp(i) =(xnhpp(i)-xnhpm(i))/( 2.0d0 * delt )
!    end do
    ! These are the original expressions from cpr.f90
    !       xnhpp(1)=2.*xnhp0(1)-xnhpm(1)+2.*( delt**2 / qnp(1) )*(ekinpr-gkbt/2.)
    !       vnhp(1) =(xnhpp(1)-xnhpm(1))/( 2.0d0 * delt )
    return
  end subroutine ions_noseupd

  

  real(DP) function ions_nose_nrg( xnhp0, vnhp, qnp, gkbt2nhp, kbt, nhpcl, nhpdim )
    implicit none
    integer :: nhpcl, nhpdim
    real(DP) :: gkbt2nhp(:), qnp(nhpcl,nhpdim),vnhp(nhpcl,nhpdim),xnhp0(nhpcl,nhpdim),kbt
    integer :: i,j
    real(DP) :: stmp
    !
    stmp = 0.0_DP
    do j=1,nhpdim
    stmp = stmp + 0.5_DP * qnp(1,j) * vnhp(1,j) * vnhp(1,j) + gkbt2nhp(j) * xnhp0(1,j)
    if (nhpcl > 1) then
       do i=2,nhpcl
          stmp = stmp + 0.5_DP * qnp(i,j) * vnhp(i,j) * vnhp(i,j) + kbt * xnhp0(i,j)
       enddo
    endif
    enddo
    ions_nose_nrg = stmp
    return
  end function ions_nose_nrg



  subroutine ions_nose_shiftvar( xnhpp, xnhp0, xnhpm )
    !  shift values of nose variables to start a new step
    implicit none
    real(DP), intent(inout) :: xnhpm(:), xnhp0(:)
    real(DP), intent(in)  :: xnhpp(:)
      !
      xnhpm = xnhp0
      xnhp0 = xnhpp
      !
    return
  end subroutine ions_nose_shiftvar
  
!------------------------------------------------------------------------------!
  END MODULE ions_nose
!------------------------------------------------------------------------------!
