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

#       
# <img src="https://raw.githubusercontent.com/scikit-hep/mplhep/master/docs/source/_static/mplhep.png" width="70%"/>
# 
# 

# # mplhep

# ## Styling
# - Easy publication grade styling
# - Distribute styles and fonts
# - Save real analyzer time
# 

# ## Plotting
# - extend `mpl` with domain specfic functions
# - maximally align with `mpl` API to be intuitive
# - upstream as much as possible

# ## Histogram plotting
# - `matplotlib` now has basic histogram plotting support - `plt.stairs`
# - some functionality remains too specific to be upstreamed - like histograms with error bars

# In[ ]:


import matplotlib.pyplot as plt
import numpy as np 
np.random.seed(2)


# In[ ]:


H = np.histogram(np.random.normal(2.5, .5, 100), bins=np.arange(0,6, 0.5))
print("Type:", type(H))
h, bins = H
print("Values:", h)
print("Bins:", bins)


# ### Simple Histograms

# In[ ]:


fig, ax = plt.subplots()

ax.stairs(*H, label='Empty')
ax.stairs(h, bins + 4, fill=True, label='Filled')
ax.legend()


# In[ ]:


import mplhep as hep
f, axs = plt.subplots(1,2, figsize=(14, 5))

hep.histplot(H, ax=axs[0])
hep.histplot(h, bins, yerr=True, ax=axs[1]);


# #### Primary goal is to stay unobtrusive
# - if you know how `plt.hist()` works, `mplhep.histplot()` should behave like you'd expect
# - kwargs you are used to should work

# In[ ]:


f, axs = plt.subplots(1,2, figsize=(14, 5))

hep.histplot(H, ax=axs[0], histtype='fill', hatch='///', edgecolor='white')
hep.histplot(H, ax=axs[1], histtype='errorbar', yerr=True, c='black', capsize=4)


# #### and be able to do same things as `plt.hist`

# In[ ]:


f, axs = plt.subplots(1,3, figsize=(21, 5))

hep.histplot([h, h*2], bins=bins, ax=axs[0], yerr=True)
hep.histplot([h, h*2], bins=bins, ax=axs[1], stack=True, histtype='fill')
hep.histplot([h, h*2], bins=bins, ax=axs[2], binwnorm=[20, 9]);


# #### Be convenient for a physicist

# In[ ]:


f, axs = plt.subplots(1,3, figsize=(21, 5))

hep.histplot([h, h*2], bins=bins, ax=axs[0], yerr=True, label=["MC1", "MC2"])
hep.histplot(np.random.poisson(h*3), bins=bins, ax=axs[1], yerr=True, label="Data")

hep.histplot([h, h*2], bins=bins, ax=axs[2], stack=True, label=["MC1", "MC2"], density=True)
hep.histplot(np.random.poisson(h*3), bins=bins, ax=axs[2], yerr=True, histtype='errorbar', label="Data", density=True, color='k')
for ax in axs:
    ax.legend()
axs[0].set_title("Some MCs")
axs[1].set_title("Draw Poisson Data")
axs[2].set_title("Data/MC Shape comparison"); 


# #### Be flexible about input types

# In[ ]:


f, axs = plt.subplots(1,3, figsize=(21, 5))

my_hist = [1,2,3,4,2]
hep.histplot(my_hist, bins=range(len(my_hist)+1), ax=axs[0])
hep.histplot(H, yerr=True, ax=axs[1])
hep.histplot(h, bins, yerr=True, ax=axs[2])

axs[0].set_title("'Manual' inputs")
axs[1].set_title("Numpy Tuple")
axs[2].set_title("Unwrapped Tuple");


# #### Process anything that implements UHI PlottableProtocol

# In[ ]:


import uproot4
from skhep_testdata import data_path
fname = data_path("uproot-hepdata-example.root")
f = uproot4.open(fname)
print(f.keys())
print(f['hpx'])
hep.histplot(f['hpx']);


# In[ ]:


import ROOT
h = ROOT.TH1F("h1", "h1", 50, -2.5, 2.5)
h.FillRandom("gaus", 10000)

hep.histplot(h);


# #### Integrate tightly with `hist`

# In[ ]:


import hist

h = hist.Hist(
    hist.axis.Regular(10, 0.0, 1.0, label='X', name='x'),
)
h.fill(np.random.normal(0.5, 0.2, 1000))

hep.histplot(h);


# In[ ]:


import hist

h = hist.Hist(
    hist.axis.IntCategory([], 
        growth=True, label='Categorical Bins', name='x'),
)
h.fill(np.random.normal(5, 1, 1000))

hep.histplot(h);


# ## 2D histograms
# - minimal wrap on `plt.pcolormesh()`, which alrady has almost everything
# - fix(flip) indexing
# - add some annotation sugar ala `sns.heatmap`

# In[ ]:


h2 = hist.Hist(
    hist.axis.Regular(10, 0.0, 1.0, label='Some label'),
    hist.axis.Regular(10, 0, 1)
)
h2.fill(np.random.normal(0.5, 0.2, 1000), np.random.normal(0.5, 0.2, 1000))
hep.hist2dplot(h2, labels=True, cbar=False);


