#Relevant only for Zgamma CR studies where we have 5 njet_bjet categories from Top-21-004
#For completeness, there are other low jet_bjet categories as well
#The function returns integer from 1 to 9 for each event. If for a given event, the integer is 8, it means that event has exactly 2 jets and 2 bjets, etc.

import awkward as ak

def multiplicityOfJetsAndbJets(num_jets, num_bjets):
    #First is 0 btagged jets
    jets0_bjets0 = ((num_jets==0) & (num_bjets==0))
    jets1_bjets0 = ((num_jets==1) & (num_bjets==0))
    jets2_bjets0 = ((num_jets==2) & (num_bjets==0))
    jets3_bjets0 = ((num_jets==3) & (num_bjets==0))
    jetsgeq4_bjets0 = ((num_jets>=4) & (num_bjets==0))


    #second is 1 btagged jet
    jets1_bjets1 = ((num_jets==1) & (num_bjets==1))
    jets2_bjets1 = ((num_jets==2) & (num_bjets==1))
    jets3_bjets1 = ((num_jets==3) & (num_bjets==1))
    jetsgeq4_bjets1 = ((num_jets>=4) & (num_bjets==1))

    #third is 2 btagged jet
    jets2_bjets2 = ((num_jets==2) & (num_bjets==2))
    jets3_bjets2 = ((num_jets==3) & (num_bjets==2))
    jetsgeq4_bjets2 = ((num_jets>=4) & (num_bjets==2))

    #forth is 3 btagged jet
    jets3_bjets3 = ((num_jets==3) & (num_bjets==3))
    jetsgeq4_bjets3 = ((num_jets>=4) & (num_bjets==3))
    
    #atleast 4 btagged jet
    jetsgeq4_bjetsgeq4 = ((num_jets>=4) & (num_bjets>=4)) 

    return 1*jets0_bjets0 + 2*jets1_bjets0 + 3*jets2_bjets0 + 4*jets3_bjets0 + 5*jetsgeq4_bjets0 + 6*jets1_bjets1 + 7*jets2_bjets1 + 8*jets3_bjets1 + 9*jetsgeq4_bjets1 + 10*jets2_bjets2 + 11*jets3_bjets2 + 12*jetsgeq4_bjets2 + 13*jets3_bjets3 + 14*jetsgeq4_bjets3 + 15*jetsgeq4_bjetsgeq4