# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Formal Update 7/20/2021
# ### <NAME>

import awkward1 as ak
import uproot
from coffea.nanoevents import NanoAODSchema,NanoEventsFactory
import coffea
from uproot3_methods import TLorentzVectorArray
import numpy as np
class HackSchema(NanoAODSchema):
    def __init__(self, base_form):
        base_form["contents"].pop("Muon_fsrPhotonIdx", None)
        base_form["contents"].pop("Electron_photonIdx", None)
        super().__init__(base_form)
files ="TTbarPowheg_Semilept_Skim_NanoAOD_1of21.root"
import coffea.processor as processor
from pprint import pprint
file=uproot.open(files)
from pprint import pprint

# ### Table of Contents:
#
# 1. Data Preprocessing and Goals for data usage
# 2. Relevant data per combo and event
# 3. Model Structure and Parameters

# ______________________________________________________________________

# The data used is event data in the NanoAOD format:

events=NanoEventsFactory.from_root(files,schemaclass=HackSchema,entry_stop=10000).events()

events.fields

#Condensing_all Lepton_cuts
jets=events.Jet
jetSel = (jets.pt>30) & (abs(jets.eta)<2.4)
tightJet = jets[jetSel]
#Selecting Jets
bJet = tightJet[tightJet.btagDeepFlavB > 0.642]
muons = events.Muon
muonSel = (muons.pt>30) & (abs(muons.eta)<2.4)
tightMuon = muons[muonSel]
#Selecting Events that have decaying leptons
ele = events.Electron
eleSel = (ele.pt>35)&(abs(ele.eta)<2.4)
tightEle = ele[eleSel]
eventSel = (((ak.num(tightMuon)==1) | (ak.num(tightEle)==1)) &
            (ak.num(tightJet)>= 3) & (ak.num(bJet)>=1)
           )
final=events[eventSel]

# We will be using the GenPart and GenJet columns to match the Jets

# We are looking for occurences of hadronically decaying tops:
# t->bW->bqq.
#
# First we will identify top quarks using pdgId.
#
# (Link to pdgid number scheme: https://pdg.lbl.gov/2007/reviews/montecarlorpp.pdf)

genPart = final.GenPart
tops = genPart[abs(genPart.pdgId)==6]
print(tops.pdgId[0:3])
#The isLastCopy Flag filters out copy Genparticles:
tops = tops[tops.hasFlags('isLastCopy')]
print(tops.pdgId[0:3])

# GenPart has the function distinctChildren and distinctParent respectively, using distinctChildren on tops returns top decays products:

tDecay = tops.distinctChildren
tDecay = tDecay[tDecay.hasFlags('isLastCopy')]
print(tDecay.pdgId[0:3])

# Isolating the W boson(+-24) and the bottom(+-5):

t_Events=tDecay[abs(tDecay.pdgId)==5]
W = tDecay[abs(tDecay.pdgId)==24]
W = W[W.hasFlags('isLastCopy')]
WDecay = W.distinctChildren
WDecay = WDecay[WDecay.hasFlags('isLastCopy')]
#t_events is the lone bottom, W_events is the -> two jets
print("W Decay")
print(WDecay[0:3].pdgId)
print("Filtering out non-hadronically deacaying product(ie light quarks only)")
#select the hadronically decaying W
W_Events=ak.flatten(WDecay[ak.all(abs(WDecay.pdgId)<=8,axis=-1)],axis=3)
print(W_Events[0:3].pdgId)
#print(qqb)

# Create a mask for the quark decay events using W_Events:

#HadW is mask for Quark deacying W boson
hadW = ak.num(W_Events,axis=2)==2
#filters out t_events that have a hadronically decayign W Boson
hadB = t_Events[hadW]
hadB = ak.flatten(hadB,axis=2)
print(len(hadB))
print("this will only contian bottom quarks, from events that have hadronic W bosons")
print(hadB.pdgId[0:3])
W_quarks = W_Events[hadW]
W_quarks = ak.flatten(W_quarks,axis=2)
print(len(hadW))
#concatentating these two arrays make an array of events with the correctly decaying GenParticles.
qqb = ak.concatenate([hadB,W_quarks],axis=1)
print(len(qqb))
print("The final correct combinations")
print(qqb[0:3].pdgId)
#lets filter for only 1 HadronicW
#This last mask just gets rid of events that conatin more than one hadronic top

