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 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()#path.dirname(path.realpath('__file__'))

pi=numpy.pi
eps =numpy.finfo(float).eps
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)}
B1MF5 = numpy.array([0.75,2.25,1.5])
B2MF5 = numpy.array([0.425,1.275,0.85])

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(1,1)

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,T2,dI_Pct,A,MF=None):
	''' Fitting function for dephasing time, uses Jesus's calculations of the AC Stark shift from the PRA'''
	Intensity = A* Intensity #Rescale intensity
	DF = numpy.zeros(Intensity.shape)
	# Choose which line to focus on, either MF=4 or 5 in N=1
	if MF ==4:
		lines = interp.interp1d(Line_data[:,0],Line_data[:,2]*1e3,kind='cubic',bounds_error=False,fill_value=(Line_data[0,2],Line_data[-1,2]))
	elif MF ==5:
		lines = interp.interp1d(Line_data[:,0],Line_data[:,1]*1e3,kind='cubic',bounds_error=False,fill_value=(Line_data[0,1],Line_data[-1,1]))
	#iterate over intensity
	for x,I in enumerate(Intensity):
		DI = I * dI_Pct #build a range of intensities
		intensities = numpy.linspace(I-DI/2,I+DI/2,1500)
		
		freq = lines(intensities) #interpolate the transition frequencies
		DF[x] = numpy.amax(freq)-numpy.amin(freq) #difference between max and min
	dephasing = 1./numpy.sqrt((1/T2)**2+(2*numpy.pi*DF)**2) #fitting function with T2
	return dephasing


def trans_fit(x,A,B,C,D,E=0,f=None):
	#return -A*((x-C)**2*numpy.exp(-x/B))+D
	if f==1:
		return A*(B*numpy.exp(-x/C)-(x-D)**2)+E
	elif f==2:
		return A*(B/(x-C)-(x-D)**2)+E
	
def dephase2(Intensity,lines_fitted,T2,dI_Pct):
	''' fitting function for dephasing, uses fitted version of the AC STark shift'''
	DF = numpy.zeros(Intensity.shape)
	Intensity = 1.0*Intensity
	for x,I in enumerate(Intensity):
		DI = I * dI_Pct #build a range of intensities
		intensities = numpy.linspace(I-DI/2,I+DI/2,1500)
		freq =1e6*trans_fit(intensities,*lines_fitted) #calculate transition frequencies
		DF[x] = numpy.amax(freq)-numpy.amin(freq) #difference between max and min
	dephasing = 1./numpy.sqrt((1/T2)**2+(2*numpy.pi*DF)**2) #fitting function with T2
	return dephasing	
	
def dephase3(Intensity,poly,T2,dI_Pct):
	''' fitting function for dephasing, uses polyfitted version of the AC STark shift'''
	DF = numpy.zeros(Intensity.shape)
	Intensity = 1.0*Intensity
	for x,I in enumerate(Intensity):
		DI = I * dI_Pct #build a range of intensities
		intensities = numpy.linspace(I-DI/2,I+DI/2,1500)
		freq =1e6*numpy.polyval(poly,intensities)#calculate transition frequencies
		DF[x] = numpy.amax(freq)-numpy.amin(freq) #difference between max and min
	dephasing = 1./numpy.sqrt((1/T2)**2+(2*numpy.pi*DF)**2) #fitting function with T2
	return dephasing	
	
T4=[]
T5=[]
T5_err=[]
T4_err=[]

F4 =[]
F4_err=[]
for I in [10,30,20]: #end on 20 - best looking dataset
	#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]))

	#PRint the results
	print("for I = "+str(I)+"Dephasing Times are...")
	print("MF = 5:", int(round(popt1[2])), "+/-", int(round(numpy.sqrt(pcov1[2,2]))), "$\mu$s")

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])
	#Guess = [5e3,200e-6,15e3,2000]
	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]))
	#popt2, pcov2 = scipy.optimize.curve_fit(lambda T,D,T0,O,N: Explicit_soln(2*pi*D,T/T0,2*pi*O,T,0.5*pi/(2*pi*numpy.sqrt(O**2+D**2)),N), N1MF4[1:,0], N1MF4[1:, 1], p0=Guess)
	#popt2 = Guess
	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])
	F4_err.append(numpy.sqrt(pcov2[1,1]))
	
	print("for I = "+str(I)+"Dephasing Times are...")
	print("MF = 4:", int(round(popt2[2])), "+/-", int(round(numpy.sqrt(pcov2[2,2]))), "us")


