#!/bin/csh -x

# __author__ = Felix Kleinert
# __date__  = '2020-04-30'
# This script creates run scripts for JUWELS or HDFML

# When you call this script directly you can use
#   $1 which has to be `juwels' or `hdfml'.
#   $2 which is the path where the run scripts should be stored

if [[ $1 != '' ]]; then
  hpcsys=$1
else
  if [[ $HOSTNAME == *"juwels"* ]]; then
    hpcsys="juwels"
  elif [[ $HOSTNAME == *"hdfml"* ]]; then
    hpcsys="hdfml"
  else
     echo  "Unknown hpc host \`$HOSTNAME\`. Pass 'juwels' or 'hdfml' as first argument."
     exit
  fi
fi

if [[ $2 != '' ]]; then
  cur=$2
else
  cur=$PWD
fi

echo "############################################################"
echo "#                                                          #"
echo "#            user interaction required                     #"
echo "#                                                          #"
echo "############################################################"
echo 

echo "This script creates the HPC batch scripts to run mlt on compute nodes on JUWELS or hdfml."
echo "You can modify the created run scripts afterwards if needed."

echo
echo
echo "Creating run script for $hpcsys:"
echo

budget=''
while [[ $budget == '' ]]
do
 echo
 read -p "Enter project budget for --account flag: " budget 
done

email=`jutil user show -o json | grep email | cut -f2 -d':' | cut -f1 -d',' | cut -f2 -d'"'`
echo
read -p "Enter e-mail address for --mail-user (default: ${email}): " new_email

if [[ -z "$new_email" ]]; then
    new_email=$email
fi

# create HPC_logging dir
hpclogging="HPC_logging/"
mkdir -p ${cur}/${hpclogging}


# ordering for looping:
# "partition nGPUs timing"
if [[ $hpcsys = "juwels" ]]; then
  for i in "develgpus 2  02:00:00" "gpus 4 08:00:00"; do
      set -- $i

cat <<EOT > ${cur}/run_${hpcsys}_$1.bash
#!/bin/bash -x
#SBATCH --account=${budget}
#SBATCH --nodes=1
#SBATCH --output=${hpclogging}mlt-out.%j
#SBATCH --error=${hpclogging}mlt-err.%j
#SBATCH --time=$3 
#SBATCH --partition=$1 
#SBATCH --gres=gpu:$2
#SBATCH --mail-type=ALL
#SBATCH --mail-user=${email}

source HPC_setup/mlt_modules_${hpcsys}.sh
source venv_${hpcsys}/bin/activate

timestamp=\`date +"%Y-%m-%d_%H%M-%S"\`

export PYTHONPATH=\${PWD}/venv_${hpcsys}/lib/python3.8/site-packages:\${PYTHONPATH}

srun --cpu-bind=none python run.py --experiment_date=\$timestamp
EOT

  echo "Created runscript: run_${hpcsys}_$1.bash"

  done

elif [[ $hpcsys = "hdfml" ]]; then
cat <<EOT > ${cur}/run_${hpcsys}_batch.bash
#!/bin/bash -x
#SBATCH --account=${budget}
#SBATCH --nodes=1
#SBATCH --output=${hpclogging}mlt-out.%j
#SBATCH --error=${hpclogging}mlt-err.%j
#SBATCH --time=08:00:00
#SBATCH --gres=gpu:4
#SBATCH --mail-type=ALL
#SBATCH --mail-user=${email}

source HPC_setup/mlt_modules_${hpcsys}.sh
source venv_${hpcsys}/bin/activate

timestamp=\`date +"%Y-%m-%d_%H%M-%S"\`

export PYTHONPATH=\${PWD}/venv_${hpcsys}/lib/python3.8/site-packages:\${PYTHONPATH}

srun --cpu-bind=none python run_HPC.py --experiment_date=\$timestamp
EOT

fi

echo "###################################################################################"
echo "#   You have to run the the following command on a login node to download data:   #"
echo "#                              \`python run_HPC.py'                                #"
echo "#                                                                                 #"

echo "#   Please execute the following command to check if the setup went well:         #"
if [[ ${hpcsys} = 'juwels' ]]; then
  echo "#                   \`sbatch run_juwels_develgpus.bash'                            #"
else
  echo "#                   \`sbatch run_hdfml_batch.bash'                                 #"
fi

echo "###################################################################################"

