#!/usr/bin/env python
import os, sys

def expand(t):
  if t.count('x'):
    pb = t.find('(')
    pe = t.find(')')
    return [t[:pb]] * int(t[pb+2:pe])
  return [t]

def main():
    if 'SLURM_JOBID' not in os.environ:
        print "You need to be running a slurm job for dispatch to work."
        sys.exit(1)
    args = sys.argv[1:]
    if len(args) < 2 or args[0] != '-r' or len(args) > 2:
        print "Only a very simple dispatch replacement is avaiable."
        print "There is only one supported format: `dispatch -r commandfile`"
        sys.exit(1)
    commandfile = args[1]
    if os.environ.get('SLURM_PROCID', '0') != '0':
        return
    # example 12(x2),13
    core_str = os.environ['SLURM_JOB_CPUS_PER_NODE']
    ncores = 0
    for part in core_str.split(","):
        ncores += sum(map(int, expand(part)))
    os.system("srun -c 1 --ntasks %i _slurm_dispatch_helper '%s'" % (ncores, commandfile))

if __name__ == "__main__":
    main()
