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.csv",delimiter=',')
Dat_EXP = numpy.genfromtxt("StirapExp.csv",delimiter=',')

Grid = GridSpec(2,1,width_ratios=[1],height_ratios=[1.2,3])
w, h = figure.figaspect(1.)
fig=pyplot.figure(figsize=(w,h), dpi=100)


### Drawing the potential curves and annotating###
ax_potentials = fig.add_subplot(Grid[:])


dat_pot=numpy.genfromtxt("potentials.csv",delimiter=',')

#draw curves
for i in range(1,len(dat_pot[0,:])-1):
	ax_potentials.plot(dat_pot[:,0]*0.1,dat_pot[:,i]*1e-3*WaveNoToTHz,color=JQC['grayblue'] )

#set axes
ax_potentials.set_ylim(-4.5*WaveNoToTHz,7*WaveNoToTHz)
ax_potentials.set_xlim(.30,1.10)
ax_potentials.set_xlabel(r"Atomic Separation (nm)")
ax_potentials.set_ylabel(r"Energy/$h$ (THz)",labelpad = -5)

#annotate
ax_potentials.hlines(0*WaveNoToTHz,0.5151,1.1,linestyle='-',color='k',zorder=2)
ax_potentials.hlines(6.422*WaveNoToTHz,.358,.56825,linestyle='-',color='k',zorder=2)
ax_potentials.hlines((6.422-10.20408)*WaveNoToTHz,.42,.4588,linestyle='-',color='k',zorder=2)
ax_potentials.text(.5,-4*WaveNoToTHz,r"$|G\rangle=|X^1\Sigma^+\;v''=0,N''=0\rangle$",fontsize=12)
ax_potentials.text(0.6,0.22*WaveNoToTHz,r"$|F\rangle$",fontsize=12)
ax_potentials.text(.59,6.25*WaveNoToTHz,r"$|E\rangle=|^3\Pi_1 \; v'=29,J'=1\rangle$",fontsize=12)

ax_potentials.annotate(s='',xy=(.53,0.0*WaveNoToTHz),xytext=(.53,6.422*WaveNoToTHz),arrowprops=dict(arrowstyle="<-",color=colour_pump,lw=2.))
ax_potentials.annotate(s='',xy=(.425,6.422*WaveNoToTHz),xytext=(.425,(6.422-10.20408)*WaveNoToTHz),arrowprops=dict(arrowstyle="<-",color=colour_stokes,lw=2.))

ax_potentials.text(.49,4*WaveNoToTHz,"Pump 1557 nm",color=colour_pump,fontsize=12,rotation='vertical')
ax_potentials.text(.38,0.5*WaveNoToTHz,"Stokes 977 nm",color=colour_stokes,fontsize=12,rotation='vertical')

ax_potentials.text(.63,-2*WaveNoToTHz,r"$X^1\Sigma^+$",fontsize=12)
ax_potentials.text(.54,-.88*WaveNoToTHz,r"$a^3\Sigma^+$",fontsize=12)
#ax_potentials.text(.425,7.500*WaveNoToTHz,r"$b^3\Pi$",fontsize=12)
#ax_potentials.text(.69,7.500*WaveNoToTHz,r"$A^1\Sigma^+$",fontsize=12)

ax_potentials.text(.79,-.76*WaveNoToTHz,r"Rb: $5S_\frac{1}{2}+$ Cs: $6S_\frac{1}{2}$",fontsize=12)


'''
#Add inset for Feshbach states
axPotInset = pyplot.axes([0.53,0.59,0.4,0.235])
figInset = pyplot.gca()
figInset.axhline(197.1, color='black', linestyle='dashed')
Data = numpy.genfromtxt("BoundStates.csv",delimiter=",")
figInset.plot(Data[:,0], Data[:,1], color=JQC['red'], linewidth=3, alpha=0.5)
figInset.plot(Data[:,0], Data[:,3], color='k', linewidth=3)
figInset.annotate("$\\left|-2(1,3)d(0,3) \\right\\rangle$",xy=(0.18, 0.17), xytext=(0.4, 0.1),xycoords="axes fraction", size=9,  
            arrowprops=dict(arrowstyle="->", color='black', linewidth=1))

figInset.annotate("$\\left|-6(2,4)d(2,4) \\right\\rangle$",xy=(0.22, 0.75), xytext=(0.4, 0.7),xycoords="axes fraction", size=9,  
            arrowprops=dict(arrowstyle="->", color='black', linewidth=1))
figInset.set_xlim(177, 200)
figInset.set_ylim(-4, 0, 1)
figInset.set_yticks([-4, -2, 0])
figInset.set_xticks([180, 190, 200])
pyplot.setp(axPotInset.get_xticklabels(), fontsize=12)
pyplot.setp(axPotInset.get_yticklabels(), fontsize=12)
figInset.set_xlabel("B (G)", size=12)
figInset.set_ylabel("Energy/$h$ (MHz)", size=12,labelpad =-2)
figInset.plot([180.35, 181.6], [-3.3, -1.5], linestyle='none', marker = 'o', ms=8, color='black')
'''

#adjust layout
fig.subplots_adjust(wspace=0.05,hspace=0.15,left=0.18,bottom=0.13,top=0.96,right=0.97)


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