# In[ ]:


print(f['hpxpy'])
hep.hist2dplot(f['hpxpy']);


# ## Styling
# - Primary purpose of `mplhep` is to serve and distribute styles 
#     - **ALICE**
#     - **ATLAS**
#     - **CMS**
#     - **LHCb**
# - To ensure plots looks the same on any framework fonts need to be included
#  - I am liable to go on a rant, so suffice to say:
#  - We package an open look-alike of Helvetica called Tex Gyre Heros

# In[ ]:


hep.style.use([hep.style.ATLAS])#, {'xtick.direction': 'out'}])
hep.histplot(np.histogram(np.random.normal(10, 3, 1000)));
hep.atlas.label();


# In[ ]:


hep.style.use("CMS")
hep.histplot(np.histogram(np.random.normal(10, 3, 1000)))
hep.cms.label() 


# In[ ]:


hep.style.use()
hep.style.use([hep.style.LHCb2])
hep.histplot(np.histogram(np.random.normal(10, 3, 1000)))
hep.lhcb.label() 


# In[ ]:


hep.style.use()
hep.style.use([hep.style.ALICE])
hep.histplot(np.histogram(np.random.normal(10, 3, 1000)))
hep.alice.label() 


# ###  Label styles

# In[ ]:


hep.style.use()
fig, axs = plt.subplots(1, 5, figsize=(18, 3))
for i, ax in enumerate(axs):
    hep.cms.label(ax=ax, loc=i)


# ## Extended examples

# #### Ratio Plot

# In[ ]:


a = hist.Hist.new.Reg(20,-2,2).Int64().fill(np.random.uniform(-2,2,size=3000))
b = hist.Hist.new.Reg(20,-2,2).Int64().fill(np.random.normal(0,0.5,size=5000))
tot = a + b
data = tot.copy()
data[...] = np.random.poisson(tot.values())

from hist.intervals import ratio_uncertainty


# In[ ]:


fig, (ax, rax) = plt.subplots(2, 1, figsize=(6,6), gridspec_kw=dict(height_ratios=[3, 1], hspace=0), sharex=True)

hep.histplot([a, b], ax=ax, stack=True, histtype='fill', label=["MC1", "MC2"])
hep.histplot(data, ax=ax, histtype='errorbar', color='k', capsize=4, yerr=True, label="Data")

errps = {'hatch':'////', 'facecolor':'none', 'lw': 0, 'color': 'k', 'alpha': 0.4}
ax.stairs(
    values=tot.values() + np.sqrt(tot.values()),
    baseline=tot.values() - np.sqrt(tot.values()),
    edges=tot.axes[0].edges, **errps, label='Stat. unc.')

yerr = ratio_uncertainty(data.values(), tot.values(), 'poisson')
rax.stairs(1+yerr[1], edges=tot.axes[0].edges, baseline=1-yerr[0], **errps)
hep.histplot(data.values()/tot.values(), tot.axes[0].edges, yerr=np.sqrt(data.values())/tot.values(),
    ax=rax, histtype='errorbar', color='k', capsize=4, label="Data")


rax.axhline(1, ls='--', color='k')
rax.set_ylim(0.7, 1.3)
ax.set_xlim(-2, 2)
ax.legend()


# #### 2D Ratio plot

# In[ ]:


a = hist.Hist.new.Reg(5,-2,2).Reg(5,-2,2).Int64().fill(*np.random.normal(0,1,size=(2,1000)))
b = hist.Hist.new.Reg(5,-2,2).Reg(5,-2,2).Int64().fill(*np.random.uniform(-2,2,size=(2,1000)))

ratio = a.values() / b.values()
err_down, err_up = ratio_uncertainty(a.values(), b.values(), 'poisson')

def unctext(ra, u, d):
    ra, u, d = f'{ra:.2f}', f'{u:.2f}', f'{d:.2f}'
    return '$'+ra+'_{-'+d+'}^{+'+u+'}$'

labels = np.vectorize(unctext)(ratio, err_up, err_down)

hep.hist2dplot(ratio, labels=labels, cmap='cividis');


# ### mplhep in publications
# Package has already helped produce plots in several publication
#  - We can create experiment TDR guidelines compatible plots in python
# 
# - [Simultaneous Jet Energy and Mass Calibrations with Neural Networks](https://cds.cern.ch/record/2706189), ATLAS Collaboration, 2019
# - [Integration and Performance of New Technologies in the CMS Simulation](https://arxiv.org/abs/2004.02327), Kevin Pedro, 2020 (Fig 3,4)
# - [GeantV: Results from the prototype of concurrent vector particle transport simulation in HEP](https://arxiv.org/abs/2005.00949), Amadio et al, 2020 (Fig 25,26)
# - [Search for the standard model Higgs boson decaying to charm quarks](https://cds.cern.ch/record/2682638), CMS Collaboration, 2019 (Fig 1)