#!/usr/bin/env python2

"""Submit multiple IcedustControl jobs to a batch and modify as neccesary the input config files
   """

import optparse
import subprocess
import sys
import tempfile
import os
import time

# Parse the command line options
description="""
Helper script to launch RunIcedustControl on a batch system.
The config file that's given to this is used 
"""
parser = optparse.OptionParser(description="Submit parametric RunIcedustControl jobs to a batch system.")

icedust_group = optparse.OptionGroup(parser, "Options passed to the RunIcedustControl command")
icedust_group.add_option("-c","--configfile",action="store",dest="configfile",
                  type="string",default = "comet.cfg",
                  help="Specify the name of the cards file for the comet job")
icedust_group.add_option("-r","--runnumber",action="store",dest="number",type="string",
                  help="Provide a number of up to 8 characters to give output "\
                  "files a unique name. This will be encorporated into the output "\
                  "files. If this option is used it will override the value in the cards file.")
icedust_group.add_option("-u","--comment",action="store",dest="comment",type="string",
                  help="Provide a comment to be used in the last field of the filename. "\
                  "This will be truncated if the file name exceeds 80 characters")
icedust_group.add_option("-t","--tmpdir",action="store",dest="tmpdir",
                  type="string",default = "/tmp",
                  help="Specify the directory for temporary files. [default: %default]")
parser.add_option_group(icedust_group)

qsub_group = optparse.OptionGroup(parser, "Options passed to qsub")
qsub_group.add_option("-n","--jobs",action="store",dest="n_jobs",
                  type="int",default = "10",
                  help="Specify the number of tasks to run for this job. [default: %default]")
qsub_group.add_option("-o","--batch-stream",action="store",dest="batch_stream", default=None,
                  help="Specify the output + error stream destination. [default: %default]")
qsub_group.add_option("-O","--qsub-option",action="append",dest="qsub_opts",
                  help="Provide commands to the underlying qsub command")
parser.add_option_group(qsub_group)

(options,args) = parser.parse_args()

#################################################################
# Check for the ICEDUST environment
needed_vars = ["ICEDUST_HOME", "ICEDUST_PROJECT","CMTROOT","CMTPATH"]
for var in needed_vars:
   if not os.environ.get(var):
       print "ERROR: Please source the appropriate setup.sh file to configure the ICEDUST software"
       sys.exit(1)

environ=os.environ.get
cmtSetupScript=os.path.join(environ('CMTROOT'),"mgr","setup.sh")
projectSetupScript=os.path.join(environ('ICEDUST_HOME'),environ('ICEDUST_PROJECT'),"cmt","setup.sh")

if not os.path.isfile(cmtSetupScript):
   print "ERROR: CMT is not configured, make sure you've installed ICEDUST properly"
   sys.exit(1)

if not os.path.isfile(projectSetupScript):
   print "ERROR: Cannot find cmt/setup.sh file for project '"+environ('ICEDUST_PROJECT')+"'.  Make sure you've installed ICEDUST properly"
   sys.exit(1)

#################################################################
# Make the run file
job="time RunIcedustControl " 
if options.configfile:
    job+="-c '"+options.configfile+"' "
if options.number:
    job+="-r '"+options.number+"' "
if options.comment:
    job+="-u '"+options.comment+"' "
if options.tmpdir:
    job+="-t '"+options.tmpdir+"' "

batchFile=""
try:
    batch=tempfile.NamedTemporaryFile(prefix=options.tmpdir+"/",delete=False)
    batchFile=batch.name
    batch.write("#!/bin/bash\n")
    batch.write("unset LD_LIBRARY_PATH\n")
    if environ('LD_LIBRARY_PATH'):
        batch.write("export LD_LIBRARY_PATH="+environ('LD_LIBRARY_PATH')+"\n")
    batch.write("source "+cmtSetupScript+"\n")
    batch.write("export CMTPATH="+environ('CMTPATH')+"\n")
    batch.write("source "+projectSetupScript+"\n")
    batch.write("cd "+os.getcwd()+"\n")
    batch.write(job+"\n")
    batch.close()
except:
    print "Cannot write batch file "
    raise 


#################################################################
# Submit jobs to the batch system
jobname="RunIcedustControl"
if options.configfile:
   jobname=os.path.basename(options.configfile)
   jobname=os.path.splitext(jobname)[0]

command = ["qsub","-V","-cwd"]
command+=["-N",jobname]
if options.batch_stream:
   command += ["-e",options.batch_stream]
   command += ["-o",options.batch_stream]
if options.n_jobs:
   command += ["-t","1-"+str(options.n_jobs)]
if options.qsub_opts:
   for qopt in options.qsub_opts:
       command += qopt.split(None,1)
command += [batchFile]

print "Command line is: " + " ".join(command)
print "Submitting job in:"
for i in range(5):
    sys.stdout.write(str(i)+" ")
    sys.stdout.flush()
    time.sleep(0.3)
print

proc = subprocess.Popen(command)

# Run teh job.
rtc = proc.wait()

sys.exit(rtc)