from __future__ import division
import numpy,os
import matplotlib.pyplot as pyplot
from scipy.integrate import ode
from scipy.signal import argrelextrema
from scipy.optimize import curve_fit
import scipy.constants as constants
import matplotlib.gridspec as gridspec
from jqc import jqc_plot

cwd = os.path.dirname(os.path.abspath(__file__))

pi = numpy.pi
h = constants.h
hbar = h/(2.*numpy.pi)
jqc_plot.plot_style("normal")


Pulse_Time = 97e-6
O_1 = 5.e3
resonance = 980.38298e6
epsilon_0 = 8.854187817e-12

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(P,Transition_element,efficiency):
    Power = efficiency*P
    return (Transition_element/hbar)*numpy.sqrt(Power)


def Optical_Bloch(t,c,par):
    c0,c1,c2 = c
    omega_a,omega_b,Delta_1,Delta_2 = 2*numpy.pi*numpy.array(par)
    #print Delta_2
    f = [-0.5j*omega_a*c1,
         -0.5j*(omega_a*c0+omega_b*c2-2.0*Delta_1*c1),
         -0.5j*(omega_b*c1-2.0*(Delta_1-Delta_2)*c2)]
    return f

def Autler_Townes(D,No,O_2,D_2):
    state1=[]
    state2=[]
    state3=[]

    E1 = []
    E2 = []
    E3 = []
    time_start = 0.0
    Initial_state = (1.,0.0,0.0)

    for d in D:
        Par = [O_1,O_2,d,D_2]
        Bloch_equations = ode(Optical_Bloch).set_integrator('zvode',method='bdf',nsteps=1000)
        Bloch_equations.set_initial_value(Initial_state,time_start).set_f_params(Par)
        dt = 0.1e-6

        while Bloch_equations.successful() and Bloch_equations.t < Pulse_Time:
            Bloch_equations.integrate(Bloch_equations.t+dt)
            total_length =numpy.sqrt(abs(Bloch_equations.y[0])**2+abs(Bloch_equations.y[1])**2+abs(Bloch_equations.y[2])**2)

        state1.append(Bloch_equations.y[0]/total_length)
        state2.append(Bloch_equations.y[1]/total_length)
        state3.append(Bloch_equations.y[2]/total_length)



    state1 = No*(numpy.abs(numpy.array(state1))**2)
    state2 = No*(numpy.abs(numpy.array(state2))**2)
    state3 = No*(numpy.abs(numpy.array(state3))**2)
    return state1

i=0
fig=pyplot.figure()
grid = gridspec.GridSpec(3,2,width_ratios=[1.5,1])
ax0 = fig.add_subplot(grid[0,0])
ax1 = fig.add_subplot(grid[1,0],sharex=ax0)
ax2 = fig.add_subplot(grid[2,0],sharex=ax0)
ax3 = fig.add_subplot(grid[:,1])
Detuning = numpy.linspace(-20e3,20e3,500)
Powers = [3000,1000,300]
i =0
O=[]
Oerr=[]

