
from compareUltimateSource import *


'''
Tool for comparing two variables from two datasets.
Yes, writting it in Python3 would be smarter, 
but default is Python2 on all servers so far
'''
def compareUltimateAllPolarities(variable1, variable2,
                     year1, year2, cut1, cut2 ,
                     MC1, TM1, ReferenceChannel1, PHSP1, Preselected1, BDTed1,
                     MC2, TM2, ReferenceChannel2, PHSP2, Preselected2, BDTed2,
                     sWeighted1, bWeighted1, b2Dweighted1,
                     sWeighted2, bWeighted2, b2Dweighted2,
                     bPrint,   KshortDecaysInVelo1, KshortDecaysInVelo2, verbose):
    for magnet in polarityList():                 
        compareUltimate(variable1 = variable1, variable2 = variable2,
                     year1 = year1, year2 = year2, cut1 = cut1, cut2 = cut2,
                     magnet1 = magnet, MC1 = MC1, TM1 = TM1, ReferenceChannel1 = ReferenceChannel1, PHSP1 = PHSP1, Preselected1 = Preselected1, BDTed1 = BDTed1,
                     magnet2 = magnet, MC2 = MC2, TM2 = TM2, ReferenceChannel2 = ReferenceChannel2, PHSP2 = PHSP2, Preselected2 = Preselected2, BDTed2 = BDTed2,
                     sWeighted1 = sWeighted1, bWeighted1 = bWeighted1, b2Dweighted1 = b2Dweighted1, weightBranch1 = firstMCweight(),
                     sWeighted2 = sWeighted2, bWeighted2 = bWeighted2, b2Dweighted2 = b2Dweighted2,  weightBranch2 = firstMCweight(),
                     bPrint = bPrint, KshortDecaysInVelo1 = KshortDecaysInVelo1, KshortDecaysInVelo2 = KshortDecaysInVelo2, verbose = verbose)
   
def compareUltimateAllPolAndYears( variable1, variable2,
                     cut1, cut2 ,
                     MC1, TM1, ReferenceChannel1, PHSP1, Preselected1, BDTed1,
                     MC2, TM2, ReferenceChannel2, PHSP2, Preselected2, BDTed2,
                     sWeighted1, bWeighted1, b2Dweighted1,
                     sWeighted2, bWeighted2, b2Dweighted2,
                     bPrint,   KshortDecaysInVelo1, KshortDecaysInVelo2, verbose):
    for year in polarityList():                 
        compareUltimateAllPolarities(variable1 = variable1, variable2 = variable2,
                     year1 = year, year2 = year, cut1 = cut1, cut2 = cut2,
                     magnet1 = magnet, MC1 = MC1, TM1 = TM1, ReferenceChannel1 = ReferenceChannel1, PHSP1 = PHSP1, Preselected1 = Preselected1, BDTed1 = BDTed1,
                     magnet2 = magnet, MC2 = MC2, TM2 = TM2, ReferenceChannel2 = ReferenceChannel2, PHSP2 = PHSP2, Preselected2 = Preselected2, BDTed2 = BDTed2,
                     sWeighted1 = sWeighted1, bWeighted1 = bWeighted1, b2Dweighted1 = b2Dweighted1, weightBranch1 = firstMCweight(),
                     sWeighted2 = sWeighted2, bWeighted2 = bWeighted2, b2Dweighted2 = b2Dweighted2,  weightBranch2 = firstMCweight(),
                     bPrint = bPrint, KshortDecaysInVelo1 = KshortDecaysInVelo1, KshortDecaysInVelo2 = KshortDecaysInVelo2, verbose = verbose)

