import numpy,os,time
from matplotlib import pyplot,gridspec
from jqc import jqc_plot
from scipy import constants,integrate,optimize

h = constants.h

pi = constants.pi

jqc_plot.plot_style("normal")

colours = jqc_plot.colours

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


########################### Set Global Parameters ########################
T0 = 1.5e-6 #Kelvin
w0 = 2*numpy.pi*(181*178*44)**(1/3) #Geometric average of trap frequency, Hz
STIRAPEFFICIENCY = 0.92
m = 220*1.660539e-27 #kg
kB = 1.38065e-23 #J/K

ConstantTemperature = False
N0Fixed =False

def ArbLoss(Input, t, b, Beta):
    N0, T0 = Input
    b = abs(b)
    Beta = abs(Beta)
    #I think this is right for number at least - using definitions from Anderlini thesis
    dN = - b*(N0**Beta)/(Beta**(3/2) * T0**((3/2) * (Beta-1)))
    dT =  b*((Beta - 1) / (2*Beta)) * N0**(Beta-1) * (T0 + 0.) / (Beta**(3/2) * T0**((3/2)*(Beta-1)))

    if ConstantTemperature == True:
        dT = 0
    Output = (dN, dT)
    return Output

def Beta2FitFunction(xData, N0i, b, Beta):
    #N0, N1 = N
    if N0Fixed == True:
        N0i = N0ii
    Init = N0i, T0 #Both start with same temperature
    Times = numpy.zeros(1)
    Ndat=len(xData)
    for i in range(1, len(xData[:])):
        if xData[i] != xData[i-1]:
            StarT = abs(xData[i] - xData[i-1])/Ndat
            Tim = numpy.linspace(xData[i-1]+StarT, xData[i], Ndat)
            Times = numpy.hstack((Times, Tim))
    #CHANGE THIS LINE TO CHANGE FIT FUNCTION!!
    FitLine = integrate.odeint(ArbLoss, Init, t = Times, args=(b, 2.))

    yData = FitLine[0::int(Ndat), 0]
    yData = numpy.array(FitLine[0, 0])
    it = 0
    for i in range(1, len(xData[:])):
        if xData[i] != xData[i-1]:
            it = it+1
        yData = numpy.hstack((yData, FitLine[int(it*Ndat), 0]))

    return yData

def CalcK2(ValueRaw, ErrorRaw, w):
    Value = ValueRaw*(((2*numpy.pi*kB)/(m*(w**2)))**1.5) #m^3 s^-1
    Error = ErrorRaw*(((2*numpy.pi*kB)/(m*(w**2)))**1.5) #m^3 s^-1
    return Value*1e6, Error*1e6 #cm^6 s^-1

def CalcBeta(Value,Error,w):
    Value = Value/(((2*numpy.pi*kB)/(m*(w**2)))**1.5) #m^3 s^-1
    Error = Error/(((2*numpy.pi*kB)/(m*(w**2)))**1.5) #m^3 s^-1
    return Value*1e-6, Error*1e-6

def twobody_fixed(t,N,T1):
    return N*(1/(1+(t/T1)*(numpy.exp(1)-1)))

def Ramsey_Fn(t,d0,phase,tau,N0i,b):
    fringes = numpy.cos(2*pi*d0*t+phase)
    decay = 0.5*numpy.exp(-t/tau)
    #lifetime = Beta2FitFunction(t,N0i,b,2)
    lifetime = twobody_fixed(t,N0i,b)
    return (decay*fringes+0.5)*lifetime

def chi_squared(y,exp,err):
    chi = (y-exp)**2/err**2
    chi = numpy.sum(chi)
    return chi
grid = gridspec.GridSpec(4,3,height_ratios=[1,1,0.2,0.5])
fig = pyplot.figure()
############################################################################
Data = numpy.genfromtxt(cwd+"\\data\\5_4_data.csv",delimiter=',')
order = numpy.argsort(Data[:,0])
Data = Data[order,:]
Nmax = numpy.amax(Data[:,1])
times = (Data[:,0])*1e-3

k2 = 4.8e-11
beta = CalcBeta(k2,0,w0)[0]
Init = Nmax, T0

x = [300,pi/2,0.08,1.5*Init[0],0.64]
x0 =[10,-pi,0,500,0.6399999]
x1 = [750,pi,0.11,2*Init[0],0.6400000001]

locs = numpy.where(Data[:,0]<=300)[0]

curve,cov = optimize.curve_fit(Ramsey_Fn,Data[locs,0]*1e-3,Data[locs,1],p0=x,
            bounds=(x0,x1))
err =numpy.sqrt(numpy.diag(cov))

print(curve[0],err[0])
print(curve[1],err[1])
times = numpy.linspace(0,0.3,3500)
Ramsey = Ramsey_Fn(times,*curve)
Ramsey_err= numpy.zeros(Ramsey.shape)

for i,e in enumerate(err):
    args_p = curve.copy()
    args_p[i]+=e
    Ramsey_err += (Ramsey_Fn(times,*args_p)-Ramsey)**2

Ramsey_err = numpy.sqrt(Ramsey_err)
#Lifetime = Beta2FitFunction(times,curve[3],curve[4],2)
Lifetime = twobody_fixed(times,*curve[3:])

ax1 = fig.add_subplot(111)
ax1.scatter(Data[:,0],Data[:,1]/1e3,color=colours['blue'],zorder=2)
ax1.plot(times*1e3,Ramsey/1e3,color=colours['grayblue'],zorder=1)
ax1.plot(times*1e3,0.5*Lifetime/1e3,color='k',zorder=1,ls='--')

ax1.fill_between(times*1e3,(Ramsey_err+Ramsey)/1e3,
                (Ramsey-Ramsey_err)/1e3,alpha=0.5,
                fc=colours['grayblue'],ec='none')
ax1.set_ylabel("Number of Molecules in $|0\\rangle$ ($10^3$)")
ax1.set_xlim(-1,82)
ax1.set_ylim(-0.05,2.65)
ax1.set_xlabel("Ramsey Time, $T$ (ms)")
######################################################################

f1 = 980.193e6
f2 = 979.991896e6

print(f1-f2-curve[0])

pyplot.tight_layout()

pyplot.savefig(cwd+"\\Ramsey.pdf")
pyplot.show()
