# Renata Kopecna


#########################################################
#							                            #
#	        Define Parser and add all options           #
#                                                       #
#########################################################
import sys
import argparse

RunList = [1,2,12]

class ShowArgumentsParser(argparse.ArgumentParser):
    def error(self, message):
        sys.stderr.write('error: %s\n\n' %message)
        parser.print_usage(sys.stderr)
        sys.stderr.write('\n'+self.description)
        sys.exit(2)

def addAllArgs(parser):
    #Basic options
    parser.add_argument("-v", "--verbosity", 
                        metavar = '<intVerbosity>',
                        default = 2, 
                        type = int, 
                        help="Set verbosity level. Default is INFO.\n"+
                             "Available levels are:\n"+
                             "0:  Trace\n"+
                             "1:  Debug\n"+
                             "2:  Info\n"+
                             "3:  Warn\n"+
                             "4:  Error\n"+
                             "5:  Critical\n\n")
    parser.add_argument("-log", "--log", 
                        metavar = '<strLog>',
                        default = "", 
                        type = str, 
                        help="Set log file. Default is no log, streaming right into console.\n\n")
    parser.add_argument("-hilfe", "--hilfe", 
                        action="store_true",
                        help="Print full help.\n\n")
    #Compile                        
    parser.add_argument("-compile", "--compile", 
                        action="store_true",
                        help="Compile the code before runing.\n\n")
    parser.add_argument("-dontCompile", "--dontCompile", 
                        action="store_true",
                        help="Don't compile the code before runing.\n\n")
    # #Set dataset                        
    parser.add_argument('-Run', '--Run', 
                        metavar='<Run>', 
                        default = 1, 
                        type = int, 
                        help="Set the Run number. Default is Run 1.\n\n")
    parser.add_argument('-allRun', '--allRun', 
                        action="store_true", 
                        default = False, 
                        help="Run on Runs 1, 2 and 12.\n\n")
    parser.add_argument('-year', '--year',  #TODO 
                        metavar='<year>', 
                        default = -1, 
                        type = int, 
                        help="Run only on one year.\n\n") #TODO
    parser.add_argument("-PHSP", "--PHSP",    
                        action="store_true", 
                        default = False, 
                        help="Run on PHSP\n\n")
    parser.add_argument("-Data", "--Data",    
                        action="store_true", 
                        default = False, 
                        help="Run on Data\n\n")
    parser.add_argument("-Ref", "--Ref",    
                        action="store_true", 
                        default = False, 
                        help="Run on reference channel Jpsi\n\n")
    parser.add_argument("-MC", "--MC",      
                        action="store_true", 
                        default = False, 
                        help="Run on signal MC\n\n")    
    parser.add_argument("-all", "--all", 
                        action="store_true", 
                        default = False,
                        help = "Run on data, signal MC, reference MC and PHSP\n\n" )
    parser.add_argument("-allMC", "--allMC", 
                        action="store_true", 
                        default = False,
                        help = "Run on signal MC, reference MC and PHSP\n\n" )
    parser.add_argument("-genMC", "--genMC", 
                        action="store_true", 
                        default = False, 
                        help="Run on generator level MC\n\n")    
    parser.add_argument("-boost", "--boost", 
                        action="store_true", 
                        default = False, 
                        help="Run on boosted sample, turn on also generator level MC option!\n\n")                          
    #Set tasks                        
    parser.add_argument("-convert", "--convert", 
                        action="store_true",
                        help="Convert selection tuples to FCNC tuples.\n\n") 
    parser.add_argument("-angCorr", "--angCorr", 
                        action="store_true",
                        help="Obtain angular acceptance correction coefficients from PHSP Monte Carlo events.\n\n") 
    parser.add_argument("-scan", "--scan", 
                        action="store_true",
                        help="Scan whatever you are doing :)\n\n") 
    parser.add_argument("-fit", "--fit", 
                        action="store_true",
                        help="Run the main fit\n\n")
    parser.add_argument("-angRes", "--angRes", 
                        action="store_true",
                        help="Get angular resolution from signal MC.\n\n") 
    parser.add_argument('-toys', '--toys',  #TODO  
                        action="store_true",
                        help="Generate toys. Add -evts to generate custom number of events.\n\n") 
    parser.add_argument('-script', '--script', 
                        action="store_true",
                        help="Run whatever script you decided.\n\n") 
    parser.add_argument('-job', '--job', 
                        metavar='<job>', 
                        default = -1, 
                        type = int, 
                        help="Set the job number\n\n")
    #Set toy options
    parser.add_argument('-evts', '--evts', 
                        metavar='<evts>', 
                        default = -1, 
                        type = int, 
                        help="Set number of used/generated events\n\n")

    #Set fit options                        
    parser.add_argument('-onlyBkg', '--onlyBkg', 
                        action="store_true",
                        help="Fit only angular background.\n\n") 
    parser.add_argument('-upper', '--upper', 
                        action="store_true",
                        help="Fit only upper-mass side-band.\n\n") 
    parser.add_argument('-lower', '--lower', 
                        action="store_true",
                        help="Fit only lower-mass side-band.\n\n") 
    parser.add_argument('-bin', '--bin', 
                        metavar='<bin>', 
                        default = -1, 
                        type = int, 
                        help="Which angular bin would you like to fit?\n\n") 
    parser.add_argument('-nBins', '--nBins', 
                        metavar='<nBins>', 
                        default = -1, 
                        type = int, 
                        help="Set total number q2 bins.\n\n") 
    parser.add_argument("-MagDown", "--MagDown",    
                        action="store_true", 
                        default = False, 
                        help="Run only on MagDown\n\n")
    parser.add_argument("-MagUp", "--MagUp",    
                        action="store_true", 
                        default = False, 
                        help="Run only on MagUp\n\n")
    parser.add_argument('-trueMC', '--trueMC', 
                        action="store_true",
                        help="Use MC true information.\n\n") 
    parser.add_argument('-likelyhood', '--likelyhood', 
                        action="store_true",
                        help="Likelihood profile scans.\n\n") 
    parser.add_argument('-massDim', '--massDim', 
                        action="store_true",
                        help="Fit mass dimension\n\n") 
    parser.add_argument('-MoM', '--MoM', 
                        action="store_true",
                        help="Use Method of Moments instead of the fit.\n\n") 
    parser.add_argument('-observeP', '--observeP', 
                        action="store_true",
                        help="use P(') angular observables in the fit instead of the S(')\n\n") 
    parser.add_argument('-folding', '--folding', 
                        metavar='<folding>', 
                        default = -1, 
                        type = int, 
                        help="Set the folding of angles.\n"
                         +"Five possibilities use range [0-4]\n"
                         +"\t0: phi' = phi + pi   for phi < 0\n"
                         +"\t1: phi' = -phi       for phi < 0\n"
                         +"\t   phi' = pi - phi   for ctl < 0\n"
                         +"\t   ctl' = -ctl       for ctl < 0\n"
                         +"\t2: phi' = -phi       for phi < 0\n"
                         +"\t   ctl' = -ctl       for ctl < 0\n"
                         +"\t3: phi' = pi - phi   for phi > pi/2\n"
                         +"\t   phi' = -pi - phi  for phi < -pi/2\n"
                         +"\t   ctl' = -ctl       for ctl < 0\n"
                         +"\t4: phi' = pi - phi   for phi > pi/2\n"
                         +"\t   phi' = -pi - phi  for phi < -pi/2\n"
                         +"\t   ctk' = -ctk       for ctk < 0\n"
                         +"\t   ctl' = -ctl       for ctl < 0\n"
                         +"Default is -1, meaning no folding.\n\n")
    parser.add_argument('-loopFolds', '--loopFolds',
                        action="store_true",
                        help="Evaluate all five foldings\n\n")
    # #Other options

    #parser.add_argument('-FC', '--FC',  #TODO 
    #                    action="store_true",
    #                    help="Run Feldman-Cousins: par=<int>/8 and q2bin=<int>%%8") #TODO

    #parser.add_argument('-pullsMC', '--pullsMC',  #TODO  
    #                    action="store_true",
    #                    help="Create -pullsMC pulls using specified signal MC events.") #TODO
    parser.add_argument('-index', '--index',  #TODO 
                        metavar='<index>', 
                        default = -1, 
                        type = int, 
                        help="Set index, used for various checks.")
    
    #Systematics
    #parser.add_argument('-systematics', '--systematics',  #TODO "-s <int>"
    #                    metavar='<systematics>', 
    #                    default = -1, 
    #                    type = int, 
    #                    help="Run a systematic study number:\n" 
    #                        +"\t1\t:Bootstrapping of PHSP MC\n"
    #                        +"\t2\t:Perform fit with non-symmetric acceptance in ctl\n"
    #                        +"\t3\t:Increase Legendre poly order by 2 for PHSP MC\n"
    #                        +"\t4\t:Randomly change PHSP MC reweighting wihtin it's uncertainties\n"
    #                        +"\t5\t:Check non-homogeneous FS distributions in q2\n"
    #                        +"\t6\t:Vary angles (ctk, ctl and phi) within the angular resoluation\n"
    #                        +"\t7\t:Generate toy events with double-gaussian profile and fit with CB\n"
    #                        +"\t8\t:Systematic study on angular background model, add 10%% contribution of higher order\n"
    #                        +"\t9\t:Investigate systematic effects due to trigger selection\n"
    #                        +"\t10\t:Reweight PHSP MC according to pion PT disagreement in DD\n"
    #                        +"\t11\t:Background systematic for mimicing the B0 -> KS mu mu veto\n")
    #Test parameter   
    parser.add_argument('-test', '--test',  #TODO 
                        metavar='<test>', 
                        default = -1, 
                        type = int, 
                        help="Test parameter. When !=-1, verbosity automatically on debug.\n\n") #TODO   

