from __future__ import division
import os
import sys
import numpy
import scipy.optimize
import matplotlib.pyplot as pyplot
import matplotlib.patheffects as PathEffects
from matplotlib.gridspec import GridSpec
from matplotlib.patches import ConnectionPatch,Rectangle
from matplotlib.lines import Line2D
import scipy.interpolate as interp
import scipy.optimize as optimise
from scipy.stats import skewnorm
from scipy.stats import chisquare as chi_squared
from mpl_toolkits.axes_grid1.inset_locator import BboxPatch, BboxConnector,BboxConnectorPatch
from matplotlib.transforms import TransformedBbox

from matplotlib import rcParams
rcParams['font.family'] = 'sans-serif'
rcParams['font.sans-serif'] = ['Arial']

#setting the directory to wherever this script is
directory = os.getcwd()
#define some constants for ease
pi=numpy.pi #pi = 3.141....
eps =numpy.finfo(float).eps #machine floating point limit (smallest number allowed)

#some colours for the plot (JQC colour scheme)
JQC = {'red'     :(198.0/255.0, 62.0/255.0, 98.0/255.0), \
       'blue'    :(0.0/255.0, 70.0/255.0, 127.0/255.0), \
       'purple'  :(126.0/255.0, 29.0/255.0, 123.0/255.0), \
       'sand'  :(244./255.0, 234./255.0, 168./255.0), \
       'grayblue'  :(145./255.0, 148./255.0, 182./255.0), \
       'green'   :(45.0/255.0, 159.0/255.0, 60.0/255.0),
	   'reddish': (198./255,131./255,136./255),\
	   'purpley': (209./255,174./255,208./255)}

#DT Set point voltages	,B1,B2 for N=1,MF=+5 dataset
B1MF5 = numpy.array([0.75,2.25,1.5])
B2MF5 = numpy.array([0.425,1.275,0.85])
#DT Set point voltages	,B1,B2 for N=1,MF=+4 dataset
B1MF4 = numpy.array([0.75,0.75,1.125,1.2,1.3,1.37,1.5,1.875,1.4])
B2MF4 = numpy.array([0.425,0.425,0.6375,0.68,0.737,0.775,0.85,1.063,0.793])

#Import style file
pyplot.style.use(r'Wide_Origin_style.mplstyle')
Grid=GridSpec(2,2)

def DampedRamseySin(t, N0, f, t0, Phi=-pi/4):
    '''Damped sine wave used to quantify dephasing'''
    N = 0.5*N0*numpy.exp(-(t/t0))*numpy.sin(t*f/(2*numpy.pi) + Phi) + N0/2
    return N

def gaussian(x, mu, sig):
    '''Gaussian distribution'''
    return numpy.exp(-(x - mu)**2. / (2 * sig**2.))

def Ramsey_Frequency(detuning,Omega,T_Pulse,T_evolution,N_mol,phi=0.):
	''' Our lineshape, derived in Ramsey_calculations.nb, makes no assumptions about the intensity of the light.'''
	detuning = detuning * 2.*pi
	Omega_eff = numpy.sqrt(Omega**2.+detuning**2.)

	line = (3*(Omega**4 + Omega_eff**4) + 2*Omega_eff**2*detuning**2 + 3*detuning**4 + (Omega**4 + (Omega_eff**2 - detuning**2)**2)*numpy.cos(2*Omega_eff*T_Pulse) + 4*numpy.cos(Omega_eff*T_Pulse)*(-Omega**4 + Omega_eff**4 - detuning**4 + \
		Omega**2*(Omega_eff**2 + detuning**2)*(-1 + numpy.cos(Omega_eff*T_Pulse))*numpy.cos(phi + T_evolution*detuning)) + 8*Omega**2*numpy.sin((Omega_eff*T_Pulse)/2.)**2*((-Omega_eff**2 + detuning**2)*numpy.cos(phi + T_evolution*detuning) + \
		2*Omega_eff*detuning*numpy.sin(Omega_eff*T_Pulse)*numpy.sin(phi + T_evolution*detuning)))/(8.*Omega_eff**4)

	return N_mol*(line)

def dephase(Intensity,lines,T2,dI_Pct):
	''' fit function for the dephasing time, eq. 21'''
	DF = numpy.zeros(Intensity.shape)

	for x,I in enumerate(Intensity):
		DI = I * dI_Pct
		intensities = numpy.linspace(I-DI/2,I+DI/2,1500)

		freq_fitted = numpy.polyval(lines,intensities)
		DF[x] = 1e6*(numpy.amax(freq_fitted)-numpy.amin(freq_fitted))

	dephasing = 1./numpy.sqrt((1/T2)**2+(2*numpy.pi*DF)**2)
	return dephasing

def chi_squared2(O,E,err,dof):
	X = (O-E)**2/err**2
	X = numpy.sum(X)
	Xsquare = X/(len(O)-dof)
	return Xsquare	
	
	
	
