from __future__ import division
import numpy
import matplotlib.pyplot as pyplot
import matplotlib.gridspec as gridspec
import matplotlib.patheffects as PathEffects
from matplotlib.patches import ConnectionPatch,FancyBboxPatch,Rectangle
from matplotlib.path import Path
from mpl_toolkits.axes_grid1.inset_locator import BboxPatch, BboxConnector,BboxConnectorPatch
from matplotlib.transforms import TransformedBbox
from scipy.optimize import curve_fit
from matplotlib import rcParams
from scipy.stats import chisquare as chi_squared
rcParams['font.sans-serif'] = ['Arial']
rcParams['font.family'] = 'sans-serif'
import sys
eps = sys.float_info.epsilon

pi = numpy.pi

Grid = gridspec.GridSpec(2,2,width_ratios=[1,1],height_ratios=[1,2])

def Ramsey_Frequency_truppe(detuning,Beta,T_pulse,T_evolution,N_mol):
	''' The Imperial Lineshape function. Makes the assumption that the pulse is a pi/2 pulse for all frequencies that are scanned - will only be valid near to the centre.'''
	detuning = detuning *2*pi
	X = numpy.sqrt(numpy.pi**2.+4.*detuning**2.*T_pulse**2.)
	Prefactor = 4.*numpy.pi**2.*numpy.sin(X/4)**2./X**4.
	Brackets = X*numpy.cos(X/4.)*numpy.cos(0.5*(detuning*T_evolution+Beta)) - 2.*detuning*T_pulse*numpy.sin(X/4.)*numpy.sin(0.5*(detuning*T_evolution+Beta))
	return (1.-Prefactor*(Brackets**2.))*N_mol


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 Falling_ramsey(t,N0,f,f0,Omega,T_pulse,tau):
	'''Ramsey linshape modulated by a (1-X^2) function to match the falling we see ''' 
	Ramsey = Ramsey_Frequency(f-f0,Omega,T_pulse,t,1.0,0)
	return N0*(1.-(t/tau)**2)*Ramsey


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), \
       'sand_alpha'  :(244./255.0, 234./255.0, 168./255.0,0.6), \
       'grayblue'  :(145./255.0, 148./255.0, 182./255.0), \
       'green'   :(45.0/255.0, 159.0/255.0, 60.0/255.0)}

#Import style file
pyplot.style.use(r'Wide_Origin_style.mplstyle')
#fig = pyplot.figure()
#Import Data Csv
Dat = numpy.genfromtxt(r"Data\Ramsey_spectroscopy_refit.csv",delimiter=',')

Det_rng = numpy.linspace(-500e3,500e3,1e6)

#Curve fitting
inc = 30
Centres =[]
err_centre=[]
N = []
err_N =[]
Dat_plotting = Dat.copy()

'''Normalise data to maximum in 30 shot batches '''
for i in range(0,len(Dat[:,0]),inc):
	if i + inc > len(Dat[:,0]):
		#Curve,err = curve_fit(lambda x,d,N: Ramsey_Frequency((x-d),2*pi*5.24e3,47.8e-6,600e-6,N),(Dat[i:-1,0])*1e6,Dat[i:-1,1],p0=[980.386e6,2000.],bounds=([980.383e6,1100.],[980.387e6,2500.]))
		Dat_plotting[i:,1] = Dat[i:,1]/numpy.amax(Dat[i:,1])#Curve[1]
	else:
		#Curve,err = curve_fit(lambda x,d,N: Ramsey_Frequency((x-d),2*pi*5.24e3,47.8e-6,600e-6,N),(Dat[i:i+inc,0])*1e6,Dat[i:i+inc,1],p0=[980.386e6,2000.],bounds=([980.383e6,1100.],[980.387e6,2500.]))
		Dat_plotting[i:i+inc,1] = Dat[i:i+inc,1]/numpy.amax(Dat[i:i+inc,1])#Curve[1]
	#print(Curve,numpy.sqrt(numpy.diag(err)))
	#err = numpy.sqrt(numpy.diag(err))
	#Centres.append(Curve[0])
	#err_centre.append(err[0])


