# Fit polynomial:
# Fit with the python interface to Minuit 2 called iminuit
# https://iminuit.readthedocs.io/en/stable/


from matplotlib import pyplot as plt
import numpy as np

# Data
x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype='d')
dx =  np.array([0.1,0.1,0.5,0.1,0.5,0.1,0.5,0.1,0.5,0.1], dtype='d')
y = np.array([1.1 ,2.3 ,2.7 ,3.2 ,3.1 ,2.4 ,1.7 ,1.5 ,1.5  ,1.7 ], dtype='d')
dy = np.array([0.15,0.22,0.29,0.39,0.31,0.21,0.13,0.15,0.19,0.13], dtype='d')

# Define fit function
def pol3(x,a0, a1, a2, a3):
    return a0 + x*a1 + a2*x**2 + a3*x**3

# least-squares function: sum of data residuals squared
def LSQ(a0, a1, a2, a3):
    return np.sum((y - pol3(x, a0, a1, a2, a3)) ** 2 / dy ** 2)


# import Minuit object

from iminuit import Minuit

LSQ.errordef = Minuit.LEAST_SQUARES

# Minuit instance using LSQ function to minimize
m = Minuit(LSQ,a0=-1.3, a1=2.6 ,a2=-0.24 ,a3=0.005)

# run migrad
m.migrad()
print(m.values , m.errors)

# Get contour
m.draw_mncontour("a2", "a3")

# Improve the fit
m.hesse()
m.minos()

# access fit results

print(m.values,m.errors)
a0_fit = m.values["a0"]
a1_fit = m.values["a1"]
a2_fit = m.values["a2"]
a3_fit = m.values["a3"]

# Generate fit curve with high resolution

x_fit = np.linspace(x.min(), x.max(), 300)
y_fit = pol3(x_fit,a0_fit, a1_fit, a2_fit, a3_fit)


# plot Data in matplotlib
plt.figure()
plt.errorbar(x, y, dy , dx, fmt="o", label='data')
plt.plot(x_fit, y_fit, label='fit')
plt.title("iminuit Fit Test")
plt.xlim(-0.1, 10.1)
plt.legend()
plt.show()