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

# <img src="http://files.oproject.org/img/HeaderOpenData.png">
# # CMS Open Data Example #6: Fitting the J/$\psi$ fit
# 
# ## Import required modules

# In[1]:


from ROOT import TFile, TCanvas, TH1F, TLorentzVector
get_ipython().run_line_magic('jsroot', 'on')


# ## Read in Data from Input File

# In[2]:


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


# ## Declare Histograms

# In[25]:


Jpsi = TH1F("Jpsi","#mu#mu mass;#mu#mu mass [GeV];Events", 1000, 2, 12)


# # Compute Invariant Mass of J/$\psi$

# In[26]:


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)
    
        InvariantMass = (muon1 + muon2).M() 
        
        if InvariantMass > 2.0 and InvariantMass < 12.0:

            Jpsi.Fill(InvariantMass)


# ## Fitting the J/$\psi$ Peak

# In[40]:


#Importing colors. All colors can be found at: https://root.cern.ch/doc/master/pict1_TColor_002.png
from ROOT import TF1, kGreen, kRed, kBlack, kBlue, kMagenta, kSpring, kOrange

Canvas = TCanvas()


# Possible fitting options: <BR>
# "gaus" gaussian <BR>
# "pol1" 1st order polynomial <BR>

# In[48]:


Gaussian = TF1("Gaussian","gaus", 3.03, 3.16)
Gaussian2 = TF1("Gaussian2","gaus", 3.61, 3.73)
Gaussian3 = TF1("Gaussian3","gaus", 9.30, 9.60)
Gaussian4 = TF1("Gaussian4","gaus", 9.90, 10.10)
Background  = TF1("Background","pol1", 2, 3.5)
Background2 = TF1("Background2","pol1", 3.5, 4.5)
Background3 = TF1("Background3","pol1", 9.0,9.80)
Background4 = TF1("Background4","pol1", 9.8,10.2)

Gaussian.SetLineColor(kMagenta)
Gaussian2.SetLineColor(kRed)
Gaussian3.SetLineColor(kGreen)
Gaussian4.SetLineColor(kOrange)
Background.SetLineColor(kMagenta-3)
Background2.SetLineColor(kRed-3)
Background3.SetLineColor(kGreen-3)
Background4.SetLineColor(kOrange-3)

#Adding all of the ranges to the graph. "R" adds one range, "R+" adds any more desired ranges on top of the graph.

Jpsi.Fit(Gaussian, "R")
Jpsi.Fit(Gaussian2,"R+")
Jpsi.Fit(Gaussian3,"R+")
Jpsi.Fit(Gaussian4,"R+")
Jpsi.Fit(Background,"R+")
Jpsi.Fit(Background2, "R+")
Jpsi.Fit(Background3, "R+")
Jpsi.Fit(Background4, "R+")

Canvas.Draw()


# In[ ]: