# last updated 2011-11-10

# this programme creates the data and outputs it to data_plot1.dat
# and then re-imports it to do the analysis

from numpy import random, sqrt, exp, log, arange, size, shape, reshape, mean, std, zeros, identity, ones, log, sin, cos, tanh, pi, clip
from scipy.optimize import leastsq
from pylab import figure, clf, show, xlabel, ylabel, rcParams, pcolor, imshow, cm, clim, meshgrid, colorbar

from csv import reader, writer

from time import clock

from matplotlib import rc
rc('font',**{'family':'serif','serif':['Times New Roman']})

params={'axes.labelsize':18,'xtick.labelsize':18,'ytick.labelsize':18,'figure.figsize':(8,6)}
rcParams.update(params)


startTime = clock()

xmax=10.0 
dx=0.025
x1=arange(-xmax,xmax,dx)
num=size(x1)
c=cos(x1)
c2=100*c*c

filename='data_plot1.dat'

# create data and read to file

offset=random.rand(1)[0]
xs=x1-50.0*offset*pi/180
cs=cos(xs)
cs2=100*cs*cs
y1=cs2+0.1*sqrt(2)*cs2*random.randn(num)

fileout=open(filename,'w')
output = writer(fileout, delimiter=',')
for i in range(0,num):
    row = [x1[i],y1[i]]
    output.writerow(row)
fileout.close()

spacing=0 #skips lines if needed (e.g., on oscilloscope files)
columns=2

fid=open(filename,'U')
inData=reader(fid,delimiter=',')
if spacing!=0:
	fid.seek(spacing)	
data=[]
for i in range(0,columns): data.append([])
	
for row in inData:
	for i in range(0,columns):
		data[i].append(float(row[i]))
fid.close()

x_data=data[0]
y_data=data[1]
num=size(x_data)

binwidth=5
bnum=num/binwidth

ax=reshape(x_data,(bnum,binwidth))
bx=mean(ax,axis=1)

ay=reshape(y_data,(bnum,binwidth))
by=mean(ay,axis=1)
sy=std(ay,axis=1)

# calculate standard error, will be used for error bars in plot

se=sy/sqrt(binwidth)

# fit function
fitfunc = lambda p, fx: 0.5*p[0]*(1.0+cos(2.0*p[1]*fx+p[2]))+ p[3]
p0 = [100.0, 1.0, 0.0, 0.0] # Initial guess for the parameters


# unweighted fit to the binned data 
#errfunc = lambda p, fx, fy: fitfunc(p, fx) - fy # Distance to the target function
#p1,cov,infodict,mesg,ier = leastsq(errfunc,p0[:],args=(bx,by),full_output=1,warning=True)


# unweighted fit to the binned data 
weights=1/(se*se)
w_errfunc = lambda p, fx, fy, w: w*(fitfunc(p, fx) - fy) # Distance to the target function
p1,cov,infodict,mesg,ier = leastsq(w_errfunc,p0[:],args=(bx,by,weights),full_output=1,warning=True)

print p1,cov.diagonal()


ty=fitfunc(p1,x1)


print size(x_data), size(ty)
tby=fitfunc(p1,bx)


r=by-tby
nr=(by-tby)/se
nr2=nr*nr
ssqs=sum(nr2)
chi2=ssqs/bnum

print ssqs

phi=-p1[2]*180.0/pi
if chi2<3.0:
    err_phi=cov.diagonal()[2]*ssqs*180/pi
else:
    err_phi=cov.diagonal()[2]*180/pi


# the rest of the programme is just to create the plot


fig = figure(1,facecolor='white')

dpi=[196.0/255.0,59.0/255.0,142.0/255.0] #  Durham pink 
dr=[170.0/255.0,43.0/255.0,74.0/255.0] #  Durham red 
dy=[232.0/255.0,227.0/255.0,145.0/255.0] #  Durham yellow 
dg=[159.0/255.0,161.0/255.0,97.0/255.0] # Durham green
db=[0,99.0/255.0,136.0/255.0] # Durham blue
dp=[126.0/255.0,49.0/255.0,123.0/255.0] # Durham purple
dv=[216.0/255.0,172.0/255.0,244.0/255.0] # Durham violet

clf()

ms=8 # makrer size

# main plot

