import pylab
import numpy
from scipy.optimize import curve_fit

# data load
ang1,eang,p1,ep=pylab.loadtxt('/Users/macbook_franc/Documents/stash/dida14_15/lab2_14_15/malus/ang1.txt',unpack=True)


# data plot 
pylab.rc('font',size=16)
pylab.xlabel('$\\theta$  [degrees]')
pylab.ylabel('S  [mV]')
pylab.title('One polaroid with two free parameters')
pylab.xlim(0,180)
pylab.ylim(-2,45)
pylab.minorticks_on()

pylab.errorbar(ang1,p1,ep,eang,linestyle = '', color = 'black', marker = '.')

# set the constant parameter(s)
B = -0.1

# define the fit function (NOTE THE INDENT!)
def fit_function(ang1, A, theta0):
    return  (A*(numpy.cos((2*numpy.pi/360)*(ang1-theta0)))**2+B)

# set the array of initial value(s)
initial_values = (50,90)

# determine the sigma through error propagation
epfit=numpy.sqrt(ep**2+((2*numpy.pi/360)*eang*initial_values[0]*numpy.sin(2*(2*numpy.pi/360)*(ang1-initial_values[1])))**2)


# call the minimization procedure (NOTE THE ARGUMENTS)
pars, covm = curve_fit(fit_function, ang1, p1, initial_values, epfit)

# calculate the resulting chisquare and the number of dof
chisq = numpy.sum(((p1-fit_function(ang1,pars[0],pars[1]))/epfit)**2)
ndof = len(ang1) - len(pars)

# plot the fitting curve
func_grid = numpy.linspace(0,180, 100)
pylab.plot(func_grid, fit_function(func_grid, pars[0],pars[1]), color = 'red')

# add the annotations
txt1=('A =  (%.1f'%pars[0]+'$\pm$ '+'%.1f'%covm[0,0]+') mV')
txt2=('$\\theta_0$ =  (%.1f'%pars[1]+'$\pm$ '+'%.1f'%covm[1,1]+') degrees')
txt3=('norm.cov. = %.2f'%(covm[0,1]/numpy.sqrt(covm[0,0]*covm[1,1])))
txt4=('chisq/ndof = %.0f'%chisq+'/%.0f'%ndof)

pylab.text(50,5,(txt1+'\n'+txt2+'\n'+txt3+'\n'+txt4),fontsize=14)

pylab.show()