#! /usr/bin/python
import os
import sys
import getopt
import time

#----------------------------------------------------------------------

def printHeader(file_out):
    print >> file_out,"# original version produced by " + os.path.basename(sys.argv[0]) + " on " + time.asctime()
    print >> file_out,"# command:"
    print >> file_out,"#   " + os.path.basename(sys.argv[0]) + " " + " ".join(sys.argv[1:])
    print >> file_out

#----------------------------------------------------------------------
# main
#----------------------------------------------------------------------

usage = """
job_maker [options]

--data/MC/particleGun              select data, MC or particleGun
--fastsim                          works with fastsim eventcontent (actually not the standard one ;-))
--AOD/RECO                         AOD should be standard
--caf                              send jobs using CAF (hadoop is std in SD)
--server                           use server to submit crab jobs
--AddPdfWeight                     add pdf weight information into the ntuple (warning: very very slow)
--energyCorrFromDB                 read energy corrections from DB (default is read from ROOT file)
--skimworz                         skim for T&P
--diphoskim                        DiPhoton HLT Skim
--pjetskim                         Do PhotonPlusJet Skim (Works for both Data and MC)
--vlpresel                         Very Loose Preselection on 2 photons
--noskim                           no skim performed (typically only for signal)
--skimmm                           skim for DYmumu
--skimmmg                          skim for mumugamma
--skim1el                          skim to select only one electron
--storage=eos/castor/hadoop        use eos/castor/hadoop as storage element
--name=                            name of the job, directory, files (without a default will be made)
--append=                          appends the name of the task
--scheduler=                       glite or glidein
--datasetpath=                     datasetpath
--runselection=                    comma separated list of runs
--events=                          number of events to process (-1 is standard)
--events_per_job=                  number of jobs per event (10000 is standard)
--json=""                          give the path of the json file to use (data only)
--lumis=                           equivalent to events in case of JSON file (it will be ignored if --json not set)
--lumis_per_job=                   equivalent to events_per_job in case of JSON file (it will be ignored if --json not set) 
--outputdir=                       output dir in the storage element (e.g. matteo/data/2010 in hadoop or /s/sani/data/2010  in castor)
--globaltag=                       change globaltag
--hlttag=                          set the HLT tag (HLT is standard, REGIGI3YX is frequent)
--blacklist=                       sites to black list (e.g. \"T2_US_UCSD,ultralight.org,hpc.ufl.edu\")
--whitelist=                       sites to white list (e.g. \"T2_US_UCSD,ultralight.org,hpc.ufl.edu\")
--useT1                            edits the crab.cfg so that T1 access is possible

--enable=                          enable processing of specified objects (comma separated list)
--disable=                         disable processing of specified objects (comma separated list)
--published                        to run on local dbs published dataset
"""


# default selection
flagData = False
flagMC = False
flagMyPreSelection = False
flagSkimDiphoton = False
flagSkimPJet = False
flagSkimworz = False
flagSkim1El = False
flagVLPreSelection = False
flagNoSkim = False
flagMuMuSkim = False
flagMMgSkim = False
flagSkimHmm = False
flagSkimHee = False
flagSkimMu  = False
flagAOD = False
flagRECO = False
flagPG = False
flagFastSim = False
flagT1 = False
flagDB = False
flagPublished = False

name = ""
datasetpath=""
append=""
runselection = [""]
total_number = -1
events_job = 10000
storage = "hadoop"
email = ""
useCAF = False
user_remote_dir = ""
globaltag = ""
hlttag = ""
scheduler = "glite"
json = ""
lumis = -1
lumis_per_job=10
blackList=""
whiteList=""
useServer = False
toEnable = []
toDisable = []
jecFromSqlite = False

(opts, args) = getopt.getopt(sys.argv[1:], '', ['data', 'MC', 'AOD', 'RECO', 'diphoskim', 'pjetskim', 'vlpresel', 'storage=',  'name=', 'datasetpath=', 'append=', 'outputdir=','runselection=', 'events=', 'events_per_job=', 'help', 'globaltag=', 'noskim', 'json=', 'lumis=', 'lumis_per_job=', 'caf','hlttag=', 'whitelist=', 'blacklist=', 'skimmmg', 'skimworz', 'useT1', 'server', 'skim1el', 'AddPdfWeight', 'scheduler=', 'fastsim', 'enable=', 'disable=', 'mypresel', 'particleGun', 'energyCorrFromDB', 'jecFromSqlite', 'published', 'skimhee', 'skimhmm', 'skimmu'])