#Plot the result
fig = pyplot.figure()

Intensities = numpy.linspace(0,45e3,11000)
Line_data = numpy.genfromtxt("lines_trimmed.csv",delimiter=',')
ax3=fig.add_subplot(Grid[:,0])

'''MF=+4'''
Coherence,err = optimise.curve_fit(lambda x,A,B,C: dephase(x,A,B,C,4),(B1MF4*6.89+B2MF4*12.2)*1e3,numpy.array(T4)*1e-6,p0=[700e-6,0.01,0.85,],bounds=([500e-6,0.01,0.1],[1.22e-3,0.1,1.]))
Coherence = numpy.append(Coherence,4)
err = numpy.sqrt(numpy.diag(err))
print(Coherence,err)

Coherence[2]=1.0
Coherence[1]=0.03

ax3.plot(Intensities*1e-3,1e3*dephase(Intensities,*Coherence),color=JQC['grayblue'],zorder=2.,label='jesus')
Ints = numpy.linspace(15e3,25e3,2500)
lines = interp.interp1d(Line_data[:,0],Line_data[:,2]*1e3,kind='cubic',bounds_error=False,fill_value=(Line_data[0,2],Line_data[-1,2]))



'''MF=+5'''
Coherence,err = optimise.curve_fit(lambda x,A,B,C: dephase(x,A,B,C,5),(B1MF5*6.89+B2MF5*12.2)*1e3,numpy.array(T5)*1e-6,p0=[500e-6,0.7,0.03],bounds=([300e-6,0.7,0.],[1500e-6,1.2,0.5]))
Coherence= numpy.append(Coherence,5)

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


ax3.set_xlabel(r"Intensity ($\mathrm{kW/cm^2}$)")
ax3.set_ylabel("Dephasing Time (ms)")
ax3.set_xticks([5,10,15,20,25,30,35])
ax3.set_xlim(5,35)
ax3.set_ylim(0,1.)
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['red'],mec=JQC['red'],mfc='w',zorder=1.5,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['blue'],zorder=1.5)

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

pyplot.figure()
grad = numpy.gradient(Line_data[:,1]*1e3,Line_data[1,0]-Line_data[0,0])
grad_func = interp.interp1d(Line_data[:,0],grad,kind='cubic')

pyplot.plot(Intensities*1e-3,grad_func(Intensities))


grad = numpy.gradient(Line_data[:,2]*1e3,Line_data[1,0]-Line_data[0,0])
grad_func = interp.interp1d(Line_data[:,0],grad,kind='cubic')

pyplot.plot(Intensities*1e-3,grad_func(Intensities))
ax = pyplot.gca()
ax.axvspan(20*0.97,20,color=(1,0,0,0.3))
pyplot.xlabel("DT intensity (kW/cm$^2$)")
pyplot.ylabel("df/dI (Hz/(W/cm$^2$))")


pyplot.figure()

line_Centres =[980.337,980.325,980.332]
line_Centres_err =[0.003,0.001,0.002]
intensities = [1.5*9.4+0.85*12.5,1.875*9.4+1.063*12.5,0.75*9.4+0.425*12.5]

freq = numpy.array([980.347,980.347,980.347,980.347,980.347,980.347,980.347,980.335,980.347])

pyplot.errorbar(B1MF4*6.89+B2MF4*12.2,(freq-numpy.array(F4)/(2*pi)**2),yerr=numpy.array(F4_err)/(2*pi)**2,fmt='o',color='k')
pyplot.plot(Line_data[:,0]*1e-3,Line_data[:,2]*1e-3,label='jesus')