#Two samples, same variable
def compareUltimateDataMC(variable, year, cut, magnet,
                     TM = True, ReferenceChannel = False, PHSP = False,
                     sWeighted = True, bWeighted  = False, b2Dweighted = True, KshortDecaysInVelo = False, verbose = False):
    return compareUltimate(variable1 = variable, variable2 = variable,
                     year1 = year, year2 = year, cut1 = cut, cut2 = cut,
                     magnet1 = magnet, MC1 = False, TM1 = False, ReferenceChannel1 = False, PHSP1 = False, Preselected1 = True, BDTed1 = False,
                     magnet2 = magnet, MC2 = True, TM2 = TM, ReferenceChannel2 = ReferenceChannel, PHSP2 = PHSP, Preselected2 = True, BDTed2 = False,
                     sWeighted1 = sWeighted, bWeighted1 = False , b2Dweighted1 = False, weightBranch1 = firstMCweight(),
                     sWeighted2 = False, bWeighted2 = bWeighted, b2Dweighted2 = b2Dweighted,  weightBranch2 = firstMCweight(),
                     bPrint = False,   KshortDecaysInVelo1 = KshortDecaysInVelo, KshortDecaysInVelo2 = KshortDecaysInVelo, verbose = verbose)

#Two MC versions, same variable
def compareUltimateMCMC(variable, year1, year2, cut1, cut2, magnet1, magnet2,
                    TM = True, ReferenceChannel = False, PHSP = False,
                    bWeighted1 = False, b2Dweighted1 = False, bWeighted2 = False, b2Dweighted2 = False, KshortDecaysInVelo = False, verbose = False):
    return compareUltimate(variable1 = variable, variable2 = variable,
                     year1 = year1, year2 = year2, cut1 = cut1, cut2 = cut2,
                     magnet1 = magnet1, MC1 = True, TM1 = TM, ReferenceChannel1 = ReferenceChannel, PHSP1 = PHSP, Preselected1 = True, BDTed1 = False,
                     magnet2 = magnet2, MC2 = True, TM2 = TM, ReferenceChannel2 = ReferenceChannel, PHSP2 = PHSP, Preselected2 = True, BDTed2 = False,
                     sWeighted1 = False, bWeighted1 = bWeighted1, b2Dweighted1 = b2Dweighted1, weightBranch1 = firstMCweight(),
                     sWeighted2 = False, bWeighted2 = bWeighted2, b2Dweighted2 = b2Dweighted2, weightBranch2 = firstMCweight(),
                     bPrint = False,   KshortDecaysInVelo1 = KshortDecaysInVelo, KshortDecaysInVelo2 = KshortDecaysInVelo, verbose = verbose)
  
#Two data versions, same variable                   
def compareUltimateDataData(variable, year1, year2, cut1, cut2, magnet1, magnet2, sWeighted1, sWeighted2, KshortDecaysInVelo1, KshortDecaysInVelo2, verbose):

    return compareUltimate(variable1 = variable, variable2 = variable,
                     year1 = year1, year2 = year2, cut1 = cut1, cut2 = cut2,
                     magnet1 = magnet1, MC1 = False, TM1 = False, ReferenceChannel1 = False, PHSP1 = False, Preselected1 = True, BDTed1 = False,
                     magnet2 = magnet2, MC2 = True, TM2 = False, ReferenceChannel2 = False, PHSP2 = False, Preselected2 = True, BDTed2 = False,
                     sWeighted1 = sWeighted1, bWeighted1 = False, b2Dweighted1 = False, weightBranch1 = firstMCweight(),
                     sWeighted2 = sWeighted2, bWeighted2 = False, b2Dweighted2 = False, weightBranch2 = firstMCweight(),
                     bPrint = False,   KshortDecaysInVelo1 = KshortDecaysInVelo, KshortDecaysInVelo2 = KshortDecaysInVelo, verbose = verbose)           
                     
#Two variables, same data sample, different cuts
def compareUltimateTwoVars(variable1, variable2, year, cut1, cut2, magnet, MC, TM, ReferenceChannel, PHSP, Preselected, BDTed, sWeighted, bWeighted, b2Dweighted, KshortDecaysInVelo, verbose = False):
    return compareUltimate(variable1 = variable1, variable2 = variable2,
                     year1 = year, year2 = year, cut1 = cut1, cut2 = cut2,
                     magnet1 = magnet, MC1 = MC, TM1 = TM, ReferenceChannel1 = ReferenceChannel, PHSP1 = PHSP, Preselected1 = Preselected, BDTed1 = BDTed,
                     magnet2 = magnet, MC2 = MC, TM2 = TM, ReferenceChannel2 = ReferenceChannel, PHSP2 = PHSP, Preselected2 = Preselected, BDTed2 = BDTed,
                     sWeighted1 = sWeighted, bWeighted1 = bWeighted, b2Dweighted1 = b2Dweighted, weightBranch1 = firstMCweight(),
                     sWeighted2 = sWeighted, bWeighted2 = bWeighted, b2Dweighted2 = b2Dweighted, weightBranch2 = firstMCweight(),
                     bPrint = False,   KshortDecaysInVelo1 = KshortDecaysInVelo, KshortDecaysInVelo2 = KshortDecaysInVelo, verbose = verbose)