for opt,argm in opts:
    if (opt == "--help"):
        print 'Usage: %s' % (usage)
        sys.exit(0)
    if (opt == "--data"):
        flagData = True
    elif (opt == "--MC"):
        flagMC = True
    elif (opt == "--fastsim"):
        flagFastSim = True
    elif (opt == "--AOD"):
        flagAOD = True
    elif (opt == "--particleGun"):
        flagPG = True
    elif (opt == "--RECO"):
        flagRECO = True
    elif (opt == "--server"):
        useServer = True
    elif (opt == "--caf"):
        useCAF = True
    elif (opt == "--diphoskim"):
        flagSkimDiphoton = True
    elif (opt == "--pjetskim"):
         flagSkimPJet = True
    elif (opt == "--vlpresel"):
        flagVLPreSelection = True
    elif (opt == "--mypresel"):
        flagMyPreSelection = True
    elif (opt == "--skimhmm"):
        flagSkimHmm = True
    elif (opt == "--skimmu"):
        flagSkimMu = True
    elif (opt == "--skimhee"):
        flagSkimHee = True        
    elif (opt == "--skimmm"):
        flagMuMuSkim = True
    elif (opt == "--skimmmg"):
        flagMMgSkim = True
    elif (opt == "--skimworz"):
        flagSkimworz = True
    elif (opt == "--skim1el"):
        flagSkim1El = True
    elif (opt == "--noskim"):
        flagNoSkim = True
    elif (opt == "--storage"):
        storage = argm
    elif (opt == "--energyCorrFromDB"):
        flagDB = True
    elif (opt == "--name"):
        name = argm
    elif (opt == "--datasetpath"):
        datasetpath = argm
    elif (opt == "--append"):
        append = argm
    elif (opt == "--runselection"):
        runselection = argm.split(",")
    elif (opt == "--events"):
        total_number = int(argm)
    elif (opt == "--events_per_job"):
        events_job = int(argm)
    elif (opt == "--outputdir"):
        user_remote_dir = argm
    elif (opt == "--globaltag"):
        globaltag = argm
    elif (opt == "--json"):
        json = argm
    elif (opt == "--lumis"):
        lumis = int(argm)
    elif (opt == "--lumis_per_job"):
        lumis_per_job = int(argm)
    elif (opt == "--hlttag"):
        hlttag = argm
    elif (opt == "--whitelist"):
        whiteList = argm
    elif (opt == "--blacklist"):
        blackList = argm
    elif (opt == "--useT1"):
        flagT1 = True
    elif (opt == "--AddPdfWeight"):
        AddPdfWeight = True
    elif (opt == "--scheduler"):
        scheduler = argm
    elif (opt == "--enable"):
        toEnable = argm.split(",")
    elif (opt == "--disable"):
        toDisable = argm.split(",")
    elif (opt == "--jecFromSqlite"):
        jecFromSqlite = True
    elif (opt == "--published"):
        flagPublished = True
    else:
        print 'Wrong options: %s' % (opt)
        sys.exit(3)
        
if (name == ""):
    parts = datasetpath.split("/")
    name = parts[1]+"_"+parts[2]

if (append != ""):
    name=name+'_'+append


if (globaltag == ""):
    print "WARNING: you are using default GLOBAL-TAG, it may not be correct."

if (not(flagNoSkim ^ flagSkimDiphoton ^ flagSkimPJet ^ flagMMgSkim ^ flagVLPreSelection ^ flagSkim1El ^ flagSkimworz ^ flagMyPreSelection ^ flagSkimHmm ^ flagSkimHee ^ flagSkimMu)):
    print "You must have one and only one of the skims."
    print "Try again when you have chosen one:  --skim2pho or --noskim or --skimmmg or --skimworz."
    sys.exit(1)

if (not (flagRECO ^ flagAOD)):
    print "You must have one and only one of these:"
    print "flagRECO is "+str(flagRECO)+"\tflagAOD is "+str(flagAOD)
    print "Try again when you have chosen one:  --RECO or --AOD."
    sys.exit(1)

if (not (flagData ^ flagMC ^ flagPG)):
    print "Try again when you have chosen one:  --Data, --MC or --particleGun."
    sys.exit(1)