for P in Powers:
    FP = cwd+r"\\Data\AT-"+str(P)+r"mW.csv"
    Data = numpy.loadtxt(FP,delimiter=',')
    Data[:,0]=(Data[:,0])*1e6
    if i ==0:
        Curve,err = curve_fit(lambda x,F,No,O_2,D_2: Autler_Townes(x-F,No,O_2,D_2),Data[:,0],Data[:,1],p0 = [980.38498e6,2e3,17.3e3,-2.4e3],bounds = ([980.383e6,1e3,100,-5e3],[980.387e6,2.5e3,20e3,-100]))
        err = numpy.sqrt(numpy.diag(err))
        print( Curve,err)
        O.append(Curve[2])
        Oerr.append(err[2])
        ax0.scatter((Data[:,0]-Curve[0])*1e-3,Data[:,1]/1e3,color=JQC['grayblue'],label = str(P*1e-3) +" W",zorder=1.25)
        curve_plot = Autler_Townes(Detuning,*Curve[1:])
        ax0.plot(Detuning*1e-3,curve_plot/1e3,color=JQC['blue'],zorder=1.5)
        ax0.set_ylim(-0.1,2.5)

        ax3.text(2500,Curve[2]*1e-3,"(i)",verticalalignment='center',fontsize=15)

        #ax0.text(0.15,0.1,"$\\Omega_{B} =$"+" {:.1f}({:.0f}) kHz".format(Curve[2]*1e-3,err[2]*1e-2),transform=ax0.transAxes,fontsize=13)
        #ax0.text(0.63,0.1,"$\\Delta_{B} =$"+" {:.0f}({:.0f}) kHz".format(Curve[3]*1e-3,err[3]*1e-3),transform=ax0.transAxes,fontsize=13)
        ax0.tick_params(labelbottom=False)
    elif i ==1:
        Curve,err = curve_fit(lambda x,F,No,O_2,D_2: Autler_Townes(x-F,No,O_2,D_2),Data[:,0],Data[:,1],p0 = [980.38498e6,2e3,10e3,-2.4e3],bounds = ([980.383e6,1e3,100,-5e3],[980.387e6,2.5e3,15e3,-100]))
        err = numpy.sqrt(numpy.diag(err))
        print( Curve,err)
        O.append(Curve[2])
        Oerr.append(err[2])
        ax1.scatter((Data[:,0]-Curve[0])*1e-3,Data[:,1]/1e3,color=JQC['purpley'],label = str(P*1e-3) +" W",zorder=1.25)
        curve_plot = Autler_Townes(Detuning,*Curve[1:])
        ax1.plot(Detuning*1e-3,curve_plot/1e3,color=JQC['purple'],zorder=1.5)
        ax1.set_ylim(-0.1,2.5)
        print(Curve[0])
        ax3.text(1500,Curve[2]*1e-3,"(ii)",verticalalignment='center',fontsize=15)
        #ax1.text(0.15,0.1,"$\\Omega_{B} =$"+" {:.1f}({:.0f}) kHz".format(Curve[2]*1e-3,err[2]*1e-2),transform=ax1.transAxes,fontsize=13)
        #ax1.text(0.63,0.1,"$\\Delta_{B} =$"+" {:.1f}({:.0f}) kHz".format(Curve[3]*1e-3,err[3]*1e-2),transform=ax1.transAxes,fontsize=13)
        ax1.tick_params(labelbottom=False)
    elif i ==2:
        Curve,err = curve_fit(lambda x,F,No,O_2,D_2: Autler_Townes(x-F,No,O_2,D_2),Data[:,0],Data[:,1],p0 = [980.38498e6,2e3,550,-2.4e3],bounds = ([980.383e6,1e3,100,-5e3],[980.387e6,2.5e3,2e3,-100]))
        err = numpy.sqrt(numpy.diag(err))
        print(Curve,err)
        O.append(Curve[2])
        Oerr.append(err[2])
        ax2.scatter((Data[:,0]-Curve[0])*1e-3,Data[:,1]/1e3,color=JQC['reddish'],label = str(P) +" mW",zorder=1.25)
        curve_plot = Autler_Townes(Detuning,*Curve[1:])
        ax2.plot(Detuning*1e-3,curve_plot/1e3,color=JQC['red'],zorder=1.5)
        ax2.set_ylim(-0.1,2.5)


        ax3.text(750,Curve[2]*1e-3,"(iii)",verticalalignment='center',fontsize=15)
        #ax2.text(0.63,0.1,"$\\Delta_{B} =$"+" {:.0f}({:.0f}) kHz".format(Curve[3]*1e-3,err[3]*1e-3),transform=ax2.transAxes,fontsize=13)
        #ax2.text(0.15,0.1,"$\\Omega_{B} =$"+" {:.0f}({:.0f}) kHz".format(Curve[2]*1e-3,err[2]*1e-3),transform=ax2.transAxes,fontsize=13)
    i+=1
    print("Done")
O = numpy.array(O)
Oerr = numpy.array(Oerr)
Powers =[3162.3,1e3,0.3162e3]
ax2.set_xlabel("Detuning (kHz)")
ax1.set_ylabel("Number of Molecules in $|0\\rangle$ ($10^3$)")
ax0.set_xlim(-22,22)

fitfn = lambda x,A: A*numpy.sqrt(x)

curve,err = curve_fit(fitfn,Powers,O,sigma= Oerr,absolute_sigma=True)
P = numpy.linspace(0,5000,100)
ax3.errorbar(Powers,1e-3*O,yerr=1e-3*Oerr,color='k',fmt='o')
ax3.plot(P,1e-3*fitfn(P,*curve),ls='--',color='k')

ax3.set_ylabel("Fitted Rabi Frequency, $\\Omega_{12}}/2\\pi$ (kHz)")
ax3.set_xlabel("Input Power (mW)")
ax3.set_xlim(0,3500)
xmin,xmax=ax3.get_xlim()
ax3.set_ylim(0,1e-3*fitfn(xmax,*curve))

fig.tight_layout()

ax2.text(0.02,0.1,"(a)",transform=ax2.transAxes,fontsize=20)
ax0.text(0.9,0.1,"(i)",transform=ax0.transAxes,fontsize=15,
                        horizontalalignment='center')
ax1.text(0.9,0.1,"(ii)",transform=ax1.transAxes,fontsize=15,
                        horizontalalignment='center')
ax2.text(0.9,0.1,"(iii)",transform=ax2.transAxes,fontsize=15,
                        horizontalalignment='center')


ax3.text(0.75,0.05,"(b)",transform=ax3.transAxes,fontsize=20)


fig.subplots_adjust(wspace =0.4,hspace=0,top=0.95,right=0.95,
                    bottom=0.15,left=0.15)



pyplot.savefig(cwd+"\\AT.pdf")
pyplot.savefig(cwd+"\\AT.png")

pyplot.show()