#Two variables, data sample, different cuts
def compareUltimateTwoVarsData(variable1, variable2, year, cut1, cut2, magnet, Preselected, BDTed, sWeighted, KshortDecaysInVelo, verbose):
    return  compareUltimateTwoVars(variable1=variable1, variable2=variable2, year=year, cut1=cut1,cut2=cut2, magnet=magnet, MC=False, TM=False, ReferenceChannel=False, PHSP=False, Preselected=Preselected, BDTed = BDTed, sWeighted=sWeighted, bWeighted=False, b2Dweighted=False, KshortDecaysInVelo=KshortDecaysInVelo, verbose = verbose)       

#Two variables, MC sample, different cuts
def compareUltimateTwoVarsMC(variable1, variable2, year, cut1, cut2, magnet, TM, ReferenceChannel, PHSP, Preselected, BDTed, bWeighted, b2Dweighted, KshortDecaysInVelo, verbose):
     return  compareUltimateTwoVars(variable1=variable1, variable2=variable2, year=year, cut1=cut1,cut2=cut2, magnet=magnet, MC=True, TM=TM, ReferenceChannel=ReferenceChannel, PHSP=PHSP, Preselected=Preselected, BDTed= BDTed, sWeighted=False, bWeighted=bWeighted, b2Dweighted=b2Dweighted, KshortDecaysInVelo=KshortDecaysInVelo, verbose = verbose)   
 
#One variable, same data sample, different cuts
def compareUltimateTwoCuts(variable, year, magnet, cut1, cut2, MC, TM, ReferenceChannel, PHSP, Preselected, BDTed, sWeighted, bWeighted, b2Dweighted, KshortDecaysInVelo, verbose):
    return compareUltimate(variable1 = variable, variable2 = variable,
                     year1 = year, year2 = year, cut1 = cut1, cut2 = cut2,
                     magnet1 = magnet, MC1 = MC, TM1 = TM, ReferenceChannel1 = ReferenceChannel, PHSP1 = PHSP, Preselected1 = Preselected, BDTed1 = BDTed,
                     magnet2 = magnet, MC2 = MC, TM2 = TM, ReferenceChannel2 = ReferenceChannel, PHSP2 = PHSP, Preselected2 = Preselected, BDTed2 = BDTed,
                     sWeighted1 = sWeighted, bWeighted1 = bWeighted, b2Dweighted1 = b2Dweighted, weightBranch1 = firstMCweight(),
                     sWeighted2 = sWeighted, bWeighted2 = bWeighted, b2Dweighted2 = b2Dweighted, weightBranch2 = firstMCweight(),
                     bPrint = False,   KshortDecaysInVelo1 = KshortDecaysInVelo, KshortDecaysInVelo2 = KshortDecaysInVelo, verbose = verbose)                  

#One variable, data sample, different cuts
def compareUltimateTwoCutsData(variable, year, magnet, cut1, cut2, Preselected, BDTed, sWeighted, KshortDecaysInVelo, verbose):
    return compareUltimateTwoCuts(variable=variable, year= year, magnet= magnet, cut1 = cut1, cut2 = cut2,MC=False, TM=False, ReferenceChannel = False, PHSP = False, Preselected= Preselected, BDTed = BDTed, sWeighted = sWeighted, bWeighted = False, b2Dweighted = False, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)

