from __future__ import division

import numpy
import scipy.optimize
import matplotlib.pyplot as pyplot
from matplotlib.gridspec import GridSpec
from matplotlib.patches import ConnectionPatch,Rectangle
from matplotlib.lines import Line2D
import scipy.interpolate as interp

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)}


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

def DampedRamseySine(t, N0, f, t0, Phi):
    N = 0.5*N0*numpy.exp(-t/t0)*numpy.sin(t*f/(2*numpy.pi) + Phi) + N0/2
    return N

T5 =[]
T4=[]

T5_err=[]
T4_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("N1MF5\\I"+str(I)+".csv",delimiter=",")

	#Fit the data and normalise
	Guess = numpy.array([1500, 0.44, 100, -numpy.pi/4])
	popt1, pcov1 = scipy.optimize.curve_fit(DampedRamseySine, N1MF5[1:,0], N1MF5[1:, 1], p0=Guess)
	#popt2 = Guess
	FitTimes5 = numpy.linspace(0, 800, int(1e3))
	FitNumber5 = DampedRamseySine(FitTimes5, popt1[0], popt1[1], popt1[2], popt1[3])
	T5.append(popt1[2])
	T5_err.append(numpy.sqrt(pcov1[2,2]))
	#pyplot.figure(I)
	#pyplot.plot(FitTimes5,FitNumber5)
	#pyplot.scatter(N1MF5[1:,0],N1MF5[1:,1])
	#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,15,'10a',20]:
	#Now get the N = 0, MF = 5 <-> N = 1, MF = 4
	N1MF4 = numpy.genfromtxt("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(DampedRamseySine, N1MF4[1:,0], N1MF4[1:, 1], p0=Guess)
	#popt2 = Guess
	FitTimes4 = numpy.linspace(0, 800, int(1e3))
	FitNumber4 = DampedRamseySine(FitTimes4, popt2[0], popt2[1], popt2[2], popt2[3])
	T4.append(popt2[2])
	T4_err.append(numpy.sqrt(pcov2[2,2]))
	#pyplot.figure(I)
	#pyplot.plot(FitTimes4,FitNumber4)
	#pyplot.scatter(N1MF4[1:,0],N1MF4[1:,1])
	#Print the result

	print("for I = "+str(I)+"Dephasing Times are...")
	print("MF = 4:", int(round(popt2[2])), "+/-", int(round(numpy.sqrt(pcov2[2,2]))), "$\mu$s")


#Plot the result
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,1])
ax2 = fig.add_subplot(Grid[0,1],sharex=ax1)

ax1.plot(FitTimes5, FitNumber5/popt1[0],color=JQC['reddish'],zorder=2)
ax1.scatter(N1MF5[1:,0], N1MF5[1:, 1]/popt1[0], marker='o', color=JQC['red'],zorder=3)

ax2.plot(FitTimes4, FitNumber4/popt2[0],color=JQC['grayblue'],zorder=2)
ax2.scatter(N1MF4[1:,0], N1MF4[1:, 1]/popt2[0], marker='o',color=JQC['blue'],zorder=3)


ax1.set_xlim(0, 120)
ax1.set_ylim(-0.1, 1.1)
ax1.set_xlabel("Time ($\mu$s)")
ax2.set_xlim(0, 700)
ax2.set_ylim(-0.1, 1.1)
#ax2.set_xlabel("Time($\mu$s)")
#ax2.spines['bottom'].set_visible(False)
ax2.tick_params(which='both',labelbottom='off')
#ax2.text(-0.17,1.5,"Molecules in N=0",rotation='vertical',transform=ax2.transAxes,backgroundcolor='white',zorder=10.)
ax2.set_ylabel("Molecules in $N=0$")
ax1.set_ylabel("Molecules in $N=0$")

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

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

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')

curvature =numpy.gradient(grad,Line_data[1,0]-Line_data[0,0])
curve_func = interp.interp1d(Line_data[:,0],curvature,kind='cubic')



coherence = 1./numpy.abs((grad_func(Intensities)*0.04*Intensities+0.5*curve_func(Intensities)*(0.04*Intensities)**2))
ax3.plot(Intensities*1e-3,0.5*coherence*1e6,color=JQC['reddish'],zorder=1.)

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')

#ax3.plot(Line_data[:,0],grad)
curvature =numpy.gradient(grad,Line_data[1,0]-Line_data[0,0])
curve_func = interp.interp1d(Line_data[:,0],curvature,kind='cubic')

coherence = 1./numpy.abs((grad_func(Intensities)*0.04*Intensities+0.5*curve_func(Intensities)*(0.04*Intensities)**2))

ax3.plot(Intensities*1e-3,0.1*coherence*1e6,color=JQC['grayblue'],zorder=1.)

ax3.set_xlabel("Intensity ($\mathrm{kW/cm^2}$)")
ax3.set_ylabel("Coherence Time ($\mu$s)")
ax3.set_xlim(5,40)
ax3.set_ylim(0,400)
ax3.errorbar([10.4,31.1,20.7],T5,yerr=T5_err,fmt='o',color=JQC['red'])
ax3.errorbar([10.4,15.6,10.4,20.7],T4,yerr=T4_err,fmt='o',color=JQC['blue'])

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



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

fig.savefig("dephasing_comparison_zooms_one_axis.pdf")
pyplot.show()