#Filtering Out events with extra tops
final=final[(ak.count(qqb.pdgId,axis=1)==3)]
finaljets=final.Jet
qqb=qqb[(ak.count(qqb.pdgId,axis=1)==3)]

#Implementing Tight Jet Cuts on Training Data
finaljetSel=(abs(finaljets.eta)<2.4)&(finaljets.pt>30)
finalJets=finaljets[finaljetSel]

print(len(ak.flatten(finalJets[(abs(finalJets.eta)<2.4)])))
print(len(ak.flatten(finalJets)))

matchedGenJets=qqb.nearest(final.GenJet)
print(matchedGenJets.eta[0:3])
print(matchedGenJets.partonFlavour[0:3])
print(qqb.pdgId[0:3])

#matchedJets=qqb.nearest(finalJets)
matchedJets=matchedGenJets.nearest(finalJets)
print(len(ak.flatten(matchedJets)))
print(len(ak.flatten(matchedJets[(abs(matchedJets.eta)<2.4)])))

print(matchedJets.partonFlavour[0:3])

# Passing the reconstructed GenPart through the eventSel will fix mismatch array lengths

test=matchedJets.genJetIdx
combs=ak.combinations(finalJets,3,replacement=False)
t1=(combs['0'].genJetIdx==test[:,0])|(combs['0'].genJetIdx==test[:,1])|(combs['0'].genJetIdx==test[:,2])
t2=(combs['1'].genJetIdx==test[:,0])|(combs['1'].genJetIdx==test[:,1])|(combs['1'].genJetIdx==test[:,2])
t3=(combs['2'].genJetIdx==test[:,0])|(combs['2'].genJetIdx==test[:,1])|(combs['2'].genJetIdx==test[:,2])
t=t1&t2&t3
trutharray=ak.flatten(t)
print(len(trutharray))

print(t1[1])
print(t2[1])
print(t3[1])
print(t[1])

print(t1[0,5])
print(t2[0,5])
print(t3[0,5])
print(t[0,5])

print(t1[1])
print(t2[1])
print(t3[1])
print(t[1])

# We can plot the resulting combos mass to compare the good combos, bad combos, and the Genpart reco mass.

jetcombos=ak.flatten(combs)
genpart_vals=ak.combinations(qqb,3)
good_vals=jetcombos[trutharray]
print(len(good_vals))
bad_vals=jetcombos[np.invert(trutharray)]
print(len(bad_vals))

import matplotlib.pyplot as plt
bins1=[]
n=0
while n*5<=500:
    bins1.append(n*10)
    n+=1
plt.figure(figsize=(20,10),dpi=80)
greco=good_vals['0']+good_vals['1']+good_vals['2']
breco=bad_vals['0']+bad_vals['1']+bad_vals['2']
genreco=genpart_vals['0']+genpart_vals['1']+genpart_vals['2']
plt.hist(greco.mass,alpha=.5,density=True,bins=bins1)
plt.hist(breco.mass,alpha=.5,density=True,bins=bins1)
plt.hist(ak.flatten(genreco.mass),alpha=.5,density=True,bins=bins1)
plt.legend(["Matched Jets Combos","Bad Event Combos","GenPart Combos"])
plt.title("Mass of M3 Reconstruction")
plt.show()
plt.savefig("Mass.png")

print(np.mean(greco.mass))
print(np.mean(breco.mass))
print(np.mean(ak.flatten(genreco.mass)))

# # Relevant Features:
#
# ## Per Combinations of Jets:
#
# ### Each Jet:
# * $P_t$
# * $ \eta $
# * $ \phi $
# * mass
# * b_tag
#
# ### $ \Delta R $ Between Jets
# * 1 & 3
# * 2 & 3
# * 1 & 2
#
# ### Potentially Included
# * Event Number
#
#

# +
### Plot Distribution of Features to be utilized by model
# -

gdelta12=good_vals['0'].delta_r(good_vals['1'])
gdelta23=good_vals['1'].delta_r(good_vals['2'])
gdelta13=good_vals['0'].delta_r(good_vals['2'])
gdelta=ak.concatenate([gdelta12,gdelta23,gdelta13])
bdelta12=bad_vals['0'].delta_r(bad_vals['1'])
bdelta23=bad_vals['1'].delta_r(bad_vals['2'])
bdelta13=bad_vals['0'].delta_r(bad_vals['2'])
bdelta=ak.concatenate([bdelta12,bdelta23,bdelta13])
plt.figure(figsize=(20,10),dpi=80)
bins1=[]
n=0
while n<80:
    bins1.append(n*.1)
    n=n+1
