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

jqc_plot.plot_style()

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

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 Rabi(x,t,O,x0,A):
    D = 2*pi*(x-x0)
    W = numpy.sqrt(O**2+D**2)
    F = (numpy.cos(W*t/2))**2+(D/W)**2*(numpy.sin(W*t/2))**2# (O/W)**2*numpy.sin(0.5*W*t)**2
    return A*(F)

t = [0,50e-6,100e-6,60e-6,80e-6,85e-6,125e-6]
r = [0,2*pi*6.23e3,2*pi*3.74e3,(2*pi)/173e-6,(2*pi)/240e-6,(2*pi)/264e-6,(2*pi)/375e-6]
e1 = [0,2*pi*6.36e3,2*pi*3.96e3,(2*pi)/171e-6,(2*pi)/237e-6,(2*pi)/262e-6,(2*pi)/372e-6]
e2 = [0,2*pi*6.1e3,2*pi*3.62e3,(2*pi)/175e-6,(2*pi)/243e-6,(2*pi)/266e-6,(2*pi)/378e-6]
for i in range(1,7):
    fname = "N"+str(i)
    print(fname)
    data = numpy.genfromtxt(cwd+"\\"+fname+".csv",delimiter=',')
    avgdata = AverageData(data)

    frequency = numpy.linspace(numpy.amin(data[:,0]),numpy.amax(data[:,0]),250)

    fitfn = lambda x,O,x0,A : Rabi(x*1e6,t[i],O,x0,A)

    X = [r[i],numpy.average(data[:,0])*1e6,numpy.amax(data[:,1])]
    X0 = [0.8*e2[i],numpy.amin(data[:,0])*1e6,0]
    X1 = [1.2*e1[i],numpy.amax(data[:,0])*1e6,2*numpy.amax(data[:,1])]

    curve,cov = curve_fit(fitfn,data[:,0],data[:,1],p0=X,bounds=(X0,X1))
    err = numpy.sqrt(numpy.diag(cov))
    figure=pyplot.figure(fname)

    ax = figure.add_subplot(111)

    ax.errorbar(avgdata[:,0],avgdata[:,1],yerr=avgdata[:,2],color='k',capsize=3.5,
                zorder=1.8,fmt='o')

    ax.plot(frequency,fitfn(frequency,*curve),color=jqc_plot.colours['reddish'],
            zorder=1.5)
    par0 = curve
    model = Rabi(frequency*1e6,t[i],*par0)

    ax.set_xlabel("Frequency (MHz)")
    ax.set_ylabel("Number of Molecules")
    print(curve,err)
    ymin,ymax = ax.get_ylim()
    ax.set_ylim(0,ymax)


    #ax.plot([theory,theory],[0,ymax],ls=':',color='gray')

    pyplot.tight_layout()
    figure.savefig(cwd+"\\"+fname+".pdf")
    figure.savefig(cwd+"\\"+fname+".png")
pyplot.show()
