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

h = constants.h
pi = numpy.pi


def AverageData(data):
    '''
        A function to take experimental data with repeats in the independent (x)
        variable and return the average and standard error of the dependent (y)
        variable in an array.

        Standard error is given by std_dev/sqrt(N-1) where N is the number of
        repeats.

        Output:
        out = numpy.ndarray((Nx,3)) where Nx is number of x variables

        out[:,0] -> unique values of independent variable (x)
        out[:,1] -> mean values of depdendent variable (y)
        out[:,2] -> standard error of depdendent variable (y)
    '''
    #order the array by first column
    order = numpy.argsort(data[:,0])
    data = data[order,:]
    #find unique occurrances of time, and how many there are
    Var,counts = numpy.unique(data[:,0],return_counts=True)
    # create a counter for where in data we are
    # we will be using that python slices do not include the last index given
    # i.e A[0:5] gives A[0],A[1],A[2],A[3],A[4] but:
    #    A[5:10] gives A[5],A[6],A[7],A[8],A[9] so we can continue from index+c
    index = 0
    #create an output array
    output = numpy.zeros((len(counts),3))
    for i,c in enumerate(counts):
        #number data is a slice from data
        number = data[index:index+c,1]
        #calculate and place data into output array
        output[i,:]=[Var[i],numpy.average(number),
                        numpy.std(number)/numpy.sqrt(c-1)]
        #increase index
        index+=c
    #all done
    return output

def lorentzian(x,x0,A,y0,w):
    '''
    simple lorentizan function as defined on wikipedia:
    https://en.wikipedia.org/wiki/Spectral_line_shape#Lorentzian
    '''
    y = (x-x0)/(w*0.5)
    L = 1/(1+y**2)
    return L*A+y0

def Rabi(x,t,O,x0,A):
    D = x-x0
    W = numpy.sqrt(O**2+D**2)
    F = (O/W)**2*numpy.sin(pi*W*t)**2
    return A*(1-F)

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

#load theory data
data = numpy.genfromtxt(cwd+"\\Hyperfine Output\\N6 Energy.csv",delimiter=',')
B = data[0,:]*1e4
MF = numpy.genfromtxt(cwd+"\\Hyperfine Output\\N6 MF.csv",delimiter=',')
#load spectroscopic summary
ss = numpy.genfromtxt(cwd+"\\spec_summary[python].csv",delimiter=',',skip_header=1)

energy181 = numpy.genfromtxt(cwd+"\\Hyperfine Output\\N9 energy [old].csv",delimiter=',',skip_header=1)
energy181_new = numpy.genfromtxt(cwd+"\\Hyperfine Output\\N9 energy [new].csv",delimiter=',',skip_header=1)

# relabel components of "data" to be easier to read
Energy =data[1:,:]

#create figure objects
jqc_plot.plot_style("wide")
#fig=pyplot.figure("ladder",figsize=(6.4,2*6.4))
fig=pyplot.figure("ladder",figsize=(6.4,12.25))
grid= gridspec.GridSpec(14,2,width_ratios=[1.2,1.],
                height_ratios=[0.15 if i%2==0 else 1.0 for i in range(2*7)])

#how much of the energy level structure to show per N
frequency_range = 0.05 #400 kHz
#frequency range for spectroscopy plots
s_range = 0.05
begin=1

frac_off =0.3

keys = ["X","b","c","d","e","f","g","h"]
time = [0,100,40,50,67.7,50,60,0]#us
time = [0,50,100,60,80,85,125,0]
rabifreq =[0,6.23e-3,3.74e-3,1/173,1/240,1/264,1/375]#MHz

output_array_summary = numpy.zeros((6,3))