Curve_trans,err_trans = optimise.curve_fit(lambda x,A,B,C,D,E: trans_fit(x,A,B,C,D,E,1),numpy.append(B1MF4*6.89+B2MF4*12.2,0),numpy.append(freq-numpy.array(F4)/(2*pi)**2,980.32047),p0=[0.3e-3,600,18,27,980.2],bounds=([0.01e-3,450,11,25,980.],[7e-3,2700,50,50.5,980.6]),sigma=numpy.append(numpy.array(F4_err)/(2*pi)**2,100e-6),absolute_sigma=True)
print Curve_trans
err_trans = numpy.sqrt(numpy.diag(err_trans))

Curve_trans=numpy.append(Curve_trans,1)

pyplot.plot(Intensities/1e3,trans_fit(Intensities/1e3,*Curve_trans),label="A*(B*exp(-x/C)-(x-D)^2)+E")
fitted_times,err_fitted = optimise.curve_fit(lambda x,A,B: dephase2(x,Curve_trans,A,B),(B1MF4*6.89+B2MF4*12.2),numpy.array(T4)*1e-6,p0=[700e-6,0.03],bounds=([100e-6,0.],[1500e-6,0.5]))
print fitted_times

ax3.plot(Intensities*1e-3,1e3*dephase2(Intensities*1e-3,Curve_trans,*fitted_times),color='orange',zorder=2.,label="A*(B*exp(-x/C)-(x-D)^2)+E")

Curve_trans,err_trans = optimise.curve_fit(lambda x,A,B,C,D,E: trans_fit(x,A,B,C,D,E,2),numpy.append(B1MF4*6.89+B2MF4*12.2,0),numpy.append(freq-numpy.array(F4)/(2*pi)**2,980.32047),p0=[1e-7,2900,-10,20,980],bounds=([1e-9,100,-25,5,970],[0.1,5000,0,30,985]),sigma=numpy.append(numpy.array(F4_err)/(2*pi)**2,100e-6),absolute_sigma=True)
err_trans = numpy.sqrt(numpy.diag(err_trans))
Curve_trans=numpy.append(Curve_trans,2)
print Curve_trans

pyplot.plot(Intensities/1e3,trans_fit(Intensities/1e3,*Curve_trans),label='A*(B/(x-C)-(x-D)^2)+E')
fitted_times,err_fitted = optimise.curve_fit(lambda x,A,B: dephase2(x,Curve_trans,A,B),(B1MF4*6.89+B2MF4*12.2),numpy.array(T4)*1e-6,p0=[700e-6,0.03],bounds=([100e-6,0.],[1500e-6,0.5]))

print fitted_times
ax3.plot(Intensities*1e-3,1e3*dephase2(Intensities*1e-3,Curve_trans,*fitted_times),color=JQC['green'],zorder=2.,label='A*(B/(x-C)-(x-D)^2)+E')

poly,polycov = numpy.polyfit(numpy.append(B1MF4*6.89+B2MF4*12.2,0),numpy.append(freq-numpy.array(F4)/(2*pi)**2,980.32047),3,w=1./numpy.append(numpy.array(F4_err)/(2*pi)**2,100e-6),cov = True)
polyerr = numpy.sqrt(numpy.diag(polycov))

fitted_times,err_fitted = optimise.curve_fit(lambda x,A,B: dephase3(x,poly,A,B),(B1MF4*6.89+B2MF4*12.2),numpy.array(T4)*1e-6,p0=[700e-6,0.03],bounds=([100e-6,0.],[1500e-6,0.5]))
err_fitted = numpy.sqrt(numpy.diag(err_fitted))

print fitted_times
pyplot.plot(Intensities/1e3,numpy.polyval(poly,(Intensities/1e3)),label='3rd order poly')
fitted_times[1]=0.035
fitted_times[0]=800e-6
ax3.plot(Intensities*1e-3,1e3*dephase3(Intensities*1e-3,poly,*fitted_times),color=JQC['red'],zorder=2.,label='3rd order poly')
ax3.legend()

pyplot.legend()

pyplot.ylabel("Transition Frequency (MHz)")
pyplot.xlabel("Peak Intensity (kW/cm$^2$)")

pyplot.show()
