from __future__ import division

import numpy
from svgpath2mpl import parse_path


import matplotlib.pyplot as pyplot
from  matplotlib.patches import PathPatch
from matplotlib.gridspec import GridSpec
from matplotlib.transforms import Affine2D
import matplotlib.image as mpimg
import matplotlib.figure as figure

MW_photon_path= "m 30.790347,34.759243 c 0.796888,-2.239887 1.593775,-5.041592 2.390663,-7.00764 0.796888,-1.966058 1.593727,-3.007147 2.390615,-2.609091 0.796888,0.398076 1.593775,2.224444 2.390663,4.563404 0.796888,2.33895 1.593776,5.091048 2.390654,6.884154 0.796887,1.793126 1.593775,2.544933 2.390663,1.885806 0.796888,-0.659138 1.593727,-2.706922 2.390663,-5.117516 0.796839,-2.410604 1.593727,-5.080816 2.390615,-6.679975 0.796887,-1.599169 1.593775,-2.052884 2.390653,-1.140402 0.796898,0.912492 1.593776,3.15764 2.390663,5.611638 0.796878,2.453987 1.593766,5.011006 2.390644,6.397457 0.796888,1.386462 3.154777,3.283549 4.952899,-3.206435 0.551123,-1.989176 1.854583,-3.193491 3.066395,-3.584158 1.211814,-0.390668 8.790642,-0.220771 8.790642,-0.220771 M 71.458072,30.624626 l -5.611415,3.239753 M 71.434257,30.466373 L 65.822842,27.22662"
WaveNoToTHz = 1e3 * 1e2 * 3e8 / 1e12

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

def rabi_fall(t,O,N,T0):
	d=0
	w = numpy.sqrt(d**2+O**2)
	phi = w*t
	Rabi_osc = numpy.cos(phi)**2+(d/w)**2*numpy.sin(phi)**2
	decay = N *(1-(t/T0)**2)
	return decay*Rabi_osc
	
def photon(x,x0,extent,end,scale,offset):
	y = scale*numpy.sin((2*numpy.pi*(x-x0))/(extent*0.55))+offset
	y[numpy.where(x>end)]=offset
	return y
	
	
pyplot.style.use('Wide_Origin_style.mplstyle')

colour_stokes = '#ff7f0e'
colour_pump = '#1f77b4'

Dat_theory = numpy.genfromtxt("StirapThe_PG.csv",delimiter=',')
Dat_EXP = numpy.genfromtxt("StirapExp.csv",delimiter=',')

Grid = GridSpec(3,3,width_ratios=[1,1,1],height_ratios=[1.2,1.2,3])

w, h = figure.figaspect(1.)
fig=pyplot.figure(figsize=(w,h), dpi=100)





ax11 = fig.add_subplot(Grid[2,2-2])
ax11.errorbar(Dat_EXP[:,0]*1e3-20,Dat_EXP[:,1]*1e-3,yerr=Dat_EXP[:,2]*1e-3,fmt='o',color=JQC['blue'],zorder=2, ms=4)
ax11.plot(Dat_theory[:,0]-20,Dat_theory[:,4]*1e-3,color=JQC['grayblue'],zorder=1)
ax11.plot(Dat_theory[:,0]-20,Dat_theory[:,5]*1e-3,color=JQC['reddish'],zorder=1)
ax11.spines['right'].set_visible(False)



ax12 = fig.add_subplot(Grid[2,3-2],sharey=ax11)
ax12.set_xticks([])
ax12.spines['right'].set_visible(False)
ax12.spines['left'].set_visible(False)
ax12.tick_params(axis='y',which='both',labelleft='off',left='off')

ax13 = fig.add_subplot(Grid[2,4-2],sharey=ax11)

ax13.errorbar(Dat_EXP[:,0]*1e3-60,Dat_EXP[:,1]*1e-3,yerr=Dat_EXP[:,2]*1e-3,fmt='o',color=JQC['blue'],zorder=2, ms=4)
ax13.plot(Dat_theory[:,0]-60,Dat_theory[:,4]*1e-3,color=JQC['grayblue'],zorder=1)
ax13.plot(Dat_theory[:,0]-60,Dat_theory[:,5]*1e-3,color=JQC['reddish'],zorder=1)
ax13.spines['left'].set_visible(False)
ax13.tick_params(axis='y',which='both',labelleft='off',left='off')








#Broken axis ticks

d = .02  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax11.transAxes, color='k', clip_on=False,zorder=11)
ax11.plot((1-d, 1+d), (1-d, 1+d), **kwargs)        # top-right diagonal
ax11.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # bottom-right diagonal

