import hist
import numpy as np
import matplotlib.pyplot as plt
import mplhep as hep

#Plot Btag Efficiency

outputSignal = util.load("btagEfficiency_SignalMC_UL2018_full.coffea")


light_total = outputSignal["hJets"][{"flavor": slice(0j, 4j,sum)}]
c_total = outputSignal["hJets"][{"flavor": slice(4j, 5j,sum)}]
b_total = outputSignal["hJets"][{"flavor": slice(5j, 6j,sum)}]

light_tagged = outputSignal["hBJets"][{"flavor": slice(0j, 4j,sum)}]
c_tagged = outputSignal["hBJets"][{"flavor": slice(4j, 5j,sum)}]
b_tagged = outputSignal["hBJets"][{"flavor": slice(5j, 6j,sum)}]


Eff_1 = light_tagged.project("pt", "eta")[...,::hist.rebin(10),::hist.rebin(30)]/light_total.project("pt", "eta")[...,::hist.rebin(10),::hist.rebin(30)]
Eff_2 = c_tagged.project("pt", "eta")[...,::hist.rebin(10),::hist.rebin(30)]/c_total.project("pt", "eta")[...,::hist.rebin(10),::hist.rebin(30)]
Eff_3 = b_tagged.project("pt", "eta")[...,::hist.rebin(10),::hist.rebin(30)]/b_total.project("pt", "eta")[...,::hist.rebin(10),::hist.rebin(30)]

fig, (ax) = plt.subplots(1,3, figsize=(15, 5), sharey=True, sharex=True)
fig.subplots_adjust(wspace=0.15)

hep.hist2dplot(Eff_1, ax=ax[0], cmin=0.0, cmax=1.0)
hep.hist2dplot(Eff_2, ax=ax[1], cmin=0.0, cmax=1.0)
hep.hist2dplot(Eff_3, ax=ax[2], cmin=0.0, cmax=1.0)
'''
hep.hist2dplot(Eff_1.project("pt", "eta"), ax=ax[0])
hep.hist2dplot(Eff_2.project("pt", "eta"), ax=ax[1])
hep.hist2dplot(Eff_3.project("pt", "eta"), ax=ax[2])
'''
ax[0].set_title("HadronFlavor=0 (light)", fontsize=15)
ax[1].set_title("HadronFlavor=4 (c)", fontsize=15)
ax[2].set_title("HadronFlavor=5 (b)", fontsize=15)

ax[0].set_xlim(0,250)
ax[1].set_xlim(0,250)
ax[2].set_xlim(0,250)

ax[0].set_ylim(0,3)
ax[1].set_ylim(0,3)
ax[2].set_ylim(0,3)

ax[2].set_xlabel(r"$p_{T}$ [GeV]", fontsize=14)
ax[0].set_ylabel(r"$\eta$", fontsize=14)

#ax.set_title("B-tag efficiency")
plt.savefig('test_btag_hep.png')