#Removal of multiple cnadidates from the sample
#Renata Kopecna

import sys
import os
import argparse
from RemoveMultipleCandidatesSource import *
sys.path.insert(0,'/home/lhcb/kopecna/B2KstarMuMu_clean/Code/Selection/ComparisonTool')
#import Utils
from Utils3 import checkYearSample
import ROOT

'''
Tool for removing multiple candidates after TMVA.
Adds a branch with a boolean for each TMVA cut, whether the event should be removed or not

'''

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)


if '__main__' == __name__:
    print ("")

    parser = ShowArgumentsParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        prog=os.path.basename(sys.argv[0]),
        description=("""
        For a full list of arguments, do: 'python {0} -h'
        """).format(os.path.basename(sys.argv[0]))
        )

  ## add the arguments
    parser.add_argument("-v", "--verbose", action="store_true",help="increase output verbosity")

    parser.add_argument('-year', metavar='<year>', default = "2011", type = str, help="Sets the year")
    parser.add_argument("-Data",    "--Data",    action="store_true", default = False, help="Data")
    parser.add_argument("-MC",      "--MC",      action="store_true", default = False, help="MC")
    parser.add_argument("-TM",      "--TM",      action="store_true", default = False, help="TM")
    parser.add_argument("-TMpid",   "--TMpid",   action="store_true", default = False, help="TMpid")
    parser.add_argument("-gammaTM", "--gammaTM", action="store_true", default = False, help="gammaTM")
    parser.add_argument("-RefChan", "--RefChan", action="store_true", default = False, help="ReferenceChannel")
    parser.add_argument("-PHSP",    "--PHSP",    action="store_true", default = False, help="PHSP")
    parser.add_argument('-Run', metavar='<Run>', default = 1, type = int, help="Sets the run")
    parser.add_argument("-all",      "--all",    action="store_true", default = False, help="Run over all years")
   
    ROOT.PyConfig.IgnoreCommandLineOptions = True
    opts = parser.parse_args()
    if (opts.verbose): print ("[DEBUG]\t",opts)

    allYearsDict ={ 2011:1,
                    2012:1,
                    2015:2,  
                    2016:2, 
                    2017:2, 
                    2018:2,
                    }    


    if (opts.all):
        for yr in allYearsDict.keys():
            isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=True,  MC=False,TM=False, TMpid=False, gammaTM=False, ReferenceChannel=False, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)

            isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=False, gammaTM=False, ReferenceChannel=False, PHSP=True, KshortDecaysInVelo=True,verbose=opts.verbose)
            #isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=True,  TMpid=False, gammaTM=False, ReferenceChannel=False, PHSP=True, KshortDecaysInVelo=True,verbose=opts.verbose)
            isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=True,  gammaTM=False, ReferenceChannel=False, PHSP=True, KshortDecaysInVelo=True,verbose=opts.verbose)
            #isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=True,  gammaTM=True,  ReferenceChannel=False, PHSP=True, KshortDecaysInVelo=True,verbose=opts.verbose)

            if (yr != 2015):
                    isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=False, gammaTM=False, ReferenceChannel=False, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)
                    #isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=True,  TMpid=False, gammaTM=False, ReferenceChannel=False, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)
                    isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=True,  gammaTM=False, ReferenceChannel=False, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)
                    #isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=True,  gammaTM=True,  ReferenceChannel=False, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)

            if (yr < 2017):        
                isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=False, gammaTM=False, ReferenceChannel=True, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)
                #isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=True,  TMpid=False, gammaTM=False, ReferenceChannel=True, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)
                isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=True,  gammaTM=False, ReferenceChannel=True, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)
                #isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=True,  gammaTM=True,  ReferenceChannel=True, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)

    else:             
        #ONLY FOR INCLUSIVE SAMPLE
        for yr in allYearsDict.keys():
            if (yr != 2015 and yr <2017):
                isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=False, gammaTM=False, ReferenceChannel=False, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)            
                isMultipleCandidate(year=yr, Run=allYearsDict[yr], Data=False, MC=True, TM=False, TMpid=True,  gammaTM=False, ReferenceChannel=False, PHSP=False, KshortDecaysInVelo=True,verbose=opts.verbose)