kwargs.update(transform=ax12.transAxes)  # switch to the bottom axes
ax12.plot((-d, +d), (1 - d, 1 + d), **kwargs)  # bottom-left diagonal
ax12.plot((-d, +d), (- d,+ d), **kwargs)  # bottom-right diagonal


d = .02  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax12.transAxes, color='k', clip_on=False,zorder=11)
ax12.plot((1-d, 1+d), (1-d, 1+d), **kwargs)        # top-right diagonal
ax12.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # bottom-right diagonal

kwargs.update(transform=ax13.transAxes)  # switch to the bottom axes
ax13.plot((-d, +d), (1 - d, 1 + d), **kwargs)  # bottom-left diagonal
ax13.plot((-d, +d), (- d,+ d), **kwargs)  # bottom-right diagonal

ax11.set_ylabel(r"Molecule No. ($\times 10^3$)")

ax01 = fig.add_subplot(Grid[1,2-2],sharex=ax11)
ax02 = fig.add_subplot(Grid[1,3-2],sharex=ax12,sharey=ax01)
ax03 = fig.add_subplot(Grid[1,4-2],sharex=ax13,sharey=ax01)

ax01.plot(Dat_theory[:,0]-20,Dat_theory[:,6]/numpy.amax(Dat_theory[:,6]),zorder=5,color=colour_pump)
ax01.plot(Dat_theory[:,0]-20,Dat_theory[:,7]/numpy.amax(Dat_theory[:,7]),color=colour_stokes)

ax02.plot([0,100],[0,0],zorder=5,color=colour_pump)
ax02.plot([0,100],[0,0],color=colour_stokes)

ax03.plot(Dat_theory[:,0]-60,Dat_theory[:,6]/numpy.amax(Dat_theory[:,6]),zorder=5,color=colour_pump)
ax03.plot(Dat_theory[:,0]-60,Dat_theory[:,7]/numpy.amax(Dat_theory[:,7]),color=colour_stokes)


################################DIPOLE TRAP############################################
ax21 = fig.add_subplot(Grid[0,2-2],sharex=ax11)
ax22 = fig.add_subplot(Grid[0,3-2],sharex=ax12,sharey=ax21)
ax23 = fig.add_subplot(Grid[0,4-2],sharex=ax13,sharey=ax21)

#Broken axis ticks

d = .02  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax21.transAxes, color='k', clip_on=False,zorder=11)
ax21.plot((1-d, 1+d), (1-3*d, 1+3*d), **kwargs)        # top-right diagonal
#ax01.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # bottom-right diagonal

kwargs.update(transform=ax22.transAxes)  # switch to the bottom axes
ax22.plot((-d, +d), (1 - 3*d, 1 + 3*d), **kwargs)  # bottom-left diagonal
#ax02.plot((-d, +d), (- d,+ d), **kwargs)  # bottom-right diagonal


d = .02  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax22.transAxes, color='k', clip_on=False,zorder=11)
ax22.plot((1-d, 1+d), (1-3*d, 1+3*d), **kwargs)        # top-right diagonal
#ax02.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # bottom-right diagonal

kwargs.update(transform=ax23.transAxes)  # switch to the bottom axes
ax23.plot((-d, +d), (1 - 3*d, 1 + 3*d), **kwargs)  # bottom-left diagonal
#ax03.plot((-d, +d), (- d,+ d), **kwargs)  # bottom-right diagonal

d = .02  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax21.transAxes, color='k', clip_on=False,zorder=11)
ax21.plot((1-d, 1+d), (-3*d, +3*d), **kwargs)        # top-right diagonal
#ax01.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # bottom-right diagonal

kwargs.update(transform=ax22.transAxes)  # switch to the bottom axes
ax22.plot((-d, +d), ( - 3*d,  + 3*d), **kwargs)  # bottom-left diagonal
#ax02.plot((-d, +d), (- d,+ d), **kwargs)  # bottom-right diagonal


d = .02  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax22.transAxes, color='k', clip_on=False,zorder=11)
ax22.plot((1-d, 1+d), (-3*d, +3*d), **kwargs)        # top-right diagonal
#ax02.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # bottom-right diagonal

kwargs.update(transform=ax23.transAxes)  # switch to the bottom axes
ax23.plot((-d, +d), ( - 3*d,  + 3*d), **kwargs)  # bottom-left diagonal
#ax03.plot((-d, +d), (- d,+ d), **kwargs)  # bottom-right diagonal

ax22.set_xticks([])
ax21.spines['right'].set_visible(False)

ax22.spines['right'].set_visible(False)
ax22.spines['left'].set_visible(False)
ax22.tick_params(axis='y',which='both',labelleft='off',left='off')

