import numpy
import os
from scipy.optimize import curve_fit
cwd = os.path.dirname(os.path.abspath(__file__))
Energies_noDv=numpy.genfromtxt(cwd+"\\Energies_noDv.csv",delimiter=',')
Energies = numpy.zeros((7,2))
Experiment = numpy.genfromtxt(cwd+"\\Summary_Rabi.csv",
                                delimiter=',',skip_header=1)

def X2(theory,exp,err):
    x=0
    for i,t in enumerate(theory):
        x+=((exp[i]-t)**2)/(err[i]**2)
    return x/(len(exp)-1)

for i in range(7):
    L = numpy.where(numpy.logical_and(Energies_noDv[:,1]==i,Energies_noDv[:,2]==i+5))[0]
    Energies[i,0] = Energies_noDv[L,1]
    Energies[i,1] = Energies_noDv[L,0]

def distort_energy(N,D):
    return -D*((N)*((N)+1))**2

def Distort(Narray,Dv):
    global Energies
    f = numpy.zeros(len(Narray))
    for i,N in enumerate(Narray):
        E1 = Energies[numpy.where(Energies[:,0]==N-1)[0],1]
        E2 = Energies[numpy.where(Energies[:,0]==N)[0],1]
        E1 = E1+distort_energy(N-1,Dv)

        E2 = E2+distort_energy(N,Dv)
        freq = E2-E1
        f[i] = freq
    return f

Curve,cov = curve_fit(Distort,Experiment[:,0],Experiment[:,1]*1e6,
            sigma=Experiment[:,2]*1e6,absolute_sigma=True,p0=[206])
err = numpy.sqrt(numpy.diag(cov))

print(Distort(Experiment[:,0],*Curve)*1e-6)
print(1e3*(Experiment[:,1]-Distort(Experiment[:,0],*Curve)*1e-6))
print(Curve,err)
print("after",X2(Distort(Experiment[:,0],*Curve),Experiment[:,1]*1e6,Experiment[:,2]*1e6))
print("b4",X2(Distort(Experiment[:,0],213.),Experiment[:,1]*1e6,Experiment[:,2]*1e6))