l=plt.hist(gdelta,alpha=.5,density=True,bins=bins1)
l=plt.hist(bdelta,alpha=.5,density=True,bins=bins1)
plt.legend(["Matched Combination $ \Delta R $","Unmatched Combination $ \Delta R $"])
plt.title("$\Delta R$ Between 3 Jets")
plt.show()
plt.savefig("DeltaR.png")

# $ \Delta R $ of the unmatched data favors a Delta R between jets. Noticeably:

bnpd12=ak.to_numpy(bdelta12)
bnpd13=ak.to_numpy(bdelta13)
bnpd23=ak.to_numpy(bdelta23)
bzipdr=ak.sort(np.stack((bnpd12,bnpd13,bnpd23),axis=1))
npd12=ak.to_numpy(gdelta12)
npd13=ak.to_numpy(gdelta13)
npd23=ak.to_numpy(gdelta23)
zipdr=ak.sort(np.stack((npd12,npd13,npd23),axis=1))
plt.figure(figsize=(20,10),dpi=80)
l=plt.hist(zipdr[:,0],density=True,bins=bins1,histtype='step')
l=plt.hist(zipdr[:,1],density=True,bins=bins1,histtype='step')
l=plt.hist(zipdr[:,2],density=True,bins=bins1,histtype='step')
plt.title("$\Delta R$ of each Jet Pair")
#l=plt.hist(bzipdr[:,0],density=True,bins=bins1,histtype='step')
plt.legend(['Leading Jet to 2nd','Leading Jet to 3nd','Delta R of 2nd and 3rd Jet'])

plt.figure(figsize=(20,10),dpi=80)
plt.title("$\Delta R$ of each Jet Pair")
l=plt.hist(bzipdr[:,0],density=True,bins=bins1,histtype='step')
l=plt.hist(bzipdr[:,1],density=True,bins=bins1,histtype='step')
l=plt.hist(bzipdr[:,2],density=True,bins=bins1,histtype='step')
plt.legend(['Leading Jet to 2nd','Leading Jet to 3nd','Delta R of 2nd and 3rd Jet'])

geta23=good_vals['1'].eta-good_vals['2'].eta
geta13=good_vals['0'].eta-good_vals['2'].eta
geta12=good_vals['0'].eta-good_vals['1'].eta
geta=ak.concatenate([geta23,geta12,geta13])
beta12=bad_vals['0'].eta-bad_vals['1'].eta
beta23=bad_vals['1'].eta-bad_vals['2'].eta
beta13=bad_vals['0'].eta-bad_vals['2'].eta
beta=ak.concatenate([beta23,beta12,beta13])
plt.figure(figsize=(20,10),dpi=80)
bins1=[]
n=-90
while n<90:
    bins1.append(n*.1)
    n=n+1
l=plt.hist(geta,alpha=.5,density=True,bins=bins1)
l=plt.hist(beta,alpha=.5,density=True,bins=bins1)
plt.legend(["Matched Combination $ \Delta \eta $","Unmatched Combination $ \Delta \eta $"])
plt.title("$ \Delta \eta $ Between 3 Jets")
plt.show()
plt.savefig("deltaeta.png")

# $ \Delta R $ of the unmatched data favors a Delta R between jets. Noticeably:

print("%.2f percent of the matched combos Delta R falls below 4" %(len(gdelta[gdelta<4])/len(gdelta)*100))

gbtagSum=good_vals['0'].btagDeepB+good_vals['2'].btagDeepB+good_vals['1'].btagDeepB
bbtagSum=bad_vals['0'].btagDeepB+bad_vals['2'].btagDeepB+bad_vals['1'].btagDeepB

bins1=[]
n=0
while n<60:
    bins1.append(n*.05)
    n=n+1
plt.figure(figsize=(20,10),dpi=80)
l=plt.hist(gbtagSum,alpha=.5,density=True,bins=bins1)
l=plt.hist(bbtagSum,alpha=.5,density=True,bins=bins1)
plt.title("Sum of B-tag events")
plt.legend(["Matched Combination b-tag probability Sum","Unmatched Combination b-tag probability Sum"])

geta=ak.concatenate([good_vals['0'].eta,good_vals['1'].eta,good_vals['2'].eta])
beta=ak.concatenate([bad_vals['0'].eta,bad_vals['1'].eta,bad_vals['2'].eta])
plt.figure(figsize=(20,10),dpi=80)
bins1=[]
n=-70
while n<70:
    bins1.append(n*.1)
    n=n+1