ax23.spines['left'].set_visible(False)
ax23.tick_params(axis='y',which='both',labelleft='off',left='off')

ax21.tick_params(
    axis='x',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off') # labels along the bottom edge are off

ax23.tick_params(
    axis='x',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off') # labels along the bottom edge are off

ax21.set_ylim(-.1,1.2)
ax21.set_ylabel("$I/I_{\mathrm{max}}$")

ax21.plot(Dat_theory[298:,0]-20,Dat_theory[298:,12]/numpy.amax(Dat_theory[:,12]),zorder=5,color=JQC['purple'], linestyle='dashed')

ax22.plot([0,100],[1,1],zorder=5,color=JQC['purple'], linestyle='dashed')

ax23.plot(Dat_theory[:418,0]-60,Dat_theory[:418,12]/numpy.amax(Dat_theory[:,12]),zorder=5,color=JQC['purple'], linestyle='dashed')



ax21.plot(Dat_theory[:,0]-20,Dat_theory[:,11]/numpy.amax(Dat_theory[:,11]),zorder=100,color=JQC['purple'])

ax22.plot([0,100],[0,0],zorder=100,color=JQC['purple'])

ax23.plot(Dat_theory[:,0]-60,Dat_theory[:,11]/numpy.amax(Dat_theory[:,11]),zorder=100,color=JQC['purple'])



ax21.text(0,0.1,"Dipole \n Trap",rotation=00,color=JQC['purple'],fontsize=11)


#############################################################################


#Broken axis ticks

d = .02  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax01.transAxes, color='k', clip_on=False,zorder=11)
ax01.plot((1-d, 1+d), (1-3*d, 1+3*d), **kwargs)        # top-right diagonal
#ax01.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # bottom-right diagonal

kwargs.update(transform=ax02.transAxes)  # switch to the bottom axes
ax02.plot((-d, +d), (1 - 3*d, 1 + 3*d), **kwargs)  # bottom-left diagonal
#ax02.plot((-d, +d), (- d,+ d), **kwargs)  # bottom-right diagonal


d = .02  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax02.transAxes, color='k', clip_on=False,zorder=11)
ax02.plot((1-d, 1+d), (1-3*d, 1+3*d), **kwargs)        # top-right diagonal
#ax02.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # bottom-right diagonal

kwargs.update(transform=ax03.transAxes)  # switch to the bottom axes
ax03.plot((-d, +d), (1 - 3*d, 1 + 3*d), **kwargs)  # bottom-left diagonal
#ax03.plot((-d, +d), (- d,+ d), **kwargs)  # bottom-right diagonal

d = .02  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax01.transAxes, color='k', clip_on=False,zorder=11)
ax01.plot((1-d, 1+d), (-3*d, +3*d), **kwargs)        # top-right diagonal
#ax01.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # bottom-right diagonal

kwargs.update(transform=ax02.transAxes)  # switch to the bottom axes
ax02.plot((-d, +d), ( - 3*d,  + 3*d), **kwargs)  # bottom-left diagonal
#ax02.plot((-d, +d), (- d,+ d), **kwargs)  # bottom-right diagonal


d = .02  # how big to make the diagonal lines in axes coordinates
# arguments to pass to plot, just so we don't keep repeating them
kwargs = dict(transform=ax02.transAxes, color='k', clip_on=False,zorder=11)
ax02.plot((1-d, 1+d), (-3*d, +3*d), **kwargs)        # top-right diagonal
#ax02.plot((1 - d, 1 + d), (-d, +d), **kwargs)  # bottom-right diagonal

kwargs.update(transform=ax03.transAxes)  # switch to the bottom axes
ax03.plot((-d, +d), ( - 3*d,  + 3*d), **kwargs)  # bottom-left diagonal
#ax03.plot((-d, +d), (- d,+ d), **kwargs)  # bottom-right diagonal






ax02.set_xticks([])
ax01.spines['right'].set_visible(False)

ax02.spines['right'].set_visible(False)
ax02.spines['left'].set_visible(False)
ax02.tick_params(axis='y',which='both',labelleft='off',left='off')

ax03.spines['left'].set_visible(False)
ax03.tick_params(axis='y',which='both',labelleft='off',left='off')

ax01.text(10,0.64,"Stokes",rotation=-50,color=colour_stokes,fontsize=11)
ax01.text(1,0.64,"Pump",rotation=50,color=colour_pump,fontsize=11)


ax11.text(-2.5,3.9,"$|F\\rangle$",color=JQC['blue'],fontsize=14,clip_on=False)
ax11.text(-2.5,0.3,"$|G\\rangle$",color=JQC['red'],fontsize=14,clip_on=False)