for run in runselection:
    job_name = name
        
    file = open("template_61X_cfg.py")
    lines = file.readlines()
    file.close()
    if (not os.path.exists("../test")):
        os.mkdir("../test")
    if (not os.path.exists("../test/crabTasks")):
        os.mkdir("../test/crabTasks")
    if (not os.path.exists("../test/crabTasks/"+job_name)):
        os.mkdir("../test/crabTasks/"+job_name)
    file_out = open("../test/crabTasks/" + job_name + "/" + job_name + "_cfg.py", "w")
    printHeader(file_out)
    for line in lines:
        if ("RootFileName" in line):
            file_out.write("process.h2ganalyzer.RootFileName = \'"+job_name+".root\'\n")
        elif("jobMaker = 'jobmaker unknown'" in line):
            file_out.write("jobMaker = \'" + os.path.basename(sys.argv[0]) + " " + " ".join(sys.argv[1:])+"\'\n")
        elif ("flagData = 'OFF'" in line and flagData):
            file_out.write("flagData = \'ON\'\n")
        elif ("flagMC = 'OFF'" in line and flagMC):
            file_out.write("flagMC = \'ON\'\n")
        elif ("flagAOD = 'OFF'" in line and flagAOD):
            file_out.write("flagAOD = \'ON\'\n")
        elif ("flagFastSim = 'OFF'" in line and flagFastSim):
            file_out.write("flagFastSim = \'ON\'\n")
        elif ("flagAOD = 'ON'" in line and not flagAOD):
            file_out.write("flagAOD = \'OFF\'\n")
        elif ("flagPG = 'OFF'" in line and flagPG):
            file_out.write("flagPG = \'ON\'\n")
        elif ("flagNoSkim = 'OFF'" in line and flagNoSkim):
            file_out.write("flagNoSkim = \'ON\'\n")
        elif ("flagSkimworz = 'OFF'" in line and flagSkimworz):
            file_out.write("flagSkimworz = \'ON\'\n")
        elif ("flagSkim1El = 'OFF'" in line and flagSkim1El):
            file_out.write("flagSkim1El = \'ON\'\n")
        elif ("flagSkimDiphoton = 'OFF'" in line and flagSkimDiphoton):
            file_out.write("flagSkimDiphoton = \'ON\'\n")
        elif ("flagSkimHmm = 'OFF'" in line and flagSkimHmm):
            file_out.write("flagSkimHmm = \'ON\'\n")
        elif ("flagSkimHee = 'OFF'" in line and flagSkimHee):
            file_out.write("flagSkimHee = \'ON\'\n")
        elif ("flagSkimMu = 'OFF'" in line and flagSkimMu):
            file_out.write("flagSkimMu = \'ON\'\n")            
        elif ("flagSkimPJet = 'OFF'" in line and flagSkimPJet):
            file_out.write("flagSkimPJet = \'ON\'\n")
        elif ("flagVLPreSelection = 'OFF'" in line and flagVLPreSelection):
            file_out.write("flagVLPreSelection = \'ON\'\n")
        elif ("flagMyPreSelection = 'OFF'" in line and flagMyPreSelection):
            file_out.write("flagMyPreSelection = \'ON\'\n")            
        elif ("flagMuMuSkim = 'OFF'" in line and flagMuMuSkim):
            file_out.write("flagMuMuSkim = \'ON\'\n")
        elif ("flagMMgSkim = 'OFF'" in line and flagMMgSkim):
            file_out.write("flagMMgSkim = \'ON\'\n")
        elif ("hltLabel = \"HLT\"" in line and hlttag != ""):
            file_out.write("hltLabel = \'"+hlttag+"\'\n")
        elif ("GlobalTag.globaltag" in line and globaltag != ""):
            file_out.write("process.GlobalTag.globaltag = \""+globaltag+"\"\n")
        else:
            file_out.write(line)
    file_out.write("\n")
    for t in toEnable:
        file_out.write("process.h2ganalyzer." + t + " = True\n")
    file_out.write("\n")
    for t in toDisable:
        file_out.write("process.h2ganalyzer." + t + " = False\n")
    if flagDB:
        file_out.write("process.h2ganalyzer.energyCorrectionsFromDB = True\n")
    if jecFromSqlite:
        file_out.write("from CMGTools.External.jec_2012_cff import use2012JecPreview\n")
        file_out.write("use2012JecPreview(process)\n")
    file_out.close()
    
    # setup crab cfg file
    file = open("crab_template.cfg")
    lines = file.readlines()
    file.close()
    file_out = open("../test/crabTasks/" + job_name + "/crab_"+job_name+".cfg", "w")

    printHeader(file_out)

    for line in lines:
        if ("runselection" in line and flagData and run is not ""):
            file_out.write("runselection="+run+"\n")
        elif ("scheduler" in line and useCAF):
            file_out.write("scheduler = caf\n")
        elif ("scheduler" in line and not useCAF):
            file_out.write("scheduler = "+scheduler)
        elif ("use_server" in line and useServer):
            file_out.write("use_server = 1\n")
        elif ("datasetpath" in line):
            file_out.write("datasetpath="+datasetpath+"\n")
        elif ("dbs_url" in line):
            if flagPublished:
                file_out.write("dbs_url = http://cmsdbsprod.cern.ch/cms_dbs_ph_analysis_01/servlet/DBSServlet\n")
        elif ("pset" in line):
            file_out.write("pset = "+job_name+"_cfg.py\n")    
        elif ("output_file" in line):
            file_out.write("output_file = "+job_name+".root\n")
        elif ("working_dir" in line):
            file_out.write("ui_working_dir = " + os.getcwd()+"/../test/crabTasks/"+ job_name + "/"+ job_name + "\n")
        elif ("storage_element" in line and storage == "castor"):
            file_out.write("storage_element = srm-cms.cern.ch\n")
        elif ("storage_element" in line and useCAF):
            file_out.write("storage_element = T2_CH_CAF\n")
        elif ("storage_path" in line and storage == "castor"):
            file_out.write("storage_path = /srm/managerv2?SFN=/castor/cern.ch/user/\n")
        elif ("user_remote_dir" in line and not useCAF):
            file_out.write("user_remote_dir=/"+user_remote_dir+"/"+job_name +"\n")
        elif ("user_remote_dir" in line and useCAF):
            file_out.write("user_remote_dir=/"+job_name +"\n")
        elif ("storage_element" in line and storage == "hadoop"):
            file_out.write("storage_element = bsrm-3.t2.ucsd.edu\n")
        elif ("storage_path" in line and storage == "hadoop"):
            file_out.write("storage_path=/srm/v2/server?SFN=/hadoop/cms/store/user/\n")
        elif ("storage_element" in line and storage == "eos"):
            file_out.write("storage_element = srm-eoscms.cern.ch\n")
        elif ("storage_path" in line and storage == "eos"):
            file_out.write("storage_path=/srm/v2/server?SFN=/eos/cms/store\n")
        elif ("total_number_of_events" in line and json == ""):
            file_out.write("total_number_of_events = "+str(total_number)+"\n")
        elif ("events_per_job" in line and json == ""):
            file_out.write("events_per_job = "+str(events_job)+"\n")
        elif ("lumi_mask" in line and json != ""):
            if json.startswith('/'):
                # absolute path specified
                file_out.write("lumi_mask = "+str(json)+"\n")
            else:
                # relative path specified
                file_out.write("lumi_mask = ../../../scripts/"+str(json)+"\n")
        elif ("lumis_per_job" in line and json != ""):
            file_out.write("lumis_per_job = "+str(lumis_per_job)+"\n")
        elif ("total_number_of_lumis" in line and json != ""):
            file_out.write("total_number_of_lumis = "+str(lumis)+"\n")
        elif ("se_white_list" in line and whiteList != ""):
            file_out.write("se_white_list = "+str(whiteList)+"\n")
        elif ("se_black_list" in line and blackList != ""):
            file_out.write("se_black_list = T0, T1,"+str(blackList)+"\n")
        elif ("show_prod" in line and flagT1):
            file_out.write("show_prod=1 \n")
        elif ("role" in line and flagT1):
            file_out.write("role=t1access \n")
        elif ("remove_default_blacklist" in line and flagT1):
            file_out.write("remove_default_blacklist=1 \n")
        else:
            file_out.write(line)
    file_out.close()

    print "cd ../test/crabTasks/"+job_name
    print "rm -r "+job_name
    print "crab -cfg crab_"+job_name+".cfg -create -submit"
    print "cd -"

    checker = open("../test/crabTasks/checker", "a+")
    checker.write("crab -c "+job_name+"/"+job_name+" -status\n")
    checker.write("crab -c "+job_name+"/"+job_name+" -get &> getting.log &\n")
    checker.close()
