#!/bin/bash
#SBATCH -J _JOBNAME_
#SBATCH -t _TIMELIMIT_:00:00
#SBATCH -n _NCORES_
#SBATCH --tasks-per-node=24
#SBATCH -e Job.err

module load fortran/intel
module load mkl
module load mpi/impi

ulimit -s unlimited

WorkDir=` pwd  `       # Submission Dir is the working dir
ScratchDir=${TMPDIR}  # Temporary dir is set (already with JobID) 
VASP="/home/miglio/VASP_5.3.5/vasp.5.3.5/vasp-SRP032"

cd ${ScratchDir}
scp -r ${WorkDir}/* .

#########################################
# Velocity adjusting preparatory step   #
#########################################

cp INCAR_ INCAR
sed -i "s/_nsw_/NSW      = 1/" INCAR

# execute vasp
srun ${VASP} >> vasp.out

# move files of interest
mv POSCAR  POSCAR_0
mv OSZICAR OSZICAR_0
mv XDATCAR XDATCAR_0
grep -A 51 "TOTAL-FORCE" OUTCAR                > Forces_0.dat

# get timing  
grep Elapsed OUTCAR >> time.dat

rm OUTCAR vasp* E* D* IB* PC* W* CON* CHG

python - <<END

import os, sys, math

MC           = 12.011
MD           = 2.014
MH           = 1.008 
MPt          = 195.080 
MassMolecule = MC + MH + 3. * MD  
Dt           = 0.4E-15
Masses       = 45 * [ MPt ] + [ MC ] + [ MH ] + 3 * [ MD ] 

POSCARNameInput     = "POSCAR_0" 
ForcesFileName      = "Forces_0.dat"
POSCARNameOutput    = "POSCAR"

lines = open( ForcesFileName, "r" ).readlines()

Forces = []

for i in range( 2, 52 ):
        Force = [ float( lines[i].split()[3] ), float( lines[i].split()[4] ), float( lines[i].split()[5] ) ]
        Forces.append( Force )

lines = open( POSCARNameInput, "r" ).readlines()

MovingDegreesOfFreedom = []

if lines[7][0] == "S" or lines[7][0] == "s":
        for i in range( 9, 59 ):
                AtomDegreesOfFreedom = [ lines[i].split()[3] , lines[i].split()[4] , lines[i].split()[5]  ]
                MovingDegreesOfFreedom.append( AtomDegreesOfFreedom )
else:
        for i in range( 9, 59 ):
                MovingDegreesOfFreedom.append( [ "T", "T", "T" ] )

Velocities = []

for i in range( 60, 110 ):
        Velocity = [ float( lines[i].split()[0] ), float( lines[i].split()[1] ), float( lines[i].split()[2] ) ]
        Velocities.append( Velocity )

DOFs = 0
Ekin = 0.

# v at t=0 is v + 1/2 * a * Dt

for i in range( len( Forces ) ):
        for j in range( 3 ):
                Acceleration = Forces[ i][ j ] / ( Masses[ i ] * 1.6605402E-027 ) *1.60217733E-019 * 1.E-10
                if MovingDegreesOfFreedom[i][j] == "T":
                        Ekin = Ekin + 0.5 * Velocities[ i][ j ] ** 2. * Masses[ i ] * 1822.88839 / 0.52917721092**2. / 41.341373337**2. * 27.21138505
                        Velocities[ i][ j ] = Velocities[ i ][ j ] - 0.5 * Acceleration * Dt /1.E-15
                        DOFs = DOFs + 1
                else:
                        Velocities[ i][ j ] = 0.0

Temp = 2. * Ekin / float( DOFs ) * 11604.9

print "Expected Ekin at first vasp step: %14.11f" %Ekin
print "Expected Temp at first vasp step: %14.11f" %Temp

FileOutput = open( POSCARNameOutput, "w" )

for i in range( 0, 60 ):
        FileOutput.write( lines[ i ] )
for i in range( len( Velocities ) ):
        string = str( "   %14.11f    %14.11f    %14.11f\n" %tuple( [ Velocities[ i][ 0], Velocities[ i][ 1], Velocities[ i][ 2] ] ) )
        FileOutput.write( string )

FileOutput.close()

END

# Now we have written a new POSCAR, we can really start

sed -i "s/NSW      = 1/NSW      = _SUBSTEPS_/" INCAR

# run checker in background
./checker >& /dev/null &
# determine pid of checker
checkerpid=$!
# execute vasp
srun ${VASP} >> vasp.out
# kill checker 
kill ${checkerpid}
#execute one more in order to see whether  a restart is needed
./check-CHD3.py

#check for the restarts number
check=0
while [ -e "XDATCAR_${check}" ]
do
        let check++
done

# get timing
grep Elapsed OUTCAR >> time.dat

# update important outputs
mv POSCAR  POSCAR_${check}
mv OSZICAR OSZICAR_${check}
mv XDATCAR XDATCAR_${check}
mv CONTCAR CONTCAR_${check}

grep -A 51 "TOTAL-FORCE" OUTCAR > Forces_${check}.dat

rm D* E* IB* PC* W* vasprun* OUTCAR vdw*

# Copy files back
scp -r ${ScratchDir}/* ${WorkDir}/.
# clean up
rm -r ${ScratchDir}