import numpy as np
import matplotlib.pyplot as plt
from iminuit import Minuit

evt = np.load('dimuon.npy')
evt = evt[abs(evt-3.1)<0.5]

xmin, xmax, xbinwidth = 2.6, 3.6, 0.01
vy,edges = np.histogram(evt, bins=100, range=(xmin,xmax))
vx = 0.5*(edges[1:]+edges[:-1])
vyerr = vy**0.5

def model(x, ns, nb, mean, sigma, c1):
    
    linear = (1. + c1*x)/((xmax-xmin) + c1*(xmax**2-xmin**2)/2.)
    
    gaussian = 1./(2.*np.pi)**0.5/sigma * \
        np.exp(-0.5*((x-mean)/sigma)**2)

    return ns*gaussian + nb*linear

def fcn(ns, nb, mean, sigma, c1):
    
    L = model(evt, ns, nb, mean, sigma, c1)
    if np.any(L<=0.): return 1E100

    return 2.*(ns+nb)-2.*np.log(L).sum()

m = Minuit(fcn, ns=6000., nb=14000., mean=3.09, sigma=0.04, c1=0.)
m.migrad()
m.minos()
m.print_param()

fig = plt.figure(figsize=(6,6), dpi=80)

plt.plot([xmin,xmax],[0.,0.],c='black',lw=2)
plt.errorbar(vx, vy, yerr = vyerr, fmt = '.')
cx = np.linspace(xmin,xmax,500)
cy = model(cx,m.values['ns'],m.values['nb'],m.values['mean'],m.values['sigma'],m.values['c1'])*xbinwidth
cy_bkg = model(cx,0.,m.values['nb'],m.values['mean'],m.values['sigma'],m.values['c1'])*xbinwidth
plt.plot(cx, cy, c='red',lw=2)
plt.plot(cx, cy_bkg, c='red',lw=2,ls='--')

plt.grid()
plt.show()