#Fit the "classic" Ramsey lineshape that we see by scanning the detuning. Only free parameter is f0
Centre,err = curve_fit(lambda x,d: Ramsey_Frequency((x-d),2*pi*5.24e3,47.8e-6,600e-6,1.),(Dat[:,0])*1e6,Dat_plotting[:,1],p0=[980.386e6],bounds=([980.383e6],[980.387e6]))

print("Transition Centered at: (", Centre[0], "+/-", numpy.sqrt(err[0,0]), ")Hz")
print("chi squared = ",chi_squared(Dat_plotting[:,1],Ramsey_Frequency((Dat_plotting[:,0]*1e6-Centre[0]),2*pi*5.24e3,47.8e-6,600e-6,1.),1)[0])


#Plotting

#Model plotting range
Det_rng = numpy.linspace(-500e3,500e3,1e6)

#Plot Data (large subplot)

fig=pyplot.figure(2)
ax = fig.add_subplot(Grid[1,0])

Curve_fitted = Ramsey_Frequency(Det_rng,2*pi*5.24e3,47.8e-6,600e-6,1.)
#err_curve = numpy.sqrt((Curve_fitted-Ramsey_Frequency(Det_rng+err[0],2*pi*5.24e3,47.8e-6,600e-6,Curve[1]))**2+(Curve_fitted-Ramsey_Frequency(Det_rng+err[0],2*pi*5.24e3,47.8e-6,600e-6,Curve[1]+err[1]))**2)



ax.plot(1e-3*(Det_rng),Curve_fitted,color=JQC['grayblue'],zorder=1)
ax.scatter(1e-3*(Dat_plotting[:,0]*1e6-Centre[0]),Dat_plotting[:,1],marker='o',color=JQC['blue'],zorder=2)

ax.set_xlabel("Detuning (kHz)")
ax.set_ylim(0,1.05*numpy.amax(Dat_plotting[:,1]))
ax.set_xticks([-20,-15,-10,-5,0,5,10,15,20])
ax.set_xlim(-18.5,18.5)
ax2=fig.add_subplot(Grid[0,0])

ax2.plot(1e-3*(Det_rng),Curve_fitted,color=JQC['grayblue'],zorder=1)
ax2.scatter(1e-3*(Dat_plotting[:,0]*1e6-Centre[0]),Dat_plotting[:,1],marker='o',color=JQC['blue'],zorder=2)

#ax2.set_xlabel("Detuning (kHz)")
#ax2.set_ylabel(r"Normalised Number of Molecules in $|0\rangle$")
ax2.set_yticks([0,0.50,1.0])
ax2.set_xticks([-3,-2,-1,0,1,2,3])
ax2.set_ylim(0,1.0)
ax2.set_xlim(-3.5,3.5)
highlight_1 = Rectangle((-3.5,0),7,1.05*numpy.amax(Dat_plotting[:,1]),fill=True,fc=JQC['sand'],alpha=0.6,transform=ax.transData,zorder = 1.0)
ax.add_patch(highlight_1)


ax_main_bbox1 = TransformedBbox(highlight_1.get_bbox(),ax.transData)
#draw a polygon to connect the highlight to the zoomed axis.
p1 = BboxConnectorPatch(ax_main_bbox1,ax2.bbox, loc1a=2, loc2a=3,loc1b=1, loc2b=4, ec='none',zorder = 1.0,lw=1.5)
p1.set_fill(True)
p1.set_facecolor(JQC['sand_alpha'])
#allow drawing outside of each axis limits
p1.set_clip_on(False)
#draw polygon
ax.add_patch(p1)
ax.set_ylabel(r"Normalised Number of Molecules in $|0\rangle$",y=0.8)

