#!/bin/bash
#SBATCH -o slurm-%j.out
#SBATCH -e slurm-%j.err
#SBATCH --nodes=1
#SBATCH --gpus=1
#SBATCH --ntasks-per-node=1
#SBATCH --partition=gpu-shared
#SBATCH -t 2-00:00:00
#SBATCH --no-requeue
#SBATCH --reservation=amber24gpu
#SBATCH --account=gue998

module load workshop/amber24/default

set -x 

echo AMBERHOME is $AMBERHOME
test -f ${AMBERHOME}/amber.sh && source ${AMBERHOME}/amber.sh
LAUNCH="srun --kill-on-bad-exit"
EXE=${AMBERHOME}/bin/pmemd.cuda


BASE=$1
PARM=${BASE}.parm7
REF=${BASE}.rst7
mdins=mdins
NAME=${BASE}

steps="\
000_Min \
001_NPT \
002_NPT \
003_NPT \
01_Min \
02_Heat \
03_Equil_NPT_solvent \
04_Min_solute \
05_Heat_solute_Rwt25 \
06_Equil_solute_Rwt25 \
07_Equil_solute_Rwt10 \
08_Equil_solute_Rwt5 \
09_Equil_solute_Rwt2 \
"

#switch for skipping existing steps
k=1
for step in $steps; do
  ICRD=${BASE}.rst7
  MDIN=${mdins}/${step}.mdin
  BASE=${step}-run
  
  checkMDout1=$(tail -n 1 ${BASE}.mdout 2>/dev/null | grep "Total wall time")
  echo "Check 1 is $checkMDout1"
  checkMDout2=$(grep "wallclock() was called" ${BASE}.mdout)
  echo "Check 2 is $checkMDout2"

  if [ "$k" == 1 ]; then
    if [ -f "${BASE}.mdout" ] && [ ! -z "$checkMDout1" ]; then
       echo "${BASE} step complete, skipping"
       continue;
    elif [ -f "${BASE}.mdout" ] && [ ! -z "$checkMDout2" ]; then
       echo "${BASE} step complete, skipping"
       continue;
    else
       k=0
       echo "Starting from step ${step} and not skipping"
    fi
  fi

  if [ "$step" == 000_Min ] || [ "$step" == 01_Min ]; then
	  restartfile=${BASE}.rst7
	  trajFile=${BASE}.rst7
	  EXE=${AMBERHOME}/bin/sander

  else
	  restartfile=${BASE}.rst7
	  trajFile=${BASE}.nc
	  EXE=${AMBERHOME}/bin/pmemd.cuda
  fi

  echo running ${MDIN}
  ${LAUNCH} ${EXE} -O -p ${PARM} -c ${ICRD} -i ${MDIN} -o ${BASE}.mdout -inf ${BASE}.mdinfo -x ${BASE}.nc -r ${BASE}.rst7 -ref ${REF}
  
  mv ${BASE}.mdinfo ${BASE}.stat
  echo "MDOUT IS ${BASE}.mdout"

  checkMDout1=$(tail -n 1 ${BASE}.mdout 2>/dev/null | grep "Total wall time")
  echo "Check 1 is $checkMDout1"
  checkMDout2=$(grep "wallclock() was called" ${BASE}.mdout)
  echo "Check 2 is $checkMDout2"

  if [ -f "${BASE}.mdout" ] && [ ! -z "$checkMDout1" ]; then
	  echo "Check 1 is $checkMDout1"
          echo "${BASE} step complete, moving on"
	  continue;
  elif [ -f "${BASE}.mdout" ] && [ ! -z "$checkMDout2" ]; then
          echo "Check 2 is $checkMDout2"
	  echo "${BASE} step complete, moving on"
	  continue;
  else
	  echo "${BASE} did not finish correctly, exiting."
	  exit 1
  fi

done