#!/bin/bash
# A wrapper of qsub
# Add CSUB_ERR and CSUB_OUT to your ~/.bashrc . These are the global locations
# for the .err and .out files
# export CSUB_ERR="${HOME}/Scratch/csub_logs"
# export CSUB_OUT="${HOME}/Scratch/csub_logs"
# A test (qsub does not like binaries like echo), so wrap it:
# echo -e -n '#!/bin/bash\nstr="$@"\necho "$str"\n' > my_echo.sh
# chmod +x my_echo.sh
# csub -f -m"20M" -t"00:00:60" test_echo -- ~/my_echo.sh 'pleeeeease work.....!'
# $ cat ~/Scratch/csub_logs/test_echo.out
# pleeeeease work.....!
# Also see csub --help


# We start off doing a sanity check to make sure the user has put the --
# between the csub arguments and the command + arguments that are being run
# shflags will probably fail anyway if they are missing but this will give
# a more coherent error message
found_sep=1

# Have we found the --
for i in "$@"; do
    if [[ "$i" == '--' ]]; then
        found_sep=0
        break
    fi
done

# If we have not found the -- we do a quick check to make sure that --help
# is not requested
if [[ found_sep -eq 1 ]]; then
    if [[ ! "$1" =~ ^-?-h ]]; then
        # help has not been requested, so we issue the error message
        # and intercept $@ to display --help
        echo -n "error: can't find command separator --" 1>&2
        echo " please use it to separate csub args/cmd args" 1>&2
        set -- '--help'
    fi
fi

# Make sure the environment variables for .err and .out files are defined
# and exist
if [[ -z "$CSUB_ERR" ]] || [[ ! -e "$CSUB_ERR" ]]; then
    echo 'error: please add `export CSUB_ERR="<.err path>"` to your ~/.bashrc'
    exit 1
fi

if [[ -z "$CSUB_OUT" ]] || [[ ! -e "$CSUB_OUT" ]]; then
    echo 'error: please add `export CSUB_OUT="<.out path>"` to your ~/.bashrc'
    exit 1
fi

. shflags

# configure shflags
DEFINE_boolean 'verbose' false 'Ask before submssion' 'v'
DEFINE_boolean 'force' false 'Force overwrite of existing .err/.out files' 'f'
DEFINE_string 'mem' '4G' 'The memory usage per slot' 'm'
DEFINE_string 'time' '01:00:00' 'The time for the job' 't'
DEFINE_string 'shell' '/bin/bash' 'The default shell for the job' 'S'
DEFINE_string 'proc-type' 'smp' 'The type of parallel job smp/mpi' 'y'
DEFINE_string 'project' '-' 'Use specific project nodes and/or policy' 'P'
DEFINE_boolean 'paid' false 'Force the use of the nodes specified in --project if not used the job could run on the main cluster'
DEFINE_integer 'proc' 0 'The number of processors (slots --mem is per slot)' 'p'

FLAGS_HELP="USAGE: $0 [flags] <job name> -- <run command>"

# parse the command-line
FLAGS "$@" || exit $?
eval set -- "${FLAGS_ARGV}"

# Just to pick up stay undefined variables
set -u

# Extract the job name from the positional args
name="$1"
shift

if [[ "$name" == "" ]]; then
    echo "error: name your job!! (good practice)" 1>&2
    exit 1
fi

# Make sure we have some command args
if [[ "$#" -eq 0 ]]; then
   echo "error: no command to run!!" 1>&2
   exit 1
fi

# This is the name of the script that will be run
cmd_arg="$1"
shift

# Generate the names for the .err and .out files
error_file="$CSUB_ERR"/"$name".err
out_file="$CSUB_OUT"/"$name".out

# If we are not blitzing over the top of any existing files, then
# we check their presence and if they exist, we error out
if [[ $FLAGS_force -eq 1 ]]; then
    if [[ -e "$error_file" ]]; then
        echo ".err file exists, use --force to overwrite: $error_file" 1>&2
        exit 1
    fi
    if [[ -e "$out_file" ]]; then
        echo ".out file exists, use --force to overwrite: $out_file" 1>&2
        exit 1
    fi
fi


# Now get the absolute path of the command that we are running. qsub does not
# seem to like them, it also does not like binaries??
if [[ "$cmd_arg" = /* ]]; then
    # Get the full path of the relative path
    cmd_arg=$(readlink -f "$cmd_arg")
elif [[ "$cmd_arg" = ./* ]]; then
    # Get the full path of the relative path
    cmd_arg=$(readlink -f "$cmd_arg")
else
    # Otherwise it is a command in PATH, so we locate it's origin
    cmd_path=$(which "$cmd_arg" 2> /dev/null)

    if [[ $? -ne 0 ]]; then
        echo "error: can't find command in the PATH: $cmd_arg" 1>&2
        exit 1
    fi
    cmd_arg="$cmd_path"
fi

# Now start to build up the qsub argument
qsub_command=(
    "qsub"
    -cwd
    -S "$FLAGS_shell"
    -N "$name"
    -l "mem=""$FLAGS_mem"
    -l "h_rt=""$FLAGS_time"
    -e "$error_file"
    -o "$out_file"
)

# Add the project flag (if defined)
if [[ "$FLAGS_project" != '-' ]]; then
    qsub_command+=(
        -P "$FLAGS_project"
    )

    # paid will only get added if project is defined
    if [[ "$FLAGS_paid" -eq 0 ]]; then
        qsub_command+=(
            -l "paid=1"
        )
    fi
fi

# Add any of the processor flags
if [[ "$FLAGS_proc" -gt 0 ]]; then
    qsub_command+=(
        -pe "$FLAGS_proc_type" "$FLAGS_proc"
    )
fi

# If we are verbose, ask for confirmation before subitting
if [[ $FLAGS_verbose -eq 0 ]]; then
    echo -n "You are about to run this qsub command"
    echo ", press <ENTER> to submit <CTRL-C> to quit:"
    echo "${qsub_command[@]}" -- "$cmd_arg" "$@"
    echo ""
    read
fi

# Initialise the .err and .out files so we do not have any
# appending from QSUB
echo -n "" > "$error_file"
echo -n "" > "$out_file"

# and run
"${qsub_command[@]}" "$cmd_arg" "$@"