def getDatasetCommands(opts):
    arg_list = []
    if (opts.all):    return ["-d 0", "-d 1", "-d 2", "-d 3", "-d 4"]
    if (opts.allMC):  return ["-d 1", "-d 2", "-d 3", "-d 4"]

    if (opts.Data):  arg_list.append("-d 0")
    if (opts.genMC):
        if (opts.PHSP): arg_list.append("-d 4")
        else:           arg_list.append("-d 5")
    else:
        if (opts.MC):  
            if (opts.Ref):  arg_list.append("-d 2")
            else:           arg_list.append("-d 1")
        if (opts.PHSP):  arg_list.append("-d 3")
    return arg_list

def setDataOpts(opts,arg):
    if (arg == "-d 0"): 
        opts.Data   = True
        opts.MC     = False
        opts.PHSP   = False

    elif (arg == "-d 1"): 
        opts.Data   = False
        opts.MC     = True
        opts.Ref    = False
        opts.PHSP   = False

    if (arg == "-d 2"): 
        opts.Data   = False
        opts.MC     = True
        opts.Ref    = True
        opts.PHSP   = False

    if (arg == "-d 3"): 
        opts.Data   = False
        opts.MC     = True
        opts.PHSP   = True

def getActionCommands(opts):
    arg_list = []
    if (opts.test!=-1):  
        arg_list.append("-z " + str(opts.test))
        opts.verbosity = 1
    if (opts.angRes):   arg_list.append("-a")
    if (opts.angCorr):  
        arg_list.append("-c")
        if (opts.scan): arg_list.append("-i 1")       

    if (opts.fit): 
        if (opts.onlyBkg):    
            arg_list.append("-f 0")
            if (opts.upper): arg_list.append("-i 2")
            if (opts.lower): arg_list.append("-i 1")
        elif (opts.massDim):  arg_list.append("-f 2")
        elif (opts.MoM):      arg_list.append("-f 3")
        elif (opts.MC or opts.PHSP):       
            arg_list.append("-f 4") #TODO: I'm working on this one
            if (opts.MagDown): arg_list.append("-i 2")
            if (opts.MagUp): arg_list.append("-i 1")
        elif (opts.genMC):    arg_list.append("-f 5")
        else :                arg_list.append("-f 1")

    if (opts.toys):         arg_list.append("-w")
    if (opts.convert):      arg_list.append("-k")
    #if (opts.pullsMC):      arg_list.append("-m")  #or -t when form toys TODO
    #if (opts.systematics != -1):  arg_list.append("-s " + str(opts.systematics)) TODO
    if (opts.script):      arg_list.append("-x")

    if (opts.upper):    arg_list.append("") #TODO
    if (opts.lower):    arg_list.append("") #TODO
    if (opts.trueMC):   arg_list.append("") #TODO
    return arg_list