#One variable, MC sample, different cuts
def compareUltimateTwoCutsMC(variable, year, magnet, cut1, cut2, TM, ReferenceChannel, PHSP, Preselected, BDTed, bWeighted, b2Dweighted, KshortDecaysInVelo, verbose):
    return compareUltimateTwoCuts(variable=variable, year= year, magnet= magnet, cut1 = cut1, cut2 = cut2,MC=True, TM=TM, ReferenceChannel = ReferenceChannel, PHSP = False, Preselected= Preselected, BDTed = BDTed, sWeighted = False, bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)


def compareUltimateBchargeAllSamples(variable, year, magnet, MC, ReferenceChannel, PHSP, Preselected, BDTed, sWeighted, bWeighted, b2Dweighted, KshortDecaysInVelo, verbose):
    if (not MC):
        compareUltimateTwoCutsData(variable=variable, year= year, magnet= magnet, cut1 = "B_plus_ID = 521", cut2 = "B_plus_ID = -521", Preselected= Preselected, BDTed = BDTed, sWeighted = sWeighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)
    else:
        compareUltimateTwoCutsMC(variable=variable, year= year, magnet= magnet, cut1 = "B_plus_ID = 521", cut2 = "B_plus_ID = -521", TM=False, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed, sWeighted = False, bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)
        compareUltimateTwoCutsMC(variable=variable, year= year, magnet= magnet, cut1 = "B_plus_ID = 521", cut2 = "B_plus_ID = 521", TM=False, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed, sWeighted = False, bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)
        compareUltimateTwoCutsMC(variable=variable, year= year, magnet= magnet, cut1 = "B_plus_ID = -521", cut2 = "B_plus_ID = -521", TM=False, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed, sWeighted = False, bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)

        compareUltimateTwoCutsMC(variable=variable, year= year, magnet= magnet, cut1 = "B_plus_TRUEID = 521", cut2 = "B_plus_TRUEID = -521", TM=True, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed, sWeighted = False, bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)

        compareUltimateTwoCutsMC(variable=variable, year= year, magnet= magnet, cut1 = "B_plus_TRUEID = 521", cut2 = "B_plus_TRUEID = -521", TM=True, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed, sWeighted = False, bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)


def compareUltimateBchargeAllSamplesTwoVars(variable1, variable2, year, magnet, MC, ReferenceChannel, PHSP, Preselected, BDTed, sWeighted, bWeighted, b2Dweighted, KshortDecaysInVelo, verbose):
    if (not MC):
        compareUltimateTwoVarsData(variable1=variable1, variable2=variable2, year= year, magnet= magnet, cut1 = "B_plus_ID = 521", cut2 = "B_plus_ID = -521", Preselected= Preselected, BDTed = BDTed, sWeighted = sWeighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)
    else:
        compareUltimateTwoVarsMC(variable1=variable1, variable2=variable2,year= year, magnet= magnet, cut1 = "B_plus_ID = 521", cut2 = "B_plus_ID = -521", TM=False, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed, bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)
        compareUltimateTwoVarsMC(variable1=variable1, variable2=variable2,year= year, magnet= magnet, cut1 = "B_plus_ID = -521", cut2 = "B_plus_ID = 521", TM=False, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed, bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)
        compareUltimateTwoVarsMC(variable1=variable1, variable2=variable2,year= year, magnet= magnet, cut1 = "B_plus_ID = 521", cut2 = "B_plus_ID = 521", TM=False, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed,  bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose) 
        compareUltimateTwoVarsMC(variable1=variable1, variable2=variable2,year= year, magnet= magnet, cut1 = "B_plus_ID = -521", cut2 = "B_plus_ID = -521", TM=False, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed,  bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)

        compareUltimateTwoVarsMC(variable1=variable1, variable2=variable2,year= year, magnet= magnet, cut1 = "B_plus_TRUEID = 521", cut2 = "B_plus_TRUEID = -521", TM=True, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed,  bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)

        compareUltimateTwoVarsMC(variable1=variable1, variable2=variable2,year= year, magnet= magnet, cut1 = "B_plus_TRUEID = 521", cut2 = "B_plus_TRUEID = -521", TM=True, ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected= Preselected, BDTed = BDTed, bWeighted = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = KshortDecaysInVelo, verbose = verbose)



