######################################################################################
##Revised Based on Yorick_Multitau_Code Obtained from Dr. Andrei Fluerasu###############
####################   Coded by Dr. Yugang Zhang #################################
####################   631-885-4714   ############################################
####################   Email: yuzhang@bnl.gov     ####################################
########################################################################################
#################### or: yuzhangnew@icloud.com      ##################################
####################  At Brookhaven National Lab   #########################################
############################## Dec 20, 2015, Ver0##########################################
##############################Jan 26, 2015 ,Ver1, Add array_method################################
#######Feb 01,Ver2, Add read_pixellist, qind, qradi from a txt file############
########Feb 18,Ver3, Add two-time correlation function from large array ########
######Feb 22,Ver4, Add two-time correlation function by image_read method#######
#######Feb 28,Ver5, corrected time-index of two-time correlation function by image_read method##########
##########March 1,Ver6 add the method get one-time from the diag of two-time##################
##########March 2,Ver7 do cleaning##############################################
##########March 7,Ver8 add frame_inten,fram_pixel_waterfall,two_time_show2##################################
#######################################################################################
######################################################################################
######################################################################################
from numpy import arange, array,zeros,linspace,round,hypot, ones,sqrt, int_, hstack, vstack
from numpy import digitize,where,copy,indices,histogram,bincount,tril,diag
from numpy import savetxt, log, column_stack,intersect1d, save, load
from math import sin,cos,pi
#from scipy.interpolate import interp1d
#import matplotlib.pyplot as plt
import time
import os
import sys
#from GetEdf import get_edf,sumfile
#from create_and_read_data import cpopen,cpdump
#from create_and_read_data import make_img_array,make_img_pixel_array
#from load_data import  read_ring_data, read_txt
#from img_process import show_img,show_mask,show_edf_with_mask, interp_zeros
from img_process import *
from Init_for_Bfiber import * # the setup file

T=True
F=False

