#!/usr/bin/env python
# coding: utf-8

# <img src="http://files.oproject.org/img/HeaderOpenData.png">
# 
# # CMS Open Data Example #3: Di-Muon Resonances
# <img src="http://cms.web.cern.ch/sites/cms.web.cern.ch/files/styles/large/public/field/image/2011-bs-1-2.jpg?itok=k5_hcnFt"></img> <BR>
# ## Import Modules and Turn on Javascript

# In[3]:


from ROOT import TTree, TFile, TCanvas, TH1F, TLorentzVector

get_ipython().run_line_magic('jsroot', 'on')


# ## Read in Data from Input File

# In[4]:


file = TFile("data/Dimuons.root","READ");


# # Compute Di-Muon Invariant Mass
# Let's calculate again the invariant mass $M$ of two muons and focus on various parts of the dimuon mass spectrum

# ## Exercise: Can You Spot Other Di-Muon Resonances in the Dimuon Spectrum?
# ### Plot a Gaussian Curve for Each Resonance 

# In[10]:


Upsilon = TH1F("ϒ", "#mu#mu mass;#mu#mu mass [GeV];Events", 100, 8.0, 12.5)


# In[11]:


for dimu in file.Dimuons:
    if dimu.Muon1_Global and dimu.Muon2_Global:
        
        muon1 = TLorentzVector(dimu.Muon1_Px, dimu.Muon1_Py, dimu.Muon1_Pz, dimu.Muon1_Energy)
    
        muon2 = TLorentzVector(dimu.Muon2_Px, dimu.Muon2_Py, dimu.Muon2_Pz, dimu.Muon2_Energy)
        
        Invariant_Mass = (muon1+muon2).M()
        
        if Invariant_Mass > 0.0 and Invariant_Mass < 120.0:
        
            Upsilon.Fill(Invariant_Mass)


# In[12]:


from ROOT import TF1, kGreen, kRed, kYellow, kOrange

Canvas = TCanvas()


# In[13]:


Gaussian = TF1("Gaussian","gaus", 9.23, 9.6)
Background  = TF1("Background","pol2", 8, 12.5)
Gaussian1 = TF1("Gaussian","gaus", 9.7, 10.2)
Gaussian2 = TF1("Gaussian","gaus",10.25, 10.45)

Gaussian.SetLineColor(kRed)
Background.SetLineColor(kGreen)
Gaussian1.SetLineColor(kYellow)
Gaussian2.SetLineColor(kOrange)

Upsilon.Fit(Gaussian,"R")
Upsilon.Fit(Background,"R+")
Upsilon.Fit(Gaussian1,"R++")
Upsilon.Fit(Gaussian2,"R+++")
Canvas.Draw()


# In[ ]: