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

# # Hello World, `pyhf` style

# **Two bin counting experiment with a background uncertainty**

# In[1]:


import pyhf


# **Returning the observed and expected** $\mathrm{CL}_{s}$

# In[2]:


model = pyhf.simplemodels.uncorrelated_background(
    signal=[12.0, 11.0], bkg=[50.0, 52.0], bkg_uncertainty=[3.0, 7.0]
)
data = [51, 48] + model.config.auxdata
test_mu = 1.0
CLs_obs, CLs_exp = pyhf.infer.hypotest(
    test_mu, data, model, test_stat="qtilde", return_expected=True
)
print(f"Observed: {CLs_obs}, Expected: {CLs_exp}")


# **Returning the observed** $\mathrm{CL}_{s}$, $\mathrm{CL}_{s+b}$, **and** $\mathrm{CL}_{b}$

# In[3]:


CLs_obs, p_values = pyhf.infer.hypotest(
    test_mu, data, model, test_stat="qtilde", return_tail_probs=True
)
print(f"Observed CL_s: {CLs_obs}, CL_sb: {p_values[0]}, CL_b: {p_values[1]}")


# A reminder that 
# $$
# \mathrm{CL}_{s} = \frac{\mathrm{CL}_{s+b}}{\mathrm{CL}_{b}} = \frac{p_{s+b}}{1-p_{b}}
# $$

# In[4]:


assert CLs_obs == p_values[0] / p_values[1]


# **Returning the expected** $\mathrm{CL}_{s}$ **band values**

# In[5]:


import numpy as np


# In[6]:


CLs_obs, CLs_exp_band = pyhf.infer.hypotest(
    test_mu, data, model, test_stat="qtilde", return_expected_set=True
)
print(f"Observed CL_s: {CLs_obs}\n")
for p_value, n_sigma in enumerate(np.arange(-2, 3)):
    print(
        "Expected CL_s{}: {}".format(
            "      " if n_sigma == 0 else f"({n_sigma} σ)",
            CLs_exp_band[p_value],
        )
    )