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

evt = np.load('dimuon.npy')

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, norm, mean, sigma, c0, c1):
    
    linear = c0 + c1*(x-xmin)/(xmax-xmin)
    
    gaussian = norm*xbinwidth/(2.*np.pi)**0.5/sigma * \
        np.exp(-0.5*((x-mean)/sigma)**2)
    
    return gaussian + linear

def fcn(norm, mean, sigma, c0, c1):
    expt = model(vx, norm, mean, sigma, c0, c1)
    delta = (vy-expt)/vyerr
    return (delta[vy>0.]**2).sum()

m = Minuit(fcn, norm=6000., mean=3.09, sigma=0.04, c0=200., 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['norm'],m.values['mean'],m.values['sigma'],m.values['c0'],m.values['c1'])
cy_bkg = model(cx,0.,m.values['mean'],m.values['sigma'],m.values['c0'],m.values['c1'])
plt.plot(cx, cy, c='red',lw=2)
plt.plot(cx, cy_bkg, c='red',lw=2,ls='--')

plt.grid()
plt.show()