'''
* Open a terminal, and run:

source /usr/remote/python/setup.sh

* or, if that command fails:

source /usr/remote/python/setup.csh

* Then start up ipython with interactive plotting enabled.

ipython --pylab

* In ipython, run this code

execfile('goodfit.py')
'''


'''
Load the data
'''
x, y, u = loadtxt('lineData.txt', unpack=True)

'''
Plot the data with errorbars, even though they'll be too small to see.
'''
f0 = figure(0)
clf()
errorbar(x, y, u, fmt='ko')
xlabel('x')
ylabel('y')
xlim(-0.5, 10.5)

'''
weighted (w=1/uncertainty**2) fit
'''

'''
You can change the order of the fit here!
'''
#||||||||||||||||||
#VVVVVVVVVVVVVVVVVV

order = 1 # <- <- <-

#^^^^^^^^^^^^^^^^^^
#||||||||||||||||||

coef, covariance = polyfit(x, y, order, w=1/u**2, cov=True) 
print '-------------------------'
print 'Polynomial coefficients'
for i in range(len(coef)):
    print 'coef[',i,'] = ', coef[i], '+/-', sqrt(covariance[i,i])

'''
Calculate the model (polynomial) and plot it as a red dashed line.
'''
model = polyval(coef, x)
f1 = figure(1)
clf()
errorbar(x, y, u, fmt='ko')
plot(x, model, 'r--')
xlabel('x')
ylabel('y')
xlim(-0.5, 10.5)

'''
Generate the residual plot
'''
residual = y - model

f2 = figure(2)
clf()
errorbar(x, residual, u, fmt='ko')
plot([-0.5,10.5], [0,0], 'r--')
xlabel('x')
ylabel('Residual = y - Model')
xlim(-0.5, 10.5)

'''
Generate plot of normalized residuals
'''
nresidual = residual / u
'''
Uncertainties are now 1
'''
dnresidual = ones(len(u))


f3 = figure(3)
clf()
errorbar(x, nresidual, dnresidual, fmt='ko')
plot([-0.5,10.5], [0,0], 'r--')
xlabel('x')
ylabel('Normalized Residual = (y - Model) / Uncertainty')
xlim(-0.5, 10.5)




'''
Calculate Chi-squared
'''
chi2 = sum((residual/u)**2)
print '-------------------------'
print 'Quality of fit report'
print '-------------------------'
N = len(x) # number of data points
k = order + 1 # number of model parameters
dof = N - k # "degrees of freedom"
print 'Number of data points = ', N
print 'chi2 = ', chi2
print 'Reduced chi2 = ', chi2 / float(N - k)
print 'Bayes Information Criterion (BIC) = ', chi2 + k * log(N)