#  ptolFit.py
#  G. Cowan / RHUL Physics / October 2017
#  Program to fit refraction data from Ptolemy

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import curve_fit
from scipy.stats import chi2

# define fit function
numFunc = 3         # select 1, 2 or 3
if numFunc == 1:
    numPar = 1
    def func (x, *theta):
        alpha = theta
        return alpha*x
elif numFunc == 2:
    numPar = 2
    def func (x, *theta):
        alpha, beta = theta
        return alpha*x - beta*x**2
elif numFunc == 3:
    numPar = 1
    def func (x, *theta):
        r = theta
        return np.arcsin(np.sin(x*np.pi/180)/r)*180/np.pi      # r = n_r/n_i

# set data values (units = degrees)
x   = np.array([10, 20, 30, 40, 50, 60, 70, 80])           # theta_i
y   = np.array([8, 15.5, 22.5, 29, 35, 40.5, 45.5, 50])    # theta_r
numPoints = len(x)
sig = np.array(numPoints*[0.5])

# set default parameter values and do the fit
p0 = np.array(numPar*[1.0])
thetaHat, cov = curve_fit(func, x, y, p0, sig, absolute_sigma=True)

# Retrieve minimized chi-squared, etc.
ndof = numPoints - numPar
chisq = sum(((y - func(x, *thetaHat))/sig)**2)
print ("chisq = ", chisq, ",     ndof = ", ndof)
pval = chi2.sf(chisq, ndof)
print ("pval = ", pval)

# Print fit parameters and covariance matrix
print ("\n", "Fitted parameters and standard deviations:")
sigThetaHat = np.sqrt(np.diag(cov))
for i in range(len(thetaHat)):
    print ("thetaHat[", i, "] = ", thetaHat[i], "  +-  ", sigThetaHat[i])

print ("\n", "i, j, cov[i,j], rho[i,j]:")
for i in range(len(thetaHat)):
    for j in range(len(thetaHat)):
        rho = cov[i][j] / (sigThetaHat[i]*sigThetaHat[j])
        print (i, "  ", j, "  ", cov[i][j], "  ", rho)

# Set up plot
matplotlib.rcParams.update({'font.size':14})     # set all font sizes
plt.clf()
fig, ax = plt.subplots(1,1)
plt.gcf().subplots_adjust(bottom=0.15)
plt.gcf().subplots_adjust(left=0.15)
plt.errorbar(x, y, yerr=sig, xerr=0, color='black', fmt='o', markersize=2, label='data')
plt.xlabel(r'$\theta_{\rm i}$  (degrees)')
plt.ylabel(r'$\theta_{\rm r}$  (degrees)', labelpad=5)
xMin = 0
xMax = 90
yMin = 0
yMax = 60
plt.xlim(xMin, xMax)
plt.ylim(yMin, yMax)
xPlot = np.linspace(xMin, xMax, 100)        # enough points for a smooth curve
fit = func(xPlot, *thetaHat)
if numFunc == 1:
    plotLabel = r'$\theta_{\rm r} = \alpha \theta_{\rm i}$'
elif numFunc == 2:
    plotLabel = r'$\theta_{\rm r} = \alpha \theta_{\rm i} - \beta \theta_{\rm i}^2$'
elif numFunc == 3:
    plotLabel = r'$\theta_{\rm r} = \sin^{-1}[(\sin \theta_{\rm i})/r]$'
plt.plot(xPlot, fit, 'red', linewidth=2, label=plotLabel)
# Tweak legend
handles, labels = ax.get_legend_handles_labels()
handles = [handles[1], handles[0]]
labels = [labels[1], labels[0]]
handles = [handles[0][0], handles[1]]      # turn off error bar for data in legend
plt.legend(handles, labels, loc='upper left', fontsize=14, frameon=False)

#  add fit info to plot
if numFunc == 1:
    plt.figtext(0.2, 0.68, r'$\chi^2 / n_{\rm dof} = 134.6 / 7$')
elif numFunc == 2:
    plt.figtext(0.2, 0.68, r'$\chi^2 / n_{\rm dof} = 0.0 / 6$')
elif numFunc == 3:
    plt.figtext(0.2, 0.68, r'$\chi^2 / n_{\rm dof} = 14.0 / 7$')

# Make and store plot
plt.show()
fileName = "ptolFit_" + str(numFunc) + ".eps"
plt.savefig(fileName, format='eps')