# load data
import pandas as pd
df = pd.read_csv ('data.csv')
x=df.get('position').to_numpy()
y=df.get('counts').to_numpy()

# simple way to convert numbers to the float format
x=x*1.0
y=y*1.0

# import modules needed for data analysis
import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit

# First we plot our data with error bars
plt.figure() 
#plt.errorbar(x,y, yerr=.005, xerr=3, fmt='o', label='data')
plt.plot(x,y, 'o', label='data')
plt.ion()

# We will attempt to fit data with inverse square dependence
def model(x, I0, a_s, d_s, x0, Background):
    pi=np.pi;
    lmbda=532.0; # in nm
    l=50; # in cm
    # xsingle = pi*a/lmbda*x/l;
    # xdouble = pi*d/lmbda*x/l;
    xsingle=a_s*(x-x0)
    xdouble=d_s*(x-x0)
    single_slit = np.power( np.sin(xsingle)/xsingle, 2);
    double_slit = np.power( np.cos(xdouble), 2);
    return (I0*double_slit*single_slit + Background)

## When you do nonlinear fitting, having a good guess is important.
# Otherwise fit my converge to unexpected parameters
guessParameters=[350, np.pi/400, np.pi/60, 505.0, 150.0]; # [I0, a_s, d_s, x0, Background ] coefficients

xf=np.linspace(100, 800, 1000)
plt.plot(xf, model(xf, *guessParameters), 'b--', label='guess fit: I0=%.4g, a_s=%.4g, d_s=%.4g, x0=%.4g, Background=%.4g' % tuple(guessParameters))
plt.legend()
# plt.show() , we might delay or not the plot rendering
# We might repeat it several time to have a reasonable guess which is close to the data
plt.savefig('data_with_a_guess_line.svg')

## Finally we do the fitting
finalParameters, finalParametersErrors = curve_fit(model, x, y, guessParameters)

# lets add the best fit line to the plot
plt.figure()
plt.clf()
plt.plot(x,y, 'o', label='data')
plt.plot(xf, model(xf, *finalParameters), 'r-', label='best fit: I0=%.4g, a_s=%.4g, d_s=%.4g, x0=%.4g, Background=%.4g' % tuple(finalParameters))
plt.legend()
plt.show(); # time to see the final plot
plt.savefig('data_with_and_fit_line.svg')


## Output fit parameter values and their uncertainties
print("Model parameters:")
print("I0  =", finalParameters[0], " +/- ",np.sqrt(finalParametersErrors[0,0]))
print("a_s  =", finalParameters[1], " +/- ",np.sqrt(finalParametersErrors[1,1]))
print("d_s  =", finalParameters[2], " +/- ",np.sqrt(finalParametersErrors[2,2]))
print("x0  =", finalParameters[3], " +/- ",np.sqrt(finalParametersErrors[3,3]))
print("Background  =", finalParameters[4], " +/- ",np.sqrt(finalParametersErrors[4,4]))