def compareUltimateBchargeAllPolarities(variable, year,MC, ReferenceChannel, PHSP, Preselected, BDTed, sWeighted, bWeighted, b2Dweighted, KshortDecaysInVelo, verbose):
    for pol in polarityList():
        compareUltimateBchargeAllSamples(variable=variable, year=year, magnet = pol,MC = MC,  ReferenceChannel= ReferenceChannel, PHSP=PHSP, Preselected = Preselected, BDTed= BDTed, sWeighted=sWeighted, bWeighted=bWeighted, b2Dweighted=b2Dweighted, KshortDecaysInVelo=KshortDecaysInVelo, verbose=verbose)
    

def compareUltimateBchargeAllYears(variable,MC, ReferenceChannel, PHSP, Preselected, BDTed, sWeighted, bWeighted, b2Dweighted, KshortDecaysInVelo, verbose):
    for year in yearList():
        compareUltimateBchargeAllPolarities(variable = variable, year = year, MC = MC,ReferenceChannel = ReferenceChannel, PHSP = PHSP, Preselected=  Preselected, BDTed = BDTed, sWeighted = sWeighted, bWeighted= bWeighted, b2Dweighted = b2Dweighted,KshortDecaysInVelo =KshortDecaysInVelo, verbose= verbose)

def compareUltimateMuonIPCHI2All(verbose):  #TODO possibly change parameters such as preselected
    compareUltimateBchargeAllYears(variable = "log(mu_plus_IPCHI2_OWNPV)", MC = False, ReferenceChannel = False, PHSP = False, Preselected=  True, BDTed = False, sWeighted = False,    bWeighted= False, b2Dweighted = False,KshortDecaysInVelo =False, verbose= verbose)
    compareUltimateBchargeAllYears(variable = "log(mu_minus_IPCHI2_OWNPV)", MC = True, ReferenceChannel = False, PHSP = False, Preselected=  True, BDTed = False, sWeighted = False,    bWeighted= False, b2Dweighted = False,KshortDecaysInVelo =False, verbose= verbose)
    compareUltimateBchargeAllYears(variable = "log(mu_minus_IPCHI2_OWNPV)", MC = True, ReferenceChannel = True, PHSP = False, Preselected=  True, BDTed = False, sWeighted = False,    bWeighted= False, b2Dweighted = False,KshortDecaysInVelo =False, verbose= verbose)
    compareUltimateBchargeAllYears(variable = "log(mu_minus_IPCHI2_OWNPV)", MC = True, ReferenceChannel = False, PHSP = True, Preselected=  True, BDTed = False, sWeighted = False,    bWeighted= False, b2Dweighted = False,KshortDecaysInVelo =False, verbose= verbose)




def compareUltimateBDTcut(variable, year, magnet, sWeighted, KshortDecaysInVelo, verbose):
    return compareUltimateTwoCuts(variable = variable, year =year, magnet =magnet,  cut1 ="MLPresponse < 0.1", cut2 = "MLPresponse > 0.9", MC=False, TM=False, ReferenceChannel=False, PHSP=False, Preselected=True, BDTed=True, sWeighted = sWeighted, bWeighted=False, b2Dweighted=False, KshortDecaysInVelo=KshortDecaysInVelo, verbose=verbose)       


def compareUltimateAllKplusDataMC(year = 2011, magnet = "both", cut = "", ReferenceChannel = False, PHSP = False, bWeighted  = False, b2Dweighted = True, whichParticle = 1, verbose = False):
    for var in ListOfKplusVars(whichParticle=whichParticle):
        compareUltimateDataMC(var, year, cut, magnet,
                     TM = True, ReferenceChannel = ReferenceChannel, PHSP = PHSP,
                     sWeighted = True, bWeighted  = bWeighted, b2Dweighted = b2Dweighted, KshortDecaysInVelo = False, verbose = verbose)
    
    return 0 #TODO

def compareUltimateAllKplusBDTcut(year = 2011, magnet = "both",  whichParticle = 1, verbose = False):
    for var in ListOfKplusVars(whichParticle=whichParticle):
        compareUltimateBDTcut(var, year, year, magnet, magnet,  KshortDecaysInVelo =False)    
    return 0 #TODO
