from ROOT import *
import sys

if len(sys.argv) != 3:
	print("Please include the name of the root file you want to run and the name of the tree")

variables = ['mJ', 'deltaRbTrkJbTrkJ', 'pTBTrkJ1', 'pTBTrkJ2', 'pTBTrkJ3', 'pTV', 
	     'NAdditionalCaloJets', 'NMatchedTrackJetLeadFatJet', 'deltaRFJTrkJ1', 
             'deltaRFJTrkJ2', 'deltaRbTrkJ1TrkJ3', 'softMET', 'deltaEtabTrkJbTrkJ', 'MET',
             'etaBTrkJ1', 'sumPtJets', 'nSigJet'] #Here's where you'll list th evariables you want to plot
xmin = [0, 0, 0, 0, 0, 400, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, -3, 0, 0] #The minimum value to plot on the histogram
xmax = [200, 2, 300, 300, 300, 1000, 3.5, 5, 5, 5, 10, 10, 1.5, 1.5, 1.5, 50, 2, 600, 3, 2200,14] #The maximum number to plot on the histogram
nbins = [20, 20, 20, 20, 20, 20, 15, 5, 5, 20, 10, 10, 10, 10, 10, 20, 10, 20, 20, 20, 14] #The number of bins in the histogram

file = TFile(sys.argv[1], "READ")
Tree = file.Get(sys.argv[2])

for i in range(1):#len(variables)):
    #Create the histograms to plot
    h_Data1 = TH1F("Data1", 'Data histogram', nbins[i], xmin[i], xmax[i])
    h_MC1 = TH1F("MC1", "MC histogram", nbins[i], xmin[i], xmax[i])
    h_Data2 = TH1F("Data2", variables[i] + " Ratio", nbins[i], xmin[i], xmax[i])
    h_MC2 = TH1F("MC2", "MC histogram", nbins[i], xmin[i], xmax[i])

    #Fill the histograms
    for entry in Tree:
        if entry.sample == 'data': #Fill the data histograms
	    #These are the cuts you'll apply, they'll vary based on what interaction you're looking at
            if entry.pTV>400:
                if entry.NBTagMatchedTrackJetLeadFatJet==2:
                    if entry.NBTagUnmatchedTrackJetLeadFatJet>0.5:
                        h_Data1.Fill(getattr(entry, variables[i]))
                        h_Data2.Fill(getattr(entry, variables[i]))
        elif entry.sample != 'data': #Fill the MC histogram
            if entry.pTV>400:
                if entry.NBTagMatchedTrackJetLeadFatJet==2:
                    if entry.NBTagUnmatchedTrackJetLeadFatJet>0.5:
			#Apply cuts for additional ttbar weight (in this case, ttbar is the dominant background)
                        if entry.sample == "ttbar2VR" or entry.sample == 'ttbar3pVR':
                            h_MC1.Fill(getattr(entry, variables[i]), entry.EventWeight * 0.83)
                            h_MC2.Fill(getattr(entry, variables[i]), entry.EventWeight * 0.83)
                        else:
                            h_MC1.Fill(getattr(entry, variables[i]), entry.EventWeight)
                            h_MC2.Fill(getattr(entry, variables[i]), entry.EventWeight)
    
    #Calculate the error
    h_Data1.Sumw2()
    h_MC1.Sumw2()

    #Stack one set of histograms so they can both be on the same graph
    hs = THStack("hs", variables[i])
    h_Data1.SetLineColor(kMagenta)
    hs.Add(h_Data1, 'E')
    h_MC1.SetLineColor(kBlack)
    hs.Add(h_MC1, 'E')

    #Divide the second data histogram by the second MC histogram --> This will be used to plot the ratio of the two histograms
    h_Data2.Divide(h_MC2)

    #Plot and save the MC comparison plot
    c1 = TCanvas('c1', 'hist', 1200, 800)
    hs.Draw('nostack')
    legend = TLegend(0.9, 0.9, 0.7, 0.7)# 0.2, 0.9, 0.4, 0.7)
    legend.AddEntry(h_Data1, "Data")
    legend.AddEntry(h_MC1, "MC")
    legend.Draw('right')
    c1.Draw('E')
    c1.SaveAs(variables[i] + 'test.jpg')

    #Plot and save the ratio plot
    c2 = TCanvas('c2', 'hist', 1200, 800)
    h_Data2.Draw('HIST')
    c2.SaveAs(variables[i] + '_ratio_test.jpg')

    #Plot and save the MC comparison and ratio plots together
    c = TCanvas('c', 'hist', 1200, 800)
    c.Draw()
    p1 = TPad()
    p1.Draw()
    p1.Divide(1, 2)
    p11 = p1.cd(1)
    hs.Draw('nostack')
    legend = TLegend(0.9, 0.9, 0.7, 0.7)
    legend.AddEntry(h_Data1, "Data")
    legend.AddEntry(h_MC1, "MC")
    legend.Draw('right')
    p11.Draw('E')
    p12 = p1.cd(2)
    h_Data2.Draw("HIST")
    p12.Draw()
    c.SaveAs(variables[i] + 'test.pdf')