ax1=fig.add_subplot(Grid[:,1])
Dat_coherence = numpy.genfromtxt(r"Data\Coherence_MF6.csv",delimiter=',')
Dat_falling = numpy.genfromtxt(r"Data\Falling_no_MW.csv",delimiter=',')

ax1.set_xlabel("Evolution Time, $T$ (ms)")
ax1.set_ylabel(r"Normalised Number of Molecules in $|0\rangle$")

times = numpy.linspace(0,2.5,1500)

#Fit our data which has no MW radiation on to get an estimate of the timescale for the falling out of the trap. 

falling_ref,err_ref = curve_fit(lambda t,N0,Tau: N0*(1.-(t/Tau)**2),Dat_falling[:-4,0]*1e-6,Dat_falling[:-4,1],p0=[3000,2.0e-3],bounds=([2000,100e-6],[5000,5e-3]))
err_ref = numpy.sqrt(numpy.diag(err_ref))

print("Tau = ",falling_ref[1]*1e3,"+/-",err_ref[1]*1e3,"ms")

times_err = numpy.abs((1.-(times*1e-3/(falling_ref[1]+err_ref[1]))**2)-(1.-(times*1e-3/(falling_ref[1]))**2))

pyplot.plot(times,(1-(times*1e-3/falling_ref[1])**2),color='k',ls='--',zorder=0.6)
pyplot.fill_between(times,(1-(times*1e-3/falling_ref[1])**2)-times_err,(1-(times*1e-3/falling_ref[1])**2)+times_err,color=JQC['grayblue'],alpha=0.4,zorder=0.5)

#fitting for subfigure (b) - use the falling ramsey lineshape and fit T0 and f0 

falling,err_fall = curve_fit(lambda t,N0,Tau,f0: Falling_ramsey(t,N0,980.390e6,f0,2*pi/(4*19e-6),19e-6,Tau),Dat_coherence[:,0]*1e-6,Dat_coherence[:,1],p0=[3000,2.5e-3,980.3856e6],bounds=([2000,0.1e-3,980.3852e6],[5000,5.0e-3,980.3857e6]))
err_fall = numpy.sqrt(numpy.diag(err_fall))

print("Falling centre = ",falling[2],"+/-",err_fall[2],"Hz")
print("falling Tau = ",falling[1]*1e3,"+/-",1e3*err_fall[1],"ms")
print("Falling detuning = ",980.390e6-falling[2],"+/-",err_fall[2],"Hz")

print("chi squared = ",chi_squared(Dat_coherence[:,1]/falling[0],Falling_ramsey(Dat_coherence[:,0]*1e-6,1.,980.390e6,falling[2],2*pi/(4*19e-6),19e-6,falling[1]),3)[0])


#print(falling[2],err_fall[2])

ax1.plot(times,Falling_ramsey(times*1e-3,1.0,980.390e6,falling[2],2*pi/(4*19e-6),19e-6,falling[1]),zorder=1.,color=JQC['grayblue'])

ax1.scatter(Dat_coherence[:,0]*1e-3,Dat_coherence[:,1]/falling[0],marker='o',color=JQC['blue'],zorder=2.)

#change rasterisation so that the shaded regions don't mess up. 
ax1.set_rasterization_zorder(1.)
ax2.set_rasterization_zorder(1.)
ax.set_rasterization_zorder(1.)

ax1.set_ylim(0,1.05)
ax1.set_xlim(0,2.31)

ax2.text(0.015,0.77,"(a)",fontsize=20,transform=ax2.transAxes)
ax1.text(0.89,0.92,"(b)",fontsize=20,transform=ax1.transAxes)

fig.subplots_adjust(left=0.1,right=0.97,bottom=0.14,top=0.95)
#pyplot.savefig("Ramsey_frequency_main.eps",rasterized=True,dpi=500)
#pyplot.savefig("Ramsey_frequency_main.pdf",rasterized=True,dpi=500)
pyplot.show()