#Lists for storing dephasing times
T4=[]
T5=[]
T5_err=[]
T4_err=[]

#Lists for storing detunings
F4 =[]
F4_err=[]
F5 =[]
F5_err=[]


for I in [10,30,20]: #end on 20 
	#Get the data for the N = 0, MF = 5 <-> N = 1, MF = 5 first
	N1MF5 = numpy.genfromtxt(str(directory)+"\\N1MF5\\I"+str(I)+".csv",delimiter=",")

	#Fit the data and normalise
	Guess = numpy.array([1500, 0.44, 100.,-pi/4])
	popt1, pcov1 = scipy.optimize.curve_fit(DampedRamseySin, N1MF5[1:,0], N1MF5[1:, 1], p0=Guess)

	FitTimes5 = numpy.linspace(0, 1500, int(1e3))
	FitNumber5 = DampedRamseySin(FitTimes5, popt1[0], popt1[1], popt1[2],popt1[3])

	T5.append(popt1[2])
	T5_err.append(numpy.sqrt(pcov1[2,2]))
	F5.append(popt1[1]/((2*pi)**2))
	F5_err.append(numpy.sqrt(pcov1[1,1])/((2*pi)**2))

	if I == 20:
		print(6.89*B1MF5[-1]+12.2*B2MF5[-1])
	#PRint the results
	print("......")
	print("for I = "+str(I)+" Dephasing Times are...")
	print("MF = 5:", int(round(popt1[2])), "+/-", int(round(numpy.sqrt(pcov1[2,2]))), r"$\mu$s")
	print("chi squared =", chi_squared(N1MF5[1:, 1]/popt1[0],DampedRamseySin(N1MF5[1:,0],1,*popt1[1:]),4.)[0])


for I in [10,"10a",15,16,18,18.5,20,25,19]: #end on 19
	#Now get the N = 0, MF = 5 <-> N = 1, MF = 4
	N1MF4 = numpy.genfromtxt(str(directory)+"\\N1MF4\\I"+str(I)+".csv", delimiter=",")

	#Fit the data and normalise
	Guess = numpy.array([1500, 0.16, 300, -numpy.pi/4])
	popt2, pcov2 = scipy.optimize.curve_fit(DampedRamseySin, N1MF4[1:,0], N1MF4[1:, 1], p0=Guess,bounds=([1000,0.05,0,-pi/2],[2500,5.,1000,pi/2]))

	FitTimes4 = numpy.linspace(0,1500, int(1e3))
	FitNumber4 = DampedRamseySin(FitTimes4, popt2[0], popt2[1], popt2[2], popt2[3])
	T4.append(popt2[2])
	T4_err.append(numpy.sqrt(pcov2[2,2]))

	F4.append(popt2[1]/((2*pi)**2))
	F4_err.append(numpy.sqrt(pcov2[1,1])/((2*pi)**2))
	print("......")
	if I == 19:
		print(6.89*B1MF4[-1]+12.2*B2MF4[-1])
	print("for I = "+str(I)+" Dephasing Times are...")
	print("MF = 4:", int(round(popt2[2])), "+/-", int(round(numpy.sqrt(pcov2[2,2]))), r"$\mu$s")
	print("chi squared =", chi_squared(N1MF4[1:, 1]/popt2[0],DampedRamseySin(N1MF4[1:,0],1,*popt2[1:]),4.)[0])

print("###########################################")

#Plot the result in small sub-figures first
fig = pyplot.figure()

fig.subplots_adjust(hspace=0,bottom=0.15,right=0.97,top=0.95,left=0.1)
ax1 = fig.add_subplot(Grid[1,0],zorder=5.)
ax2 = fig.add_subplot(Grid[0,0],sharex=ax1,zorder=5.0)

ax1.plot(FitTimes5*1e-3, FitNumber5/popt1[0],color=JQC['grayblue'],zorder=2)
ax1.plot(N1MF5[1:,0]*1e-3, N1MF5[1:, 1]/popt1[0], marker='o',ls='none',mec=JQC['blue'],mfc='w',zorder=3,mew=2.5)
ax2.plot(FitTimes4*1e-3, FitNumber4/popt2[0],color=JQC['reddish'],zorder=2)
ax2.scatter(N1MF4[1:,0]*1e-3, N1MF4[1:, 1]/popt2[0], marker='o',color=JQC['red'],zorder=3)


ax1.set_xlim(0, 120)
ax1.set_ylim(-0, 1.)
ax1.set_xlabel("Evolution Time (ms)")
ax2.set_xlim(0, 1)
ax2.set_ylim(-0., 1.)

ax1.set_ylabel(r"Normalised Number of Molecules in $|0\rangle$",y=1.15,zorder=5.)

ax1.text(0.35,0.1,r"$|1\rangle = |N=1,M_F=+5\rangle$",transform=ax1.transAxes)
ax2.text(0.35,0.1,r"$|1\rangle = |N=1,M_F=+4\rangle$",transform=ax2.transAxes)

