from importlibs import *
from statf.signalprocessing import *
import statf.dataprocflow as procflow


__name__ = "lfpts_fromNex"
__revision__ = "$Revision: 0.54 - 31 May 2011$"
__author__ = "Andrea Maffezzoli"


class readNex(object):
    """Apply struct.Struct to unpack/decode *.nex binary files
       Define the FileHeader and Spike Channel headers"""

    def __get_NexFileHeader(self): return struct.Struct(endian + '2i256sd4i256s')
    def __get_NexVarHeader(self): return struct.Struct(endian + '2i64s6i4d3i68s')
    evalbool = lambda s: s=='True' #not(s=="False") or (s=="True")
    NexFileHeader = property(__get_NexFileHeader)
    NexVarHeader = property(__get_NexVarHeader)
    fxattr = {'Frequency': float,
              'dttakeLFP': float,
              'type': int,
              'SRbin': float,
              'ISImin': float,
              'ILFPI': float,
              'Convolve': lambda x: True,
              'window_shape': lambda x: x,
              'window_size': int,
              'IIRhigh': float,
              'polenr': int,
              'CCFdelay': float,
              'SRscalefactor': float,
              'Mahalanobis_threshold': float,
              'storesingleWf': evalbool,
              'storeTs': evalbool,
              'storeTsStartEndWf': evalbool,
              'storeTsStartWfraw': evalbool,
              'storediscardedWf': evalbool}
              #'prethreshold': float,

    def __init__(self, path, filename):

        self.path, self.filename = path, filename
        for line in open(basepath + 'lfpts_settings.txt', 'r'):
            if line[0]=='#': continue
            tu = line.translate(None,'\n\r ').split('=') 
            if not(tu) or len(tu) != 2: continue
            tu1 = tu[1].partition('#')[0].strip()
            if tu[0] in self.fxattr: setattr(self, tu[0], self.fxattr[tu[0]](tu1))

    def __call__(self, Event_on = False, Continuous_on = False, filemapping = False):
        """Using memory mapping, or if it fails classic I/O 
           Get header info, process LFP/timestamps, close I/O streams"""
        
        f = open('%s\/MEAdata_plx\/%s.nex' % (self.path, self.filename), 'rb')
        if filemapping:
            try: self._f = mmap.mmap(f.fileno(), 0, None, mmap.ACCESS_READ)
            except WindowsError: self._f = f
        else: self._f = f
        self.readFileHeader()  
        self.readVariablesHeader()
        self.processData()
        self._f.close(); f.close()
   
    def readFileHeader(self):
        """.nex files report only spikes ADC freq.
           LFP freq. from settings or 1000Hz (previously set to spikes ADC freq.)"""

        MagicNumber, self.NexVersion, Comment, TSFrequency, \
            self.Beg, self.End, self.NumVars, self.NextFileHeader, Padding = \
            self.NexFileHeader.unpack(self._f.read(self.NexFileHeader.size))
        self.TSFrequency = float(TSFrequency)
        if not(getattr(self, 'Frequency', None)): self.Frequency = 1000 #self.TSFrequency

    def readVariablesHeader(self):
        """Only for neurons stored as variables"""
        
        _unpack, _size = self.NexVarHeader.unpack, self.NexVarHeader.size
        self.NumDSPChannelsFiring, self.TSCountsSum = 0, 0
        self.MaxTsCount, self.MaxLfpCount = 0, 0
        self.SpikeTs, self.Wf = {}, collections.defaultdict(list)
        LFPtype = getattr(self, 'type', 0)
        LFPtypestr = "wf_" + str(LFPtype)*(LFPtype>0)
        for v in xrange(self.NumVars):
            tu = _unpack(self._f.read(_size))
            Name, Count = tu[2].rstrip(chr(0)), tu[4]
            if 'Sweep' in Name: continue
            ##################################
            # Collecting spikes Ts
            ##################################
            if tu[0] in (0, 1) and '_0_1' not in Name and '_0_2' not in Name:
                _Name = Name.rstrip(chr(0)).replace('_unit_0', '')
                self.SpikeTs[_Name] = tu[3:5]
                self.MaxTsCount = max(self.MaxTsCount, Count)
            ##################################
            # Collecting Lfp's first Ts 
            ##################################
            if tu[0] == 3 and LFPtypestr in Name:
                self.Wf[Name[:-12]].append(tuple(tu[3:5] + tu[12:14]))
                self.MaxLfpCount = max(self.MaxLfpCount, Count)
                self.NPointsWaveLfp = tu[-4]
    
    def getTsStartWf(self):

        _f, longsz = self._f, long_size 
        def getTs(ts): return ts/self.TSFrequency
        def func(count): return map(getTs, array.array('l', _f.read(count*longsz)))
        return func

    def getWfraw(self):
        
        _f, shortsz, smoother, Array = self._f, short_size, self.smoother(), numpy.asarray
        def func(NPointsWave, ADtoMV, myarray):
            myarray.extend(smoother(Array(array.array('h', _f.read(NPointsWave*shortsz)))*ADtoMV))
        return func

    def getWfTsdata(self, n):

        def selectLFP(dt):
            """Input: start_end timestamps, data
               Sort input and scan its entries to find two subsequent start
               timestamps, then hold only the one with longest data array
               Output: a selection of the input"""
            
            it, Ix = iter(sorted((j[0], i) for i, j in enumerate(start_end_Wf))), []
            try: t0, i0 = it.next()
            except StopIteration: return set(Ix)
            while True:
                Ix.append(i0)
                try: t1, i1 = it.next()
                except StopIteration: return set(Ix)
                if t1-t0>dt: t0, i0 = t1, i1; continue
                try: t0, i0 = it.next()
                except StopIteration: return set(Ix)
        
        def getallTs(n):
            
            def getTs(ts): return ts/self.TSFrequency
            Spikes_t = self.SpikeTs.get(n, None)
            if not(Spikes_t): return []
            _f.seek(Spikes_t[0], 0)
            all_TS.extend(purgeISI(map(getTs, array.array('i', _f.read(Spikes_t[1]*int_size))), 
                                   getattr(self, 'ISImin', 0.003)))

        _f, Freq, ILFPI, samples_nr = self._f, self.Frequency, self.ILFPI, self.NPointsWaveLfp
        concat = numpy.concatenate

        #########################################################################
        # Get spikes and LFPs start timestamps, LFPs Waveform.
        #########################################################################
        allLfp_t = self.Wf.get(n, None)
        if not(allLfp_t): return None, None, None, None, None

        all_TS = []; getallTs(n)
        getWfraw = self.getWfraw()
        dataWf, all_TsStartWf, start_end_Wf = [], [], []
        for Lfp_t in allLfp_t:
            #########################################################################
            # Get only start TS of each LFP. Get Timestamps.
            # If Spikes_t not null: offset, tscount = Spikes_t
            # NB Call functions only in the order below!
            #########################################################################
            offset, count, ADtoMV, NPointsWave = Lfp_t 
            _f.seek(offset, 0)
            getTsStartWf = self.getTsStartWf()
            TsStartWf = getTsStartWf(count)
            #########################################################################
            # A deliberately big negative value is given to previousSTart, and None
            # to previousWf_raw before to start reading LFP Wfs
            # getWfraw is the function grabbing LFP data, curried with filter
            #########################################################################
            previousStart, previousWf_raw = -1e10, None
            for Start in TsStartWf:
                Wf_raw = []; getWfraw(NPointsWave, ADtoMV, Wf_raw)
                deltaT = Start - previousStart
                if deltaT <= ILFPI:
                    nextWf = concat((dataWf[-1][:int(Freq*deltaT)-NPointsWave], Wf_raw))
                    dataWf[-1] = nextWf[:Wfnr*samples_nr]
                    Wfnr += 1
                    start_end_Wf[-1][1] = start_end_Wf[-1][0] + len(Wf_raw) / Freq
                else:
                    dataWf.append(Wf_raw)
                    Wfnr = 1
                    start_end_Wf.append([Start, Start + NPointsWave/Freq])
                previousStart = Start
            all_TsStartWf.extend(TsStartWf)

        start_end_Wf = numpy.asarray(start_end_Wf)
        prevLFPnr = len(start_end_Wf)
        Ix = selectLFP(getattr(self, "dttakeLFP", self.ILFPI))
        print 'LFP discarded on %s: %d[%d-%d]' % (n, prevLFPnr-len(Ix),prevLFPnr,len(Ix))

        #########################################################################
        # I prefer to proceed only with numpy.arrays, so first of all, I have to
        # fillshortest arrays of zeros at their ends
        #########################################################################
        maxl = max(len(d) for i, d in enumerate(dataWf) if i in Ix)
        self.dataWf = numpy.zeros((len(Ix), maxl))
        j = 0
        for i, wf in enumerate(dataWf):
            if i in Ix: self.dataWf[j][:len(wf)] = wf[:]; j += 1
        #self.dataWf = numpy.asarray([concat((z, numpy.zeros(maxl-len(z))))
        #                             for i, z in enumerate(dataWf) if i in Ix])
        self.dataWfpurged = [tuple(wf) for i, z in enumerate(dataWf) if i not in Ix] \
                             if self.storediscardedWf else []
        self.start_end_Wf = start_end_Wf[sorted(Ix)]
        self.all_TS = numpy.array(all_TS)
        self.TsStartWf = sorted(all_TsStartWf)
        del dataWf

    def smoothfunc(self, Wf, wsize, wshape):
        
        if len(Wf) <= wsize: return numpy.asarray(Wf)
        return smooth(numpy.asarray(Wf), wsize, wshape)
    
    def smoother(self):

        if getattr(self, 'Convolve', None):
            def smoothWf(Wf): return self.smoothfunc(Wf, self.window_size, str(self.window_shape))
        else:
            def smoothWf(Wf): return numpy.asarray(Wf) 
        return smoothWf
    
    def filterWf(self, sampleFreq): return IIRfilter(getattr(self, 'IIRhigh', None), sampleFreq,
                                                     poleNr = getattr(self, 'polenr', 4))

    def histogram(self, T, tbin=None):

        if not(tbin): tbin = 1./self.Frequency
        return numpy.histogram(T,bins=numpy.arange(0,T[-1]+tbin,tbin))[0]

    def sortLFP(self):
        
        print "Setting data for LFP sorting.."
        _sorting = procflow.sorting(Mahalanobis_threshold = self.Mahalanobis_threshold)
        def doSorting(Data): return _sorting.__clustering__(Data, range(len(Data)))
        return doSorting
    
    def getSR(self, all_TS, start_end_Wf):
        """Get timestamps, and then SR, during LFP"""

        Ts_on_LFP = []
        binsearch = numpy.searchsorted
        for start, end in start_end_Wf:
            i, j = binsearch(all_TS, [start, end])
            if i != j: Ts_on_LFP.append((i, j, start))
        if not(Ts_on_LFP): return ()
        smoother = self.smoother()
        TS_h = []
        for i,j,start in Ts_on_LFP:
            all_TS_back = all_TS[i:j]-start
            TS_h.append(self.histogram(all_TS_back))
        SRbin = getattr(self, 'SRbin', None)
        TS_h = [self.histogram(all_TS[i:j]-start, SRbin) for i,j,start in Ts_on_LFP]
        TS_h_arr = numpy.zeros((len(TS_h), max(map(len, TS_h))))
        for h0, h1 in zip(TS_h, TS_h_arr): h1[:len(h0)] = h0[:]
        #return tuple(smoother(numpy.mean(TS_h_arr,axis = 0)*self.Frequency/self.SRscalefactor))
        self.SRs.append(tuple(smoother(numpy.mean(TS_h_arr,axis = 0)*self.Frequency/self.SRscalefactor)))

    def processData(self):
        
        #write_path = '%s\/MEAdata_cluster_results\/LFPdata\/' % self.path 
        #if not(os.path.exists(write_path)): os.mkdir(write_path, 0777)
        #########################################################################
        # Get all neurons involved in LFP and/or spiking activity
        ######################################################################### 
        self.allNeurons = sorted(set(self.SpikeTs.keys()+self.Wf.keys()), cmp=sort_neuro_names)
        sortLFP = self.sortLFP()
        filterWf = self.filterWf(self.Frequency)
        _bin = 1/self.Frequency
        _mean = lambda x: numpy.mean(x,axis = 0).tolist()   
        _sem = lambda x: (numpy.sqrt(numpy.var(x, axis=0)/float(len(x)))).tolist()   
         
        print 'Sorting LFPs data and saving results..'; 
        for n in self.allNeurons:
            self.getWfTsdata(n)
            if self.dataWf is None: continue

            #########################################################################
            # Find mean of (filtered) LFPs (for the latter SEM too) and LFPs.
            #########################################################################
            meanWf          = _mean(self.dataWf)
            all_filteredWf  = numpy.asarray(map(filterWf, self.dataWf)) # dataWf
            try: sortingWf = sortLFP(all_filteredWf)
            except:
                print 'Some error occurred when sorting LFPs of channel', n
                sortingWf = {}; continue
            dataclustWf = []
            for i, j in sortingWf.iteritems():
                dataclustWf.extend([[len(j[0])] + _mean(all_filteredWf[j[0]]), \
                                    [len(j[1])] + _sem(all_filteredWf[j[0]])])
            
            #########################################################################
            # Find SR for Ts on LFPs distinguished by cluster
            #########################################################################
            self.SRs = []
            self.getSR(self.all_TS, self.start_end_Wf)
            for j in sortingWf.itervalues(): self.getSR(self.all_TS, self.start_end_Wf[j[0]])
            self.SRs.extend((len(sortingWf)<3)*[('',)])
            
            #########################################################################
            # Various data writing (if enabled):
            # - time bins, mean of LFPs, its filtered version 
            # - binned spike timestamps histogram during LFPs
            # - spike rate of the latter timestamps
            # - all LFPs traces, all spikes timestamps
            #########################################################################
            #TS_shift = [0] * (int(self.prethreshold / bin) + 1)
            data2csv = [list(myXrange(0, max(t[1]-t[0] for t in self.start_end_Wf) + _bin, _bin)), \
                        [str(len(self.dataWf))] + meanWf[1:], \
                         _mean(all_filteredWf), _sem(all_filteredWf)] + \
                       self.SRs + dataclustWf + (len(sortingWf)<3)*[('',), ('',)]
            if self.storeTs: data2csv.append(tuple(self.all_TS))
            if self.storeTsStartEndWf: data2csv.extend(zip(self.start_end_Wf))
            if self.storeTsStartWfraw: data2csv.append(tuple(self.TsStartWf))
            if self.storesingleWf: data2csv.extend(self.dataWf.tolist())
            if self.storediscardedWf: data2csv.extend(self.dataWfpurged.tolist())
            data2csv = map(None, *data2csv) if len(data2csv) != 1 else [[x,' '] for x in data2csv[0]]
            csvwriterows_safe('%s.csv' % n, data2csv,
                              '%s\/MEAdata_cluster_results\/LFPdata\/' % self.path)
            del self.dataWf
            del self.start_end_Wf
            del self.all_TS
            del self.TsStartWf
            del self.dataWfpurged
            del data2csv
        