l=plt.hist(geta,alpha=.5,density=True,bins=bins1)
l=plt.hist(beta,alpha=.5,density=True,bins=bins1)
plt.title("Eta of Jets")
plt.legend(["Matched Jets Eta","UnMatched Jet"])

geta=ak.concatenate([good_vals['0'].phi,good_vals['1'].phi,good_vals['2'].phi])
beta=ak.concatenate([bad_vals['0'].phi,bad_vals['1'].phi,bad_vals['2'].phi])
plt.figure(figsize=(20,10),dpi=80)
bins1=[]
n=-40
while n<40:
    bins1.append(n*.1)
    n=n+1
l=plt.hist(geta,alpha=.5,density=True,bins=bins1)
l=plt.hist(beta,alpha=.5,density=True,bins=bins1)
plt.title("Phi of Jets")
plt.legend(["Matched Jets","UnMatched Jet"])

# +
#Zipping Event into .csv file
# -

j1,j2,j3=ak.unzip(jetcombos)
dR1_2=j1.delta_r(j2)
dR1_3=j1.delta_r(j3)
dR2_3=j2.delta_r(j3)
j1b_tag=j1.btagCSVV2
j2b_tag=j2.btagCSVV2
j3b_tag=j3.btagCSVV2
j1area=j1.area
j2area=j2.area
j3area=j3.area
j1p4=ak.zip({"pt":j1.pt,"phi":j1.phi,"eta":j1.eta,"mass":j1.mass},with_name="PtEtaPhiMLorentzVector")
j2p4=ak.zip({"pt":j2.pt,"phi":j2.phi,"eta":j2.eta,"mass":j2.mass},with_name="PtEtaPhiMLorentzVector")
j3p4=ak.zip({"pt":j3.pt,"phi":j3.phi,"eta":j3.eta,"mass":j2.mass},with_name="PtEtaPhiMLorentzVector")
processedMLdata=ak.zip({"j1p4":j1p4,"j2p4":j2p4,"j3p4":j3p4,
                        "dR12":dR1_2,"dR13":dR1_3,"dR23":dR2_3,
                       "j1btag":j1b_tag,"j2btag":j2b_tag,"j3btag":j3b_tag})

len(jetcombos['0'])

len(jetcombos['0'][jetcombos['1'].pt<=jetcombos['2'].pt])

df=ak.to_pandas(processedMLdata)
df.to_csv("data/ML_data.csv",index=False)
vf=ak.to_pandas(trutharray)
vf.to_csv("data/Valid_data.csv",index=False)

# ## Creating simple Boosted Decision Tree:

# ![ROC_Curve_Validation_Boosted_Decision_Tree-2.png](attachment:ROC_Curve_Validation_Boosted_Decision_Tree-2.png)

print("ROC Curve is %.4f" %0.787479747138643)

# This is a rough tree with only 2 Parallel trees, 10 boosting rounds using a Binary:logistic Objective and alpha of .8. This is a slow learning XgBoost tree. I am exploring more intricate tree structures
#
# ### Questions for Boosted Decision Tree
#
# However, Should the tree be a binary:hinge(0,1) or binary:logistic(continuous probability from 0 to 1)?
#
# ### General Questions
#    
# Should we include or remove certain features like the eta and phi; and instead use $\Delta \phi$?

# ## To Do:
# * PreProcessing Coffea Processor
# * Use Boosted Decision Tree
# * Compare Efficiency to simple data cuts
# * ROC + Confusion Matrix

# # Sources/Further Reading:
#
# CMS Draft Analysis Note(FIT, AAS, Wayne State, NISTR): https://drive.google.com/file/d/1XEOLyZ-Q1HdEQY379RpyyQkOF1Q8KlsL/view
#
# Invariant Mass Wikipedia Article: https://en.wikipedia.org/wiki/Invariant_mass
#
# Measurement of the semileptonic tt + γ production cross-section in pp collisions at √s = 8 Tev: https://arxiv.org/pdf/1706.08128.pdf
#
# Coffea Documentation: https://coffeateam.github.io/coffea/
#
# Awkward Array Documentation: https://awkward-array.readthedocs.io/en/latest/
#
# AOD DATA GUIDE: https://twiki.cern.ch/twiki/bin/view/CMSPublic/SWGuideAodDataTable

StarcoderdataJupyterScriptsDedupFiltered