ax11.set_xticks([0,20,40])
ax11.set_xlim(-4,23.1)
ax12.set_xlim(55,56)
ax13.set_xticks([0,20,40])
ax13.set_xlim(-3.1,24)

ax01.set_yticks([0,1,2])
ax01.set_ylim(-.1,1.2)
ax11.set_yticks([0,2,4,6])
ax11.set_ylim(-0.2,5)

ax01.tick_params(
    axis='x',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off') # labels along the bottom edge are off

ax03.tick_params(
    axis='x',          # changes apply to the x-axis
    which='both',      # both major and minor ticks are affected
    bottom='off',      # ticks along the bottom edge are off
    top='off',         # ticks along the top edge are off
    labelbottom='off') # labels along the bottom edge are off



ax12.text(0.15,-0.25,"Time ($\mu$s)",transform=ax12.transAxes)
#ax02.text(0.1,-0.4,"Time ($\mu$s)",transform=ax02.transAxes)

###drawing microwave interactions ###
VertDisplace1 = -1.2
VertDisplace2 = -0.5
Hshift = -0.05
ax12.hlines(2.1+VertDisplace2,55.2+Hshift,55.9+Hshift)
ax12.hlines(5.15+VertDisplace1,55.2+Hshift,55.9+Hshift)

ax12.text(55.25+Hshift,4.75+VertDisplace1,r"$N=1$",fontsize=12,clip_on=False)
ax12.text(55.25+Hshift,2.15+VertDisplace2,r"$N=0$",fontsize=12,clip_on=False)
ax12.text(55.1,0.55,r"Variable Hold",fontsize=12)
ax12.text(54.95+Hshift,3.85+VertDisplace1,"MW",fontsize=12)
ax12.annotate(s='',xy=(55,0.2),xytext=(56,0.2),arrowprops=dict(arrowstyle="<->",color='k',lw=2.))


ax12.annotate(s='',xy=(55.7+Hshift,2.1+VertDisplace2),xytext=(55.7+Hshift,5.15+VertDisplace1),arrowprops=dict(arrowstyle="<-",color='k',lw=2.))
ax12.annotate(s='',xy=(55.8+Hshift,5.1+VertDisplace1),xytext=(55.8+Hshift,2.1+VertDisplace2),arrowprops=dict(arrowstyle="<-",color='k',lw=2.))




##

#ax_potentials.text(0.02,1.01,r"$\times 10^3$",clip_on=False,transform=ax_potentials.transAxes)
ax01.set_ylabel("$\Omega/\Omega_\mathrm{max}$")





Interp = 'none'

#Add images of molecule wavefunctions
ax_N0MN0 = pyplot.axes([0.415, 0.12, 0.25, 0.25])
IMG = mpimg.imread('Fig_N0MN0.png')
ax_N0MN0.imshow(IMG, interpolation = Interp )
ax_N0MN0.axis('off')

#Add images of molecule wavefunctions
ax_N1MNp1 = pyplot.axes([0.365, 0.37, 0.25, 0.25])
IMG = mpimg.imread('Fig_N1MN1.png')
ax_N1MNp1.imshow(IMG, interpolation = Interp )
ax_N1MNp1.axis('off')

#Add images of molecule wavefunctions
ax_N1MN0 = pyplot.axes([0.415, 0.37, 0.25, 0.25])
IMG = mpimg.imread('Fig_N1MN0.png')
ax_N1MN0.imshow(IMG, interpolation = Interp )
ax_N1MN0.axis('off')

#Add images of molecule wavefunctions
ax_N1MNm1 = pyplot.axes([0.465, 0.37, 0.25, 0.25])
IMG = mpimg.imread('Fig_N1MN-1.png')
ax_N1MNm1.imshow(IMG, interpolation = Interp )
ax_N1MNm1.axis('off')
#adjust layout
fig.subplots_adjust(wspace=0.05,hspace=0.15,left=0.12,bottom=0.13,top=0.94,right=0.97)

axArrow = pyplot.axes([0.345, 0.3, 0.5, 0.5])
collection = PathPatch(path = parse_path(MW_photon_path).transformed(Affine2D().translate(-30,-25).scale(1/50,1/20)),facecolor='none',lw=2.,ec='k',transform=Affine2D().scale(1/2,-1/1.5).translate(55.2+Hshift,4.2+VertDisplace1)+ax12.transData,joinstyle='round',capstyle='round')
axArrow.add_artist(collection)
axArrow.axis('off')
#ax11.add_artist(collection)

#done -save and show
#pyplot.savefig("SimonCCMI.png", dpi=600)
pyplot.show()
pyplot.close()