def getParameterCommands(opts):
    arg_list = []
    if (opts.bin!=-1):      arg_list.append("-b " + str(opts.bin))
    if (opts.evts!=-1):     arg_list.append("-e " + str(opts.evts))
    if (opts.folding!=-1):  arg_list.append("-g " + str(opts.folding))
    if (opts.loopFolds):    arg_list.append("-g 5")
    if (opts.likelyhood):   arg_list.append("-l")
    if (opts.observeP):     arg_list.append("-p")
    #if (opts.FC):           arg_list.append("-o") TODO
    if (opts.nBins!=-1):    arg_list.append("-u " + str(opts.nBins))
    if (opts.Ref):    arg_list.append("-q")
    if (opts.index !=-1):   arg_list.append("-i " + str(opts.index))
    if (opts.job!=-1):      arg_list.append("-j " + str(opts.job))
    
    return arg_list

def fullArgListPerRun(opts):
    runs = []
    if (opts.allRun): runs = RunList
    else: runs = [opts.Run]

       
    argListList = []
    for dataArg in getDatasetCommands(opts):
        for r in runs:
            arg_list = []
            arg_list.append('-v ' + str(opts.verbosity))
            setDataOpts(opts,dataArg)
            if (opts.Run!=-1):     arg_list.append("-r " + str(r))
            if (opts.year!=-1):    arg_list.append('-y' + str(opts.year))
            if (len(arg_list)==0): print ("The script needs at least run or year! Crashing now")
            arg_list.extend([dataArg])
            arg_list.extend(getActionCommands(opts))
            arg_list.extend(getParameterCommands(opts))
            argListList.append(arg_list)

    return argListList