ax = fig.add_axes([0.15, 0.6, 0.6, 0.3]) # full plot
p1a=ax.plot(x_data,y_data,color=dg,linewidth=1.0) 
p1b=ax.plot(x_data,c2,'black',linewidth=0.5)
p1c=ax.plot(x_data,ty,color=dpi,linewidth=2.0)
p1d=ax.plot(bx,by,'.',markersize=ms,color=db)

# add error bars manually but only for every 5th point

for ept in range(0,bnum,5):
    ax.plot([bx[ept],bx[ept]],[by[ept]-se[ept],by[ept]+se[ept]],color=db,linewidth=3)

ax.set_xticklabels([])
ax.set_xlim(-2.25*pi,2.25*pi)
ax.set_ylim(-5,115)
ylabel('Counts (/s)')

# This is the top right plot which we will leave blank.

ax = fig.add_axes([0.76, 0.6, 0.2, 0.3]) 
ax.axison = False
ax.set_xlim(0.0,1.0)
ax.set_ylim(0.0,1.0)
ax.text(0.0,0.5,'$\phi =$%.2f' % phi, fontsize=18,color=[0, 0, 0])
ax.text(0.62,0.51,'$\pm$%.2f$^\circ$' % err_phi, fontsize=18,color=[0, 0, 0])


# plot residuals

ax = fig.add_axes([0.15, 0.35, 0.6, 0.25])
p3a=ax.plot(x_data,y_data-ty,color=dg,linewidth=1.0)
p3b=ax.plot(bx,r,'.',markersize=ms,color=db)

# add error bars to residuals plot

for ept in range(0,bnum,5):
    ax.plot([bx[ept],bx[ept]],[by[ept]-se[ept]-tby[ept],by[ept]+se[ept]-tby[ept]],color=db,linewidth=3)

ax.set_xticklabels([])
ax.set_yticks([-10, 0, 10])
ax.set_yticklabels([-10, 0, 10])
ax.set_xlim(-2.25*pi,2.25*pi)
ax.set_ylim(-18,18)
ylabel('Residual')

# histogram of residuals

ax = fig.add_axes([0.76, 0.35, 0.2, 0.25])
p4b=ax.hist(r,40,range=(-20,20),orientation='horizontal',facecolor=dv,edgecolor=dv)

ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_xlim(0,num/10)
ax.set_ylim(-18,18)

# plot of nomalised residulas

ax = fig.add_axes([0.15, 0.1, 0.6, 0.25])
p5a=ax.plot(x_data,y_data-ty,color=dg,linewidth=1.0)
p5b=ax.plot(bx,nr,'.',markersize=ms,color=db)
ax.set_xticks([-2*pi, -pi, 0, pi, 2*pi])
ax.set_xticklabels([-360, -180, 0, 180, 360])
ax.set_yticks([-10, 0, 10])
ax.set_yticklabels([-10, 0, 10])
ax.set_xlim(-2.25*pi,2.25*pi)
ax.set_ylim(-18,18)
xlabel('Polarizer angle ($^\circ$)')
ax.text(-2.25*pi-2.5,-15,'Normalised',rotation=90,fontsize=18,color=[0, 0, 0])
ylabel('Residual')

# Lower right plot: Histogram of the normalised residuals

ax = fig.add_axes([0.76, 0.1, 0.2, 0.25])
p6a=ax.hist(nr,40,range=(-20,20),orientation='horizontal',facecolor=dv,edgecolor=dv)

gy=arange(-10.0,10.0,0.1)
gx=bnum*1/sqrt(2*pi)*exp(-gy*gy/2)
p6b=ax.plot(gx,gy,color=[0.0,0.0,0.0],linewidth=2)
ax.text(10,8,'$\chi^2=$ %.1f' % chi2,fontsize=18,color=[0, 0, 0])
#ax.text(45,10,round(chi2,1),fontsize=18,color=[0, 0, 0])

ax.set_xticks([0, 25,50])
ax.set_xticklabels([0, 25, 50])
ax.set_yticklabels([])
ax.set_xlim(0,num/10)
ax.set_ylim(-18,18)
xlabel('Occurrence')

fig.savefig('test.png',dpi=160,facecolor='None')
show()

print 'elapsed time: ', (clock() - startTime)