for i in range(7):
    #create axis objects
    if i !=0:
        old = ax
        ax = fig.add_subplot(grid[2*(6-i):2*(6-i)+2,0],sharex=old)
    else:
        ax =fig.add_subplot(grid[2*(6-i):2*(6-i)+2,0])
    #create broken axes
    d=0.02

    if i ==0:
        offset=-0.00084e3#0.0006904e3
        poff = 0
        ax.set_xlabel("Magnetic Field (G)")
        ax.text(0.03,0.1,"(a)",transform=ax.transAxes,fontsize=20)
    elif i ==1:
        offset=979.529
        oldoff=poff
        poff = 0.98e3
    elif i ==2:
        offset=2.94025e3
        oldoff=poff
        poff = 2.94e3
    elif i ==3:
        offset =5.88128e3
        oldoff=poff
        poff = 5.88e3
        ax.set_ylabel("Energy/$h$ (MHz)\n+{:.2f} GHz".format(poff*1e-3))
        #ax.set_ylabel("Energy/$h$ (MHz)")
    elif i ==4:
        offset=9.80262e3
        oldoff=poff
        poff = 9.8e3
    elif i ==5:
        offset=14.7043e3
        oldoff=poff
        poff = 14.70e3
    elif i ==6:
        offset=20.5862e3
        oldoff=poff
        poff = 20.58e3
    if i!=3 and i !=0:
        ax.set_ylabel("+{:.2f} GHz".format(poff*1e-3))

    if i ==0:
        #remove broken axes
        ax.spines['top'].set_visible(False)
        ax.plot([-d,d],[1-d,1+d],color='k',transform=ax.transAxes,lw=1.5,
                                clip_on=False)
        ax.plot([1-d,1+d],[1-d,1+d],color='k',transform=ax.transAxes,lw=1.5,
                                clip_on=False)
        #break markers
        begin=0
        e0 = -6.348031792722524842e-1
        err_e0=0

        ax.errorbar([181.5],[e0-poff],yerr=[err_e0],color='k',fmt='o')

    elif i ==6:
        #remove broken axes
        ax.tick_params(which='both',bottom=False,labelbottom=False)
        ax.spines['bottom'].set_visible(False)
        ax.plot([-d,d],[-d,+d],color='k',transform=ax.transAxes,lw=1.5,
                                clip_on=False)
        ax.plot([1-d,1+d],[-d,+d],color='k',transform=ax.transAxes,lw=1.5,
                                clip_on=False)

        err_e0 = numpy.sqrt(err_e0**2+(ss[i-1,2])**2)

        ax.errorbar([181.5],[e0+ss[i-1,1]-poff],yerr=[err_e0],color='k',fmt='o')
        old.text(165,e0+0.4-oldoff,keys[i],fontsize=18)

        cp =patches.ConnectionPatch([181.5,e0+0.05-oldoff],
                            [181.5,e0+ss[i-1,1]-0.05-poff],
                            coordsA='data',coordsB ='data', color='k',
                            axesA=old,axesB=ax,arrowstyle="->",lw=1.5)
        fig.add_artist(cp)
        e0 = e0+ss[i-1,1]

    else:
        #remove broken axes
        ax.tick_params(which='both',bottom=False,labelbottom=False)
        ax.spines['bottom'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.plot([-d,d],[-d,+d],color='k',transform=ax.transAxes,lw=1.5,
                                clip_on=False)
        ax.plot([1-d,1+d],[-d,+d],color='k',transform=ax.transAxes,lw=1.5,
                                clip_on=False)
        ax.plot([-d,d],[1-d,1+d],color='k',transform=ax.transAxes,lw=1.5,
                                clip_on=False)
        ax.plot([1-d,1+d],[1-d,1+d],color='k',transform=ax.transAxes,lw=1.5,
                                clip_on=False)

        begin = end
        #plot points from spectroscopy summary file
        err_e0 = numpy.sqrt(err_e0**2+(1e-3*ss[i-1,2])**2)
        ax.errorbar([181.5],[e0+ss[i-1,1]-poff],yerr=[err_e0],color='k',fmt='o')
        old.text(165,e0+0.4-oldoff,keys[i],fontsize=18)

        cp =patches.ConnectionPatch([181.5,e0+0.05-oldoff],
                            [181.5,e0+ss[i-1,1]-0.05-poff],
                            coordsA='data',coordsB ='data', color='k',
                            axesA=old,axesB=ax,arrowstyle="->",lw=1.5)
        fig.add_artist(cp)
        e0 = e0+ss[i-1,1]


    ax.set_ylim(offset-poff,offset+1-poff)
    ax.text(0.05,0.3,"$N=${:d}".format(i),transform=ax.transAxes)
    #decide which lines to plot, I want 32*(2*N+1) levels for each N
    end = begin + 32*(2*i+1)
    e181= energy181[numpy.where(energy181[:,1]==i)[0],:]
    e181_new= energy181_new[numpy.where(energy181_new[:,1]==i)[0],:]
    #plot the energy levels
    for j in range(begin,end):
        if j!=begin:
            x = numpy.amin([E,1e-9*data[j,:]/h])-30e-6
        E = 1e-6*Energy[j,:]/h
        if MF[j] ==i+5:
            #spin-stretched transitions
            ax.plot(B,E-poff,color=jqc_plot.colours['red'],
                zorder=1.7)
            X = numpy.where(e181==i+5)[0]
            if i>=1:
                f_theory = (e181[X,0]-e_theory)*1e-6
            e_theory=e181[X,0]

            X = numpy.where(e181_new==i+5)[0]
            if i>=1:
                f_theory_new = (e181_new[X,0]-e_theory_new)*1e-6
            e_theory_new=e181_new[X,0]

        elif MF[j] ==i-1+5 and i>0:
            #pi transitions from a spin-stretched state
            ax.plot(B,E-poff,color=jqc_plot.colours['blue'],
                zorder=1.6)
        elif MF[j] ==i-2+5 and i >0:
            #sigma- transitions from a spin-stretched state
            ax.plot(B,E-poff,color=jqc_plot.colours['green'],
                zorder=1.5)
        else:
            #everything else
            ax.plot(B,E-poff,color=jqc_plot.colours['sand'],
                    zorder=1.2)
    begin=end
    if i>=1 and i<7:
        #plotting the spectroscopy we have data for
        if i!=1:
            x1,y1 = spect.transAxes.transform((0,1.+frac_off))
            x1,y1 = fig.transFigure.inverted().transform((x1,y1))
            x2,y2 = spect.transAxes.transform((1,2.+frac_off))
            x2,y2 = fig.transFigure.inverted().transform((x2,y2))
            #print(x1,y1,x2,y2)
            place = [x1,y1,x2-x1,y2-y1]
            spect = fig.add_axes(place)
            #spect = fig.add_subplot(grid[(7-i),1])
        else:
            spect = fig.add_subplot(grid[-1,1]) #generate axis
            spect.set_xlabel("Frequency (MHz)")
            pyplot.tight_layout()
            pyplot.subplots_adjust(hspace=0.1,left=0.18,bottom=0.07,right=0.94,top=0.99,
                                wspace=0.31)
        #load data:
        spectroscopy= numpy.genfromtxt(cwd+"\\New Data\\N{:d}.csv".format(i),delimiter=',')
        avgspectroscopy = AverageData(spectroscopy)
        #plot data:
        spect.errorbar(avgspectroscopy[:,0],1e-3*avgspectroscopy[:,1],yerr=1e-3*avgspectroscopy[:,2],
                        color=jqc_plot.colours['red'],fmt='o',capsize=3.5)
        #fitting: start with guess paramters:
        #guess = [ss[i-1,1],-1.5,numpy.amax(spectroscopy[:,1])*1e-3,ss[i-1,3]]
        #fit data to lorentzian
        #curve,cov = curve_fit(lorentzian,spectroscopy[:,0],
        #                            spectroscopy[:,1]*1e-3,p0=guess)

        fitfn = lambda x,A,B,O: Rabi(x,time[i],O,A,B)
        guess = [ss[i-1,1],numpy.amax(spectroscopy[:,1])*1e-3,rabifreq[i]]

        curve,cov = curve_fit(fitfn,spectroscopy[:,0],
                                    spectroscopy[:,1]*1e-3,p0=guess)


        err =numpy.sqrt(numpy.diag(cov))
        print(curve,err)
        output_array_summary[i-1,:]=[i,curve[0],err[0]]
        #plotting range
        freqs = numpy.linspace(numpy.amin(spectroscopy[:,0])-0.01,
                        numpy.amax(spectroscopy[:,0])+0.01,500)

        #plot model
        model = fitfn(freqs,*curve)
        spect.plot(freqs,model,color=jqc_plot.colours['reddish'],zorder=0.9)

        #include predicted transtion Frequency from energy levels
        spect.plot([f_theory,f_theory],[-5,5],ls=':',color='gray',zorder=0.8)
        spect.plot([f_theory_new,f_theory_new],[-5,5],ls='--',color='k',zorder=0.8)
        print("theory",f_theory)
        print("theory",f_theory_new)
        spect.set_yticks([0,2,4])
        spect.set_ylim(0,max([2.1,1.1e-3*numpy.amax(spectroscopy[:,1])]))

        if i ==1:
            xmin = 980.35
            spect.text(0.9,0.1,"(b)",transform=spect.transAxes,fontsize=20,
                        horizontalalignment='center')
        elif i==2:
            xmin = 1960.66
            spect.text(0.9,0.1,"(c)",transform=spect.transAxes,fontsize=20,
                        horizontalalignment='center')
        elif i ==3:
            xmin = 2940.98
            spect.text(0.9,0.1,"(d)",transform=spect.transAxes,fontsize=20,
                        horizontalalignment='center')
        elif i ==4:
            xmin = 3921.3
            spect.set_ylabel("Molecule Number ($10^3$)")
            spect.text(0.9,0.1,"(e)",transform=spect.transAxes,fontsize=20,
                        horizontalalignment='center')
        elif i ==5:
            xmin = 4901.6
            spect.text(0.9,0.1,"(f)",transform=spect.transAxes,fontsize=20,
                        horizontalalignment='center')
        elif i ==6:
            xmin = 5881.85
            spect.text(0.9,0.1,"(g)",transform=spect.transAxes,fontsize=20,
                        horizontalalignment='center')
        X =numpy.round(2*curve[0],2)/2

        spect.set_xticks([X-0.02,X+0.02])
        spect.set_xlim(curve[0]-frequency_range/2,curve[0]+frequency_range/2)

x1,y1 = spect.transAxes.transform((0,1+1.6*frac_off))
x1,y1 = fig.transFigure.inverted().transform((x1,y1))
x2,y2 = spect.transAxes.transform((1,2+frac_off))
x2,y2 = fig.transFigure.inverted().transform((x2,y2))

x3,y3 = ax.transAxes.transform((1,1))
x3,y3 = fig.transFigure.inverted().transform((x3,y3))
#print(x1,y1,x2,y2)
place = [x1,y1,x2-x1,y3-y1]

spect = fig.add_axes(place)
data = numpy.genfromtxt(cwd+"\\New Data\\N6Rabi.csv",delimiter=',')
spect.text(0.9,0.1,"(h)",transform=spect.transAxes,fontsize=20,horizontalalignment='center')
spect.set_xlabel("Pulse Time (µs)")
spect.set_xlim(0,300)
data = AverageData(data)
spect.errorbar(data[:,0],1e-3*data[:,1],yerr=1e-3*data[:,2],fmt='o',capsize=3.5
                ,color=jqc_plot.colours['green'],zorder=1.5)

fitfn = lambda x,w,A: A*numpy.cos(numpy.pi*w*x)**2
test = [2e-3,1750]
curve,cov = curve_fit(fitfn,data[:,0],data[:,1],p0=test)
print(curve)
times = numpy.linspace(0,800,300)

spect.plot(times,1e-3*fitfn(times,*curve),color=jqc_plot.colours['greenish'],zorder=1.1)
spect.set_yticks([0,1,2])
spect.set_ylim(0,2.2)
spect.set_xlim(0,800)
ax.set_xlim(0,200)
numpy.savetxt(cwd+"\\Summary_Rabi.csv",output_array_summary,delimiter=',',header="N,f,Err")

pyplot.savefig(cwd+"\\ladder.pdf")
pyplot.savefig(cwd+"\\ladder.png")
pyplot.show()
