from matplotlib import pyplot,gridspec
from scipy.optimize import curve_fit
import numpy,os
from jqc import jqc_plot

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

jqc_plot.plot_style("normal")

fig = pyplot.figure("Rabi")

colour = jqc_plot.colours
grid =gridspec.GridSpec(2,2,height_ratios=[1,1.5])

ax1 = fig.add_subplot(grid[0,0])

data = numpy.genfromtxt(cwd+"\\data\\MW1.csv",delimiter=',')*1e-3

times = numpy.linspace(-0.05,1,500)


fitfn = lambda x,w,A: A*(numpy.cos(numpy.pi*x*w))**2
x = [8.500,4]
curve1,cov1 = curve_fit(fitfn,data[:,0],data[:,1],p0=x)
err1 = numpy.sqrt(numpy.diag(cov1))

print(curve1,err1)

ax1.scatter(data[:,0],data[:,1],color=jqc_plot.colours['red'],zorder=1.2)
ax1.plot(times,fitfn(times,*curve1),color=jqc_plot.colours['reddish'],zorder=1.1)

ax1.set_xlim(0,0.36)

ax1.set_ylabel("$N_\\mathrm{mol}$ ($10^3$)")
ax1.set_xlabel("$T_\\mathrm{Pulse}$ (ms)")

ax2 = fig.add_subplot(grid[0,1],sharey=ax1)

data = numpy.genfromtxt(cwd+"\\data\\MW2.csv",delimiter=',')*1e-3

x = [8.500,4]
curve2,cov2 = curve_fit(fitfn,data[:,0],data[:,1],p0=x)
err2 = numpy.sqrt(numpy.diag(cov2))

print(curve2,err2)

ax2.scatter(data[:,0],data[:,1],color=jqc_plot.colours['blue'],zorder=1.2)
ax2.plot(times,fitfn(times,*curve2),color=jqc_plot.colours['grayblue'],zorder=1.1)

ax2.set_xlim(0,0.36)
ax2.set_xlabel("$T_\\mathrm{Pulse}$ (ms)")
ax2.set_ylim(0,3.5)
ax2.tick_params(labelleft=False)

ax_Raman = fig.add_subplot(grid[1,:])

data = numpy.genfromtxt(cwd+"\\data\\N0_rabi.csv",delimiter=',',skip_header=1)

data = data*1e-3 #convert time to ms and number to 1000s

fitfn = lambda x,w,t,A: A*(1-(x/t)**2)*(numpy.cos(numpy.pi*x*w))**2
guess =[0.6,4.000,3.250]
curve,cov = curve_fit(fitfn,data[:,0],data[:,1],p0=guess,
                        bounds=([0.2,.500,0],[5,5.000,5.000]))
err = numpy.sqrt(numpy.diag(cov))
print(curve[0],err[0])

times = numpy.linspace(-0.5,1.1*numpy.amax(data[:,0]),500)
model = fitfn(times,*curve)

model_err = numpy.zeros(times.shape)
for i,x in enumerate(err):
    pars = curve.copy()
    pars[i]+= x
    model_err_x = fitfn(times,*pars)-model
    model_err = model_err + model_err_x**2
model_err = numpy.sqrt(model_err)/curve[2]

ax_Raman.scatter(data[:,0],data[:,1],zorder=1.6,color=jqc_plot.colours['purple'])
ax_Raman.plot(times,model,zorder=1.2,color=jqc_plot.colours['purpley'])


ax_Raman.fill_between(times,model-model_err,model+model_err,
            facecolor=jqc_plot.colours['purpley'],alpha=0.5,edgecolor='w')

fitfn = lambda x,A,t: A*(1-(x/t)**2)
data = numpy.genfromtxt(cwd+"\\data\\loss.csv",delimiter=',',skip_header=1)*1e-3
curve2,cov = curve_fit(fitfn,data[:,0],data[:,1])
curve2[0] = curve[2]
err2 = numpy.sqrt(numpy.diag(cov))
#ax_Raman.scatter(data[:,0],data[:,1])
ax_Raman.plot(times,fitfn(times,*curve2),ls='--',color='k')
#ax_Raman.fill_between(times,fitfn(times,curve[2],curve2[1]-err2[1]),fitfn(times,curve[2],curve2[1]+err2[1]))

ax_Raman.set_xlim(-0.05,4.1)
ax_Raman.set_ylim(0,4)
ax_Raman.set_xlabel("Pulse Time, $T_\\mathrm{Pulse}$ (ms)")
ax_Raman.set_ylabel("$N_\\mathrm{mol}$ ($10^3$)")

ax1.text(0.74,0.8,"(b)",transform=ax1.transAxes,fontsize=20)
ax2.text(0.75,0.8,"(c)",transform=ax2.transAxes,fontsize=20)
ax_Raman.text(0.87,0.87,"(d)",transform=ax_Raman.transAxes,fontsize=20)


fig.tight_layout()
fig.subplots_adjust(wspace=0.13,hspace=0.4,right=0.95,left=0.5)


pyplot.savefig(cwd+"\\Rabi_raman.pdf")

pyplot.show()
