#!/usr/bin/bash

if [ "$1" == "--help" ]; then
  echo "================================================================================================="
  echo ""
  echo "  scounts - SLURM job counts. Computes the number of jobs completed by a user, group, or account."
  echo ""
  echo "  Usage: scounts --user=[user]|--group=[group]|--account=[account] [start_date_mmddyy]]          "
  echo ""
  echo "  Note: for industry accounts, start_date must be 010112 or after 073015                         "
  echo "================================================================================================="

  exit
fi

# determine STUBL install location
# Copied from Apache Ant:
# https://git-wip-us.apache.org/repos/asf?p=ant.git;a=blob;f=src/script/ant;h=b5ed5be6a8fe3a08d26dea53ea0fb3f5fab45e3f
if [ -z "$STUBL_HOME" -o ! -d "$STUBL_HOME" ] ; then
  ## resolve links - $0 may be a link to stubl's home
  PRG="$0"
  progname=`basename "$0"`

  # need this for relative symlinks
  while [ -h "$PRG" ] ; do
    ls=`ls -ld "$PRG"`
    link=`expr "$ls" : '.*-> \(.*\)$'`
    if expr "$link" : '/.*' > /dev/null; then
    PRG="$link"
    else
    PRG=`dirname "$PRG"`"/$link"
    fi
  done

  STUBL_HOME=`dirname "$PRG"`/..

  # make it fully qualified
  STUBL_HOME=`cd "$STUBL_HOME" > /dev/null && pwd`
fi

export STUBL_HOME=$STUBL_HOME

# setup STUBL environment
. $STUBL_HOME/conf/stubl 

# Location of helper scripts
MYDIR=$STUBL_HOME/bin/sausage_helpers

LISTGROUP=$MYDIR/ListGroup.sh

USEREXCLUDE=$STUBL_USER_EXCLUDE
PRESUM_FILE=$STUBL_HOME/$STUBL_PRESUM_FILE

if [ "$1" == "" ]; then
  echo "Please specify a user group (--group=) or account (--account=)!"
  exit
fi

option=`echo $1 | cut -d= -f1`

if [ "$2" == "" ]; then
  start=010112
else
  start=$2
fi

# only aggretate values are available prior to 07/30/2015
# so don't allow intermediate start dates
bIndustry=`echo "$1" | cut -d= -f2 | grep '^pi-' | wc -l`
if [ "$bIndustry" == "1" ]; then
  if [ "$start" -gt "010112" -a "$start" -lt "073015" ]; then
    echo "Start date must be either 010112 or later than 073015"
    exit
  fi
fi

if [ "$option" == "--group" ]; then
  acct=""
  grp=`echo $1 | cut -d= -f2`
  ssh -q vortex "export STUBL_HOME=$STUBL_HOME; $LISTGROUP $grp" | grep -v $USEREXCLUDE | awk '{ print $2 }' > tmp1
  ssh -q presto "export STUBL_HOME=$STUBL_HOME; $LISTGROUP $grp" | grep -v $USEREXCLUDE | awk '{ print $2 }' >> tmp1
  userList=`sort -u tmp1 | tr '\n' ' '`
  rm -f tmp1
elif [ "$option" == "--account" ]; then
  acct=`echo $1 | cut -d= -f2`
  ssh -q vortex "sacctmgr show associations account=$acct format=Cluster,Account,User" | awk '{ print $3 }' | sed -n '4,$p' | sort -u | grep -v '^$' > tmp1
  ssh -q presto "sacctmgr show associations account=$acct format=Cluster,Account,User" | awk '{ print $3 }' | sed -n '4,$p' | sort -u | grep -v '^$' >> tmp1
  userList=`sort -u tmp1 | tr '\n' ' '`
  rm -f tmp1

  prenumEntry=`grep "$acct" $PRESUM_FILE`

  if [ "$prenumEntry" != "" ]; then
    prenum=`echo $prenumEntry | cut -d: -f3`
  else
    prenum=0
  fi

  acct="--accounts=$acct"
elif [ "$option" == "--user" ]; then
  userList=`echo $1 | cut -d= -f2`
else
  acct=""
  grp=$1
  ssh -q vortex "export STUBL_HOME=$STUBL_HOME; $LISTGROUP $grp" | grep -v $USEREXCLUDE | awk '{ print $2 }' > tmp1
  ssh -q presto "export STUBL_HOME=$STUBL_HOME; $LISTGROUP $grp" | grep -v $USEREXCLUDE | awk '{ print $2 }' >> tmp1
  userList=`sort -u tmp1 | tr '\n' ' '`
  rm -f tmp1
fi
echo "users = $userList"

grandSum=0
sum=0
for u in $userList; do
  nj=`ssh -q vortex "module load stubl; slogs $start $u -X $acct" | awk '{ print $1 }' | grep '^[0-9]' | wc -l`
  sum=`expr $sum + $nj`
done
echo "Jobs run on vortex since $start : $sum"
grandSum=$sum

sum=0
for u in $userList; do
  nj=`ssh -q presto "module load stubl; slogs $start $u --clusters=industry -X $acct" | awk '{ print $1 }' | grep '^[0-9]' | wc -l`
  sum=`expr $sum + $nj`
done

if [ "$start" -le "010112" ]; then
   sum=`echo "$sum $prenum" | awk '{ printf("%d\n", $1+$2); }'`
fi

echo "Jobs run on presto since $start : $sum"

grandSum=`expr $sum + $grandSum`
echo "Grand Total : $grandSum"

