########################################################################
#Implements a multiple-time correlation scheme...
#f(tau) = <A(t)A(t+tau)>
#
#Adapted from:
#J. Ramirez, S. K. Sukumaran, B. Vorselaars, A. E. Likhtman
#J. Chem. Phys. 133, 154103 (2010).
########################################################################
########################################################################
#STANDARD USAGE:
#                       Note that the first two lines following this one are required syntax
#c=Correlator()
#c.setsize(S,p,m)       Here S is the number of correlators. Default values are S=32,p=16,m=2
#c.setdt(deltaT)        Optional call which changes the time interval from the default of 1.0
#c.add(val1)
#c.add(val2)
#...
#...
#c.evaluate(n=0)        After "evaluate" is called, c.t (the time) and c.f (value of TCF) are available
#                       Argument controls whether to subtract off (accumulate/N_c)^2 from the answer (n==0) or not
#
#                       After being called, c.output() will print these values, or the arrays can be directly accessed
########################################################################
#Also available is a saving function:
#c.savecor('/DESTINATION_DIRECTORY/FILENAME'), which will save the current state of the correlator as a text file,
#
########################################################################
#And an accompanying loading function:
#USAGE:
#
#c2=Correlator()
#c2.loadcor('/DESTINATION_DIRECTORY/FILENAME')
#
#                       This will load a correlator saved with the savecor method
#                       As seen above, the usage of this does not require setsize, or other initialization of the correlator;
#				this info is saved in the file and loaded
#########################################################################