class xpcs( object):
    def __init__(self, PAR=PAR):
        """ DOCUMENT __init__( PAR )
        the initilization of the XPCS class
        KEYWORD:  PAR, gives the parameters, including nobuf,nolev,begframe,
        noframes, if dark_img, it will use DK as dark_img """
            
        global nobuf,nolev,begframe,noframes
        self.version='version_11'
        #print PAR
        if PAR['dark_img']:
            if DK==None:print 'Get dark_img from data.';self.backg()  #to get a background
            else:print 'Get dark_img from a file.'
        else:print 'No Dark Images are used!'
        nobuf=PAR['nobuf']
        nolev=PAR['nolev']
        begframe=PAR['begframe'] 
        noframes=PAR['noframes'] 
        self.noframes=noframes
        self.begframe = begframe
    def clear(self,):
        ''' DOCUMENT clear( )
        clear all the global variables
        KEYWORD:  None   '''
        global buf,G,IAP,IAF,num,cts,cur,g2,gmax,sigg2
        buf=0;G=0;IAP=0;
        IAF=0;num=0;cts=0;
        cur=0;g2=0;gmax=0;sigg2=0
        ttx=0;Ndel=0;Npix=0;
        g2a=0;g2a2=0;g12=0;g12a=0;g12L=0;g12s=0;
        g12x=0;g12y=0;g12z=0
        g12r=0;
        
        
    def backg(self,):
        ''' DOCUMENT backg( )
        return DK: the average of a series of dark images.
        KEYWORD:  None   '''  
        global DK
        DK = sumfile( FILENAME, first_dark, first_dark+number_darks-1, avgflag=1, ext='.edf')

    def delays(self,time=1,nolevs=None,nobufs=None, correct=True): 
        ''' DOCUMENT delays(time=)
        return array of delays.
        KEYWORD:  time: scale delays by time ( should be time between frames)
        '''
        global nolev,nobuf,dly, dict_dly
        if nolevs==None:nolevs=nolev
        if nobufs==None:nobufs=nobuf
        if nobufs%2!=0:print "nobuf must be even!!!"    
        dly=zeros( (nolevs+1)*nobufs/2 +1  )        
        dict_dly ={}
        for i in range( 1,nolevs+1):
            if i==1:imin= 1
            else:imin=nobufs/2+1
            ptr=(i-1)*nobufs/2+ arange(imin,nobufs+1)
            dly[ptr]= arange( imin, nobufs+1) *2**(i-1)            
            dict_dly[i] = dly[ptr-1]            
        dly*=time
        #self.dly=dly*timeperframe
        if correct:self.dly=dly[1:]*timeperframe
        else:self.dly=dly[:-1]*timeperframe
        
    def get_min_dly(self):
        ''' DOCUMENT get_min_dly( )
        return the min value of the dly_array.
        KEYWORD:  None
        '''
                
        global min_ind_dly  
        if len(where(dly-noframes >=0 )[0])==0:min_ind_dly = len(dly)
        else:min_ind_dly =  min( where(dly-noframes >0)[0] ) -1
        #print min_ind_dly
        
    def make_qlist(self):
        ''' DOCUMENT make_qlist( )
        return a qlist by giving the noqs, qstart,qend,qwidth.
        KEYWORD:  None
        ''' 
        global qwidth,noqs,qend,qstart,qnum, qlist, qradi
        global qlist,qradi
        qradi = linspace(qstart,qend,noqs)
        qlist=zeros(2*noqs)
        qlist[::2]= round(qradi-qwidth/2)  #render  even value
        qlist[1::2]= round(qradi+(1+qwidth)/2) #render odd value

        qlist[::2]= int_(qradi-qwidth/2)  #render  even value
        qlist[1::2]= int_(qradi+(1+qwidth)/2) #render odd value
        if qlist_!=None:qlist=qlist_       

    def calqlist(self, qmask=None ,  shape='circle',
                 asym=False,theta=0, center=None):
        ''' DOCUMENT calqlist( qmask=,shape=, )
        calculate the equvilent pixel with a shape,
        return pixellist,nopixels,qind,nopr        
        KEYWORD:  qmask, a mask file;
            shape='circle', give a circle shaped qlist
            shape='column', give a column shaped qlist
            shape='row', give a row shaped qlist
            shape='slope', give a line-shape qlist  by
            the rotation of a row line with center and tan(theta) degree 
            "to be finished yet for the theta"
            
        '''
        global pixellist,nopixels,qind,nopr
        self.make_qlist()
        y, x = indices( [dimx,dimy] )
        if shape=='circle':
            y_= y- ybar +1;x_=x-xbar+1        
            r= int_( hypot(x_, y_)    + 0.5  )#why add 0.5?
        elif shape=='column': 
            r= x
        elif shape=='row':
            r=y
        elif shape=='slope':
            if center==None:cx,cy=[(dimy/2.),(dimx/2.)]
            else:cx,cy=center[1]/1.,center[0]/1.
            a=  cx/(cx+theta*cy) 
            b=theta*cx/(cx+theta*cy)
            print a,b,cx,cy
            r = x*a + y*b
            r=array(r,dtype=int32)
            #print r
        r= r.flatten() 
        #print qlist
        noqrs = len(qlist)    
        qind = digitize(r, qlist)        
        if qmask==None:
            w_= where( (qind)%2 )# qind should be odd;print 'Yes'
            w=w_[0]
        else:
            a=where( (qind)%2 )[0]            
            b=where(  mask.flatten() )[0]            
            #print a.shape,b.shape
            w= intersect1d(a,b)
        
        nopixels=len(w)
        qind=qind[w]/2
        #pixellist= (   y*imgwidth +x ).flatten() [w]
        pixellist= (   y*dimy +x ).flatten() [w]
        nopr,bins=histogram( qind, bins= range( len(qradi) +1 ))       

        
    ###########################################################################
    ########for frame_inten######################################
    ##################################################################
    def fram_inten( self, frame_start, frame_end):
        global fint
        nofra = frame_end+1-frame_start
        fint= zeros( [nofra, noqs])
        start_time = time.time()        
        for i in range(0, nofra ):
            n= frame_start + i
            fp = FILENAME + '%04d'%n 
            
            if img_format=='EDF':fp+='.edf';img= get_edf( fp ) -  DK
            elif img_format=='TIFF':fp+='.tiff';img = scipy.misc.imread(fp,flatten=1)
            else:img= cpopen( n= n,prefix= 'data_', inDir=DATA_DIR)
            
            imgf=img.flatten()[pixellist]
            img=[]            
            fint[i]=  (histogram(qind, bins=noqs, weights= imgf))[0]
            if  int(i %( (frame_end+1-frame_start) /10)) ==0:                
                    sys.stdout.write("#")
                    sys.stdout.flush()
        elapsed_time = time.time() - start_time
        print 'Total time: %.2f min' %(elapsed_time/60.)
        #return fint 
        
    ###########################################################################
    ########for fram_pixel_waterfall######################################
    ##################################################################
    ##fram_pixel_waterfall
        
    def fram_pixel_waterfall( self, frame_start, frame_end, q=None):
        global pintq,pint_dict        
        start_time = time.time()
        nofra = frame_end+1-frame_start        
        if q==None:
            pint_dict={}
            for i in range(noqs):
                pint_dict[i]=zeros( [nofra, nopr[i]])
        else:
            pintq = zeros( [nofra, nopr[q]])
            #print pintq.shape
        for i in range(0, nofra ):
            n= frame_start + i
            fp = FILENAME + '%04d'%n

            if img_format=='EDF':fp+='.edf';img= get_edf( fp ) -  DK
            elif img_format=='TIFF':fp+='.tiff';img = scipy.misc.imread(fp,flatten=1)
            else:img= cpopen( n= n,prefix= 'data_', inDir=DATA_DIR)        
                      
            imgf=img.flatten()[pixellist]
            img=[]            
            if q!=None:                
                pintq[i]= imgf[ qind==q]
                #print pintq[i].shape
            else:
                for q_ in range(noqs):
                    pint_dict[q_][i] = imgf[ qind==q_] 
            if  int(n %( (frame_end+1-frame_start) /10)) ==0:                
                    sys.stdout.write("#")