#!/bin/bash
# Small simple wrapper around qsub to make simple jobs run easily using openpbs
#                                     yoh@pscyhology.rutgers.edu

#rcfile=$HOME/.ravanizerc
#[ -f $rcfile ] && source $rcfile

queue=${PBS_DEFAULT_QUEUE:-long}
nodes=${PBS_DEFAULT_NODES:-1}
ppn=${PBS_DEFAULT_PPN:-1}
name=${PBS_DEFAULT_NAME:-`whoami`.$$}
lparams=""						# additional parameters for resourcelist
verbose=${PBS_DEFAULT_VERBOSE:-1}
make_directory=${PBS_DEFAULT_DIRECTORY:-0}
dry_run=0
background=0
ext='pbs'

append=

usage()
{
    cat <<EOF
Usage: $0 [OPTIONS] COMMAND [COMMAND_OPTIONS]

 is used to simplify scheduling of the tasks for PBS. Taking command
 line parameters and environment variables it creates a .pbs file
  which is ran by qsub

OPTIONS is an optional list of command line arguments for $0. The very
first argument after $0 which does not start with "-" is considered to be a
COMMAND. The rest of the arguments are considered to be optional
COMMAND_OPTIONS.

The description of the arguments for $0 follows the format:

-short, --long [value] description (environment_variable) {current value}

-q, --queue  s     queue to use (PBS_DEFAULT_QUEUE) {$queue}
-n, --nodes  n     number of nodes to request (PBS_DEFAULT_NODES) {$nodes}
-p, --ppn    n     processors per node (PBS_DEFAULT_PPN) {$ppn}
-l, --res    s     additional resources requested
-N, --name   s     name of the job (PBS_DEFAULT_NAME) {$name}
-v, --verbose      toggles verbosity (PBS_DEFAULT_VERBOSE) {$verbose}
-d, --directory    toggles either to create a uniquely named directory, in
                   which actual execution will happen
                   (PBS_DEFAULT_DIRECTORY) {$make_directory}
-D, --dry-run      dry run: pbs file is created but not ran {$dry_run}
-a, --append s     a line to append into a list of PBS rules,
                   can be used multiple times {$append}
-b, --background   do not run via pbs but rather in separate shell
-f, --foreground   do not run via pbs but rather the same shell
EOF

}

while [ $# -ge 0 ] && [ x${1:0:1} == x- ]
  do
  case $1 in
      -a|--append)
	  append="${append:+$append
}$2"
	  shift
	  ;;
      -q|--queue)
          queue=$2; shift
	  ;;
      -n|--nodes)
          nodes=$2; shift
	  ;;
      -l|--res)
          lparams="${lparams}:$2"; shift
	  ;;
      -p|--ppn)
          ppn=$2;  shift
	  ;;
      -N|--name)
          name=$2;  shift
	  ;;
      -v|--verbose)		# toggle verbosity
	  verbose=$((1-$verbose))
	  ;;
      -h|--help)
	  usage
	  exit 0
	  ;;
      -d|--directory)
	  make_directory=$((1-$make_directory))
	  ;;
      -D|--dry-run)
	  dry_run=1
	  ;;
      -b|--background)
	  background=1
	  ext='pbssh'
	  ;;
      -f|--foreground)
	  background=2
	  ext='pbssh'
	  ;;

      *)
	  echo "Unknown parameter for ravanize $1"
	  exit 1
  esac
  shift
done

if [ "x$1" = "x" ]
then
 echo "ERROR: Needs an argument - command to run"
 usage
 exit 1
fi

command="$@";
commandname=`basename ${1// */}`.$$

if [ x$make_directory == x1 ]
then
    mtaskid=pbs-$commandname-`date +%Y%m%d-%H:%M:%S:%N`
    [ x$verbose == x1 ] && echo "Task Directory: $mtaskid running $command"
    mkdir -p "$mtaskid"
    cd "$mtaskid"
fi

if [ x$verbose == x1 ]
then
    echo "Name: $name  Queue: $queue   Nodes: $nodes   PPN: $ppn"
    echo "Command: $command"
fi

cat  > $commandname.$ext  <<EOF
#PBS -N $name
#PBS -l nodes=$nodes:ppn=$ppn${lparams}
#PBS -q $queue
#PBS -o `pwd`
#PBS -j oe
#PBS -V
$append
cd `pwd`
echo "Running $command"
$command
EOF


if [ x$background == x1 ]
	then
	chmod +x $commandname.$ext
	sh -c "( DISPLAY='' $commandname.$ext >$commandname.$ext.OU 2>&1 & )"
elif [ x$background == x2 ]
	then
	chmod +x $commandname.$ext
	sh -c "DISPLAY='' $commandname.$ext >$commandname.$ext.OU 2>&1"
else
	if [ x$dry_run == x0 ]
		then
		rerun=1
		while [ x$rerun == x1 ]
		  do 
		  DISPLAY='' qsub $commandname.$ext && rerun=0 || ( echo 'Sleeping to rerun'; sleep 10 )
		done
	fi
fi