#!/bin/bash
#SBATCH --time=_TIMELIMIT_:00:00
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=28
#SBATCH --job-name=_NAME_-_FIRSTJOB_-_LASTJOB_

module load i-PI/2.3.0-intel-2020a-Python-3.8.3

# workdirectory
WORKDIR=${SLURM_SUBMIT_DIR}
cd ${WORKDIR}

# scratch directory
SCRATCH=${VSC_SCRATCH_NODE}/${SLURM_JOB_ID}
mkdir -p ${SCRATCH}
cd ${SCRATCH}

# copy files from the workdir to the scratch
scp -r ${WORKDIR}/RuNNer ${SCRATCH}/.
scp -r ${WORKDIR}/sim.lmp ${SCRATCH}/.
scp -r ${WORKDIR}/Pt111_CH4.lmp ${SCRATCH}/.
scp -r ${WORKDIR}/Pt111_CH4.xyz ${SCRATCH}/.
scp -r ${WORKDIR}/input.xml ${SCRATCH}/.
scp -r ${WORKDIR}/check-CHD3-RPMD.py ${SCRATCH}/.
scp -r ${WORKDIR}/create_poscar_scratch_.py ${SCRATCH}/.
scp -r ${WORKDIR}/DynamicsINPUTS/ ${SCRATCH}/.

# lammps location
#lammps=${VSC_DATA}/Programs/lammps-29Oct20/src/lmp_mpi
lammps=${VSC_DATA}/Programs/lammps-29Oct20/src/lmp_serial

# A few times this error has popped up; simple fix
if compgen -G "/tmp/ipi_*" > /dev/null; then
	rm /tmp/ipi_*
fi

sed -e "s|xFIRSTJOBx|$(echo '_FIRSTJOB_' | sed 's/^0*\([1-9]\)/\1/;s/^0*$/0/')|g" create_poscar_scratch_.py > create_poscar_scratch.py
sed -i "s|xLASTJOBx|$(echo '_LASTJOB_' | sed 's/^0*\([1-9]\)/\1/;s/^0*$/0/')|g" create_poscar_scratch.py
python create_poscar_scratch.py

for i in $(seq -w _FIRSTJOB_ _LASTJOB_)
	do
	
	cd ${i}

	if compgen -G "/tmp/ipi_*" > /dev/null; then
		rm /tmp/ipi_*
	fi

	# start i-PI
	i-pi ../input.xml &
	ipipid=$!

	# The sleep is necessary, otherwise LAMMPS won't be able to connect and will crash since i-PI hasn't started up yet
	sleep 10

	# Start multiple instances of LAMMPS
	for idx in `seq -w 0 23`;
		do
		
		#mpirun -np 14 ${lammps} -in ../sim.lmp > /dev/null 2>&1 &
		#${lammps} -in ../sim.lmp > /dev/null 2>&1 &
		${lammps} -in ../sim.lmp > log.lammps_${idx} &
		
	done
	
	# This block regularly checks if the trajectory has reached an outcome, or if it's stuck, and kills i-pi if it's not dead yet
	sleep 20	# Wait a bit before running the first check to make sure the trajectory has advanced a few steps
	python ../check-CHD3-RPMD.py	# Make sure there is an outcome file
	outcome=$(cat outcome)
	stuck=`tail -n52 simulation.pos_centroid.xyz`
	while [[ "$outcome" == *"UNCLEAR"* ]]; do
		sleep 15
		# Check if there is an outcome
		python ../check-CHD3-RPMD.py
		outcome=$(cat outcome)
		
		# Here we check if the trajectory is stuck
		if [[ "$stuck" == `tail -n52 simulation.pos_centroid.xyz` ]];
			then
			outcome="STUCK"
			else
			stuck=`tail -n52 simulation.pos_centroid.xyz`
		fi
	done
	kill ${ipipid}	# kill i-pi
	python ../check-CHD3-RPMD.py	# run the analysis one final time
	outcome=$(cat outcome)

	wait	# Wait until all LAMMPS instances are exited as well
	
	# Save the structure that causes an extrapolation error
	if [[ "$outcome" == *"UNCLEAR"* ]]; then
		for idx in `seq -w 0 23`; do
			EW=`tail -n3 log.lammps_${idx} | head -n1`
			if [[ "$EW" != *"EXTRAPOLATION WARNING"* ]]; then
				rm simulation.*_${idx}.xyz log.lammps_${idx}
			else
				rm simulation.vel_bead_${idx}.xyz
				tail -n50 simulation.pos_bead_${idx}.xyz > tmp
				mv tmp simulation.pos_bead_${idx}.xyz
			fi
		done
	else
		rm simulation.*_bead_*.xyz log.lammps_*
	fi
	rm simulation.out simulation.pos_centroid.xyz simulation.vel_centroid.xyz simulation.pos_beads.xyz simulation.vel_beads.xyz EW.txt RESTART log.lammps
	
	# if [[ ! "$outcome" == *"UNCLEAR"* ]]; then
		# rm EW.txt
	# fi
	# rm simulation.* log.lammps* RESTART
	
	cd ..

done

cd ${SCRATCH}
rm slurm-*.out

# copy files back
rsync -r {_FIRSTJOB_.._LASTJOB_} ${WORKDIR}/.

# clean up
rm -r ${SCRATCH}

# ready