import matplotlib.pyplot as plt
import numpy as np
from ROOT import TFile, TH1F

def main():

    sumet()
    metres()

def metres():

    def res(x):

        return (-3.79944e-10*x + 3.13886E-06*x + 0.00125496)*x + 36.133

    metrange = np.linspace(0, 5000, 30)
    y = []

    for x in metrange:

        y.append(res(x))

    fig = plt.figure()
    plt.scatter(metrange, y, color='steelblue')
    plt.xlabel(r'$E_{T}^{miss}$ (Truth + pileup) [GeV]')
    plt.ylabel(r'$E_{T}^{miss}$ resolution [GeV]')
    plt.savefig('metresolution.png', dpi=500, quality=95)    

def sumet():

    tfile = TFile('sumetPU_mu200_ttbar_gold.root')
    sumeth = tfile.Get('h_sumetPU')
    sumeth.Rebin(4)

    nbins = sumeth.GetNbinsX()

    x = []
    y = []
    hist = []

    for bin in range(1, nbins+1):

        bincenter = sumeth.GetBinCenter(bin) - (sumeth.GetBinWidth(bin)/2.0)
        bincontent = sumeth.GetBinContent(bin)
        #print('{}: {}'.format(bincenter, bincontent))
        for n in range(int(bincontent)):
            hist.append(bincenter)
        #y.append(sumeth.GetBinContent(bin))
        #print(sumeth.GetBinContent(bin))
        x.append(bincenter)
        y.append(bincontent)

    fig = plt.figure()
    #plt.hist(y, nbins, histtype='step')
    #plt.scatter(x, y)
    plt.hist(hist, np.linspace(0, 5000, nbins), histtype='step', linewidth=1.2, color='steelblue')
    plt.xlabel(r'Missing $E_{T}$ from pileup [GeV]')
    plt.ylabel('Events')

    plt.savefig('met.png', dpi=500, quality=95)    






if __name__=='__main__':
    main()