class Correlator:

    ###############
    ##CONSTRUCTOR##
    ###############
    def __init__(self):
        self.shift = []         #where the coming values are stored
        self.correlation = []   #array with the actual correlation function
        self.ncorrelation = []   #number of values accomulated in cor
        self.accumulator = []   #accumulator in each correlator
        self.naccumulator = []  #index controlling accumulation in each correlator
        self.insertindex = []    #index pointing at the position where the current value is inserted

        self.numcorrelators = 0 #number of correlators
        self.dmin = 0           #minimum distance between points for correlators k>0; dmin = p/m
        self.length = 0         #length of result arrays
        self.kmax = 0           #maximum correlator attained during simulation
        self.deltat = 1.0       #size of the timestep

        self.p = 0              #points per correlator
        self.m = 2              #number of points over which to average; Recommended p mod m = 0; m = 2
        self.t = []
        self.f = []
        self.npcorr = 0
        self.accval = 0.0       #accumulated result of incoming values



    ###############
    ###FUNCTIONS###
    ###############


    ##################################################
    #Set size of correlator data structures
    ##################################################
    def setsize(self,numcorrin=32,pin=16,mval=2):
        self.numcorrelators = numcorrin
        self.p = pin
        self.m = mval
        self.dmin = self.p/self.m
        self.length = self.numcorrelators*self.p

        self.insertindex = [0]*self.numcorrelators
        self.naccumulator = [0]*self.numcorrelators
        self.accumulator = [0.0]*self.numcorrelators
        self.ncorrelation = []
        self.correlation = []
        self.shift = []

        self.t = [0]*self.length
        self.f = [0]*self.length

        self.initialize()


    ##################################################
    #initialize all values (current and averages) to zero
    ##################################################
    def initialize(self):
        self.shift=[]
        self.correlation=[]
        self.ncorrelation=[]
        for j in range(self.numcorrelators):
            self.shift.append([float(-2e10)]*self.p)
            self.correlation.append([float(0.0)]*self.p)
            self.ncorrelation.append([int(0)]*self.p)
            self.accumulator[j] = float(0.0)
            self.naccumulator[j] = int(0)
            self.insertindex[j] = int(0)
        for i in range(self.length):
            self.t[i] = float(0.0)
            self.f[i] = float(0.0)
        self.npcorr = 0
        self.kmax = 0
        self.accval = 0.0


    ##################################################
    #Change the time step away from the default
    ##################################################
    def setdt(self,dt):
        self.deltat = float(dt)


    ##################################################
    #take the stored array values and convert
    #them to the value of the autocorrelation function
    ##################################################
    def evaluate(self,norm=0):
        im = 0
        aux = 0.0
        if(norm==0):
            aux=(self.accval/float(self.ncorrelation[0][0]))*(self.accval/float(self.ncorrelation[0][0]))
        #first correlator
        for i in range(self.p):
            if (self.ncorrelation[0][i]>0):
                self.t[im] = i*self.deltat
                self.f[im] = self.correlation[0][i]/float(self.ncorrelation[0][i])-aux
                im = im + 1
        #subsequent correlators
        for k in range(1,self.kmax):
            for i in range(self.dmin,self.p):
                if (self.ncorrelation[k][i]>0):
                    self.t[im] =float( i*pow(self.m,k))*self.deltat
                    self.f[im] = self.correlation[k][i]/float(self.ncorrelation[k][i])-aux
                    im = im +1
        self.npcorr = im


    ##################################################
    #add a scalar to correlator number k
    ##################################################
    def add(self,w,k=0):
        #discard new values beyond the maximum scope of time correlator
        if (k == self.numcorrelators):
            return None
        if (k > self.kmax):
            self.kmax = k

        #first, insert new value in shift array
        self.shift[k][self.insertindex[k]] = w

        #add to average value
        if(k == 0):
            self.accval = self.accval + w

        #add to accumulator and, if needed, add to next correlator
        self.accumulator[k] = self.accumulator[k]+w
        self.naccumulator[k] = self.naccumulator[k] + 1
        if (self.naccumulator[k] == self.m):
            self.add((self.accumulator[k]/self.m),k+1)
            self.accumulator[k]=0.0
            self.naccumulator[k] = 0

        #calculate correlation functions
        ind1 = self.insertindex[k]
        if (k == 0):
            ind2 = ind1
            for j in range(self.p):
                if(self.shift[k][ind2]> -1e10):
                    self.correlation[k][j] = self.correlation[k][j] + self.shift[k][ind1]*self.shift[k][ind2]
                    self.ncorrelation[k][j] = self.ncorrelation[k][j] + 1
                ind2 = ind2 - 1
                if (ind2 < 0):
                    ind2 = ind2 + self.p
        else:
            ind2 = ind1 - self.dmin
            for j in range(self.dmin,self.p):
                if (ind2 <0):
                    ind2 = ind2 + self.p
                if (self.shift[k][ind2] > -1e10):
                    self.correlation[k][j] = self.correlation[k][j] + self.shift[k][ind1]*self.shift[k][ind2]
                    self.ncorrelation[k][j] = self.ncorrelation[k][j] + 1
                ind2 = ind2 - 1
        self.insertindex[k] = self.insertindex[k] + 1
        if (self.insertindex[k] == self.p):
            self.insertindex[k] = 0


    ##################################################
    #Output already-evaluated self.t and self.f values to screen
    ##################################################
    def output(self):
        for i in range(self.npcorr):
            print 'time = {:f}  C(t) = {:f}'.format(self.t[i],self.f[i])


    ##################################################
    #Load a correlator object from file
    ##################################################
    def loadcor(self,filename):
        f=open(filename,'r')
        lines=f.readlines()
        f.close()
        #read single values
        self.p=int((lines[1].split())[1])
        self.m=int((lines[2].split())[1])
        self.numcorrelators=int((lines[5].split())[1])
        #setsize and initialize
        self.setsize(self.numcorrelators,self.p,self.m)
        self.initialize()
        #continue reading single values
        self.npcorr=int((lines[3].split())[1])
        self.kmax=int((lines[4].split())[1])
        self.dmin=int((lines[6].split())[1])
        self.length=int((lines[7].split())[1])
        self.accval=float((lines[8].split())[1])
        self.deltat=float((lines[9].split())[1])


        #load arrays, starting with the arrays of length self.numcorrelators
        curline = 11
        for j in range(self.numcorrelators):
            line=lines[curline].split()
            self.insertindex[j] = int(line[0])
            self.accumulator[j] = float(line[1])
            self.naccumulator[j]= int(line[2])
            curline = curline+1
        #load t and f arrays (of length self.length)
        curline = curline+1
        for j in range(self.length):
            line=lines[curline].split()
            self.t[j] = float(line[0])
            self.f[j] = float(line[1])
            curline = curline+1
        #load correlator, ncorrelator, and shift arrays
        curline = curline+1
        for j in range(self.numcorrelators):
            for i in range(self.p):
                line=lines[curline].split()
                self.ncorrelation[j][i] = int(line[0])
                self.correlation[j][i] =float(line[1])
                self.shift[j][i]       =float(line[2])
                curline=curline+1


    ##################################################
    #Save the correlator object to a file
    ##################################################
    def savecor(self,filename):                #save the current state to a file
        import time
        f=open(filename,'w')
        f.write('Correlator saved on {:d}/{:d} at {:d}:{:d}\n'.format(time.localtime()[1],time.localtime()[2],time.localtime()[3],time.localtime()[4]))
        #first write single values
        f.write('p\t{:d}\n'.format(self.p))
        f.write('m\t{:d}\n'.format(self.m))
        f.write('npcorr\t{:d}\n'.format(self.npcorr))
        f.write('kmax\t{:d}\n'.format(self.kmax))
        f.write('numcorrelators\t{:d}\n'.format(self.numcorrelators))
        f.write('dmin\t{:d}\n'.format(self.dmin))
        f.write('length\t{:d}\n'.format(self.length))
        f.write('accval\t{:.16f}\n'.format(self.accval))
        f.write('deltat\t{:.16f}\n'.format(self.deltat))

        # write arrays of length self.numcorrelators
        f.write('insertindex\t accumulator\t naccumulator\n')
        for j in range(self.numcorrelators):
            f.write('{:d}\t{:.16f}\t{:d}\n'.format(self.insertindex[j],self.accumulator[j],self.naccumulator[j]))

        #write arrays of length self.length
        f.write('t\tf\n')
        for j in range(self.length):
            f.write('{:f}\t{:f}\n'.format(self.t[j],self.f[j]))

        #write two-dimensional arrays of length (numcorrelators,p)
        f.write('ncorrelation\tcorrelation\tshift\n')
        for j in range(self.numcorrelators):
            for i in range(self.p):
                f.write('{:d}\t{:.16f}\t{:.16f}\n'.format(self.ncorrelation[j][i],self.correlation[j][i],self.shift[j][i]))

        f.close()




"""
Here's some random sample
c=Correlator();c.setsize(6,16,2);c.initialize()

for j in range(1,50000):
    c.add(pow(float(j),-.5))
"""