Intensities = numpy.linspace(500,45e3,11000)


#main figure
ax3=fig.add_subplot(Grid[:,1])

'''MF=+4'''
#start by fitting the transition frequency data, I use numpy.polyfit because I have found that it is simplest and fastest for evaluating arbitrary polynomials
freq_MF4 = numpy.array([980.347,980.347,980.347,980.347,980.347,980.347,980.347,980.335,980.347]) #array containing the "detuned" frequencies that we used for observing oscillations

poly_MF4,polycov = numpy.polyfit(1e3*numpy.append(B1MF4*6.89+B2MF4*12.2,0),numpy.append(freq_MF4-numpy.array(F4),980.32047),3,w=1./numpy.append(numpy.array(F4_err),100e-6),cov = True)
polyerr_MF4 = numpy.sqrt(numpy.diag(polycov))

Coherence,err = optimise.curve_fit(lambda x,A,B: dephase(x,poly_MF4,A,B),(B1MF4*6.89+B2MF4*12.2)*1e3,numpy.array(T4)*1e-6,p0=[1e-3,0.04],bounds=([500e-6,0],[1.52,0.5]))

err = numpy.sqrt(numpy.diag(err))

print(Coherence,err)
print("Chi-squared =",chi_squared2(numpy.array(T4)*1e-6,dephase((B1MF4*6.89+B2MF4*12.2)*1e3,poly_MF4,*Coherence),numpy.array(T4_err)*1e-6,2))

ax3.plot(Intensities*1e-3,1e3*dephase(Intensities,poly_MF4,*Coherence),color=JQC['reddish'],zorder=1.)

Ints = numpy.linspace(10,20,1500)

'''MF=+5'''
freq_MF5 = numpy.array([980.156,979.98,980.07]) #array containing the "detuned" frequencies that we used for observing oscillations

poly_MF5 = numpy.polyfit(1e3*numpy.append(B1MF5*6.89+B2MF5*12.2,0),numpy.append(freq_MF5,980.231),1.,w=1./numpy.append(numpy.array([1e-3,2e-3,6e-3]),0.01e-3))

Coherence,err = optimise.curve_fit(lambda x,A,B: dephase(x,poly_MF5,A,B),(B1MF5*6.89+B2MF5*12.2)*1e3,numpy.array(T5)*1e-6,p0=[0.7e-3,0.02],bounds=([100e-6,0],[1.52,0.5]))

err = numpy.sqrt(numpy.diag(err))

print(Coherence,err)
print("Chi-squared =",chi_squared2(numpy.array(T5)*1e-6,dephase((B1MF5*6.89+B2MF5*12.2)*1e3,poly_MF5,*Coherence),numpy.array(T5_err)*1e-6,2))


ax3.plot(Intensities*1e-3,1e3*dephase(Intensities,poly_MF5,*Coherence),color=JQC['grayblue'],zorder=1.,ls='--')

ax3.set_xlabel("Peak Intensity, $I_0$ ($\mathrm{kW/cm^2}$)")
ax3.set_ylabel("Coherence Time (ms)")

ax3.set_xticks([5,10,15,20,25,30,35])
ax3.set_xlim(5,35)
ax3.set_ylim(0,0.9)

#plotting intensities from our calibration of the 1550 DT servo IB1 = 6.89 (kW/cm^2)/V, IB2= 12.2 (kW/cm^2)/V
ax3.errorbar(B1MF5*6.89+B2MF5*12.2,numpy.array(T5)*1e-3,yerr=numpy.array(T5_err)*1e-3, fmt='o',ls='none',ecolor=JQC['blue'],mec=JQC['blue'],mfc='w',zorder=3.,mew=2.5)
ax3.errorbar(B1MF4*6.89+B2MF4*12.2,numpy.array(T4)*1e-3,yerr=numpy.array(T4_err)*1e-3,fmt='o',color=JQC['red'],zorder=3.4)

ax2.tick_params(axis='both',zorder=400)

fig.subplots_adjust(hspace=0.22,bottom=0.15,right=0.97,top=0.97,left=0.1,wspace=0.3)

ax3.annotate("(a)",xy=(19.1,1e-3*T4[-1]),xytext=(12.5,0.8),fontsize=15,arrowprops=dict(lw=1.5,facecolor='black',arrowstyle = "->"))
ax3.annotate("(b)",xy=(20.5,1e-3*T5[-1]),xytext=(13.5,0.13),fontsize=15,arrowprops=dict(lw=1.5,facecolor='black',arrowstyle = "->"))

ax1.text(0.9,0.865,"(b)",fontsize=20,transform=ax1.transAxes)
ax2.text(0.9,0.865,"(a)",fontsize=20,transform=ax2.transAxes)
ax3.text(0.01,0.935,"(c)",fontsize=20,transform=ax3.transAxes)

#fig.savefig("dephasing_comparison_zooms_one_axis.eps") #optional savefig for directly saving without having to use python GUI


pyplot.show()