class lfpts_run(object):

    def __get_exp_to_process_file(self): return 'experiment_to_process_lfp.txt'
    exp_to_process_file = property(__get_exp_to_process_file)
    
    @repoRelease
    @catchException
    @addMEAdata
    def __init__(self):

        for path, name, _ in return_exp_to_process(self.exp_to_process_file): #,['.nex']):
            self.moreinfo = 'Experiment name: %s' % name
            self.stderr_file = path + '\/MEAdata_cluster_results\/logError.txt'
            print 'Experiment name: %s' % name
            self.processlfp(path, name)
            print '%s%s' % ('-'*70, '\n')
    
    @timer
    def processlfp(self, path, name): readNex(path, name)()


if __name__ == 'lfpts_fromNex':

    lfpts_run()



"""
    #def getTs(self, ts):
    #   prethresh, TSFreq = getattr(self,prethreshold,0), self.TSFrequency
    #   if prethresh: def func(ts): return (ts-prethresh) / TSFreq
    #   else: def func(ts): return ts / TSFreq
    #   return func

def smootherdec(self):
        
        r = raw_input("To smooth LFPs digit smoothing window size then Enter, otherwise only Enter: ")
        try:
            window_size, window_shape = int(r), str(self.window_shape)
            def smoothWf(Wf): return self.smoothfunc(Wf, window_size, window_shape)
        except ValueError:
            def smoothWf(Wf): return numpy.asarray(Wf)
        return smoothWf
"""
