import os
import re
import time

import numpy as np

from TDT import *

from IPython.core.debugger import Tracer

def SEV2py(SEV_DIR, *, CHANNEL=0, EVENTNAME='', T1=0, T2=0, FS=0, RANGES=None,
           VERBOSE=0, JUSTNAMES=0):
    '''
    SEV2MAT TDT SEV file format extraction.
    data = SEV2mat(SEV_DIR), where SEV_DIR is a string, retrieves
    all sev data from specified directory in struct format. SEV_DIR can
    also be a single file. SEV files are generated by an RS4 Data Streamer,
    or by enabling the Discrete Files option in the Synapse Stream Data
    Storage gizmo, or by setting the Unique Channel Files option in 
    Stream_Store_MC or Stream_Store_MC2 macro to Yes in OpenEx.

    data    contains all continuous data (sampling rate and raw data)
    
    optional parameters:
        T1          scalar, retrieve data starting at T1 (default = 0 for
                        beginning of recording)
        T2          scalar, retrieve data ending at T2 (default = 0 for end
                        of recording)
        CHANNEL     integer, returns the sev data from specified channel
                        only (default = 0 for all channels)
        RANGES      array of valid time range column vectors
        JUSTNAMES   boolean, retrieve only the valid event names
        EVENTNAME   string, specific event name to retrieve data from
        VERBOSE     boolean, set to false to disable console output
        FS          float, sampling rate override. Useful for lower
                        sampling rates that aren't correctly written into
                        the SEV header.
    '''

    data = structtype()

    MAX_UINT64 = np.iinfo(np.uint64).max
    
    if os.path.isfile(SEV_DIR):
        # treat as single file only
        sev_files = [SEV_DIR]
    elif os.path.isdir(SEV_DIR):
        # treat as directory 
        SEV_DIR = os.path.join(SEV_DIR, '')
        sev_files = get_files(SEV_DIR, '.sev')
    else:
        print('Error: Unable to find sev file or directory:\n\t', SEV_DIR)
        return None
   
    nfiles = len(sev_files)
    if nfiles < 1:
        if JUSTNAMES:
            return []
        print('Warning: no sev files found in', SEV_DIR)
        return None

    if FS > 0:
        print('Using {:.4f} Hz as SEV sampling rate for {}'.format(FS, EVENTNAME))
    
    file_list = []
    for file in sev_files:
        [filename, ext] = os.path.splitext(file)
        [path, name] = os.path.split(filename)
        file_list.append({'fullname':file,
                          'folder':path,
                          'name':name})
    
    
    chan_search = re.compile('_[Cc]h([0-9]*)')
    hour_search = re.compile('-([0-9]*)h')
    name_search = re.compile('_(.{4})_')
    
    # find out what data we think is here
    for file in file_list:
        
        # find channel number
        match_result = chan_search.findall(file['name'])
        if match_result:
            file['chan'] = int(match_result[-1])
        else:
            file['chan'] = -1
        
        # find starting hour
        match_result = hour_search.findall(file['name'])
        if match_result:
            file['hour'] = int(match_result[-1])
        else:
            file['hour'] = 0
        
        # event name of stream
        match_result = name_search.findall(file['name'])
        if match_result:
            file['eventName'] = match_result[-1]
        else:
            file['eventName'] = file['name']
                
        # check file size
        file['data_size'] = os.stat(file['fullname']).st_size - 40
        
        with open(file['fullname'], 'rb') as sev:
            # create and fill streamHeader struct
            streamHeader = structtype()
            
            streamHeader.fileSizeBytes   = np.fromfile(sev, dtype=np.uint64, count=1)[0]
            streamHeader.fileType        = np.fromfile(sev, dtype=np.uint8, count=3)
            streamHeader.fileType        = ''.join([chr(item) for item in streamHeader.fileType])
            streamHeader.fileVersion     = np.fromfile(sev, dtype=np.uint8, count=1)[0]
            streamHeader.eventName       = file['eventName']
            
            if streamHeader.fileVersion < 4:
                
                # prior to v3, OpenEx and RS4 were not setting this properly
                # (one of them was flipping it), so only trust the event name in 
                # header if fileVersion is 3 or higher
                eventName = np.fromfile(sev, dtype=np.uint8, count=4)
                eventName = ''.join([chr(item) for item in eventName])
                if streamHeader.fileVersion >= 3:
                    streamHeader.eventName = eventName
                
                # current channel of stream
                streamHeader.channelNum = np.fromfile(sev, dtype=np.uint16, count=1)[0]
                file['chan'] = streamHeader.channelNum
                # total number of channels in the stream
                streamHeader.totalNumChannels = np.fromfile(sev, dtype=np.uint16, count=1)[0]
                # number of bytes per sample
                streamHeader.sampleWidthBytes = np.fromfile(sev, dtype=np.uint16, count=1)[0]
                reserved = np.fromfile(sev, dtype=np.uint16, count=1)[0]
                
                # data format of stream in lower 3 bits
                dform = np.fromfile(sev, dtype=np.uint8, count=1)[0]
                dform &= 0b111
                streamHeader.dForm = ALLOWED_FORMATS[dform]
                
                # used to compute actual sampling rate
                streamHeader.decimate   = np.fromfile(sev, dtype=np.uint8, count=1)[0]
                streamHeader.rate       = np.fromfile(sev, dtype=np.uint16, count=1)[0]
            else:
                print('Error: unknown version', streamHeader.fileVersion)
                return None
            
            # compute sampling rate
            if streamHeader.fileVersion > 0:
                streamHeader.fs = np.power(2.,(streamHeader.rate - 12)) * 25000000 / streamHeader.decimate
            else:
                # make some assumptions if we don't have a real header
                streamHeader.dForm = 'single'
                streamHeader.fs = 24414.0625
                streamHeader.channelNum = file['chan']
                print('''Warning: {0} has empty header;
assuming {1} ch {2} format {3}
upgrade to OpenEx v2.18 or above\n'''.format(file['name'],
                                             streamHeader.eventName,
                                             streamHeader.channelNum,
                                             streamHeader.dForm))
            
            if FS > 0:
                streamHeader.fs = FS
            
            varname = fix_var_name(streamHeader.eventName)
            file['itemsize'] = np.uint64(np.dtype(streamHeader.dForm).itemsize)
            file['npts'] = file['data_size'] // file['itemsize']
            file['fs'] = streamHeader.fs
            file['dForm'] = streamHeader.dForm
            file['eventName'] = streamHeader.eventName
            file['varName'] = varname

    eventNames = list(set([file['eventName'] for file in file_list]))
    if JUSTNAMES:
        return eventNames

    if T2 > 0:
        validTimeRange = np.array([[T1], [T2]])
    else:
        validTimeRange = np.array([[T1], [np.inf]])

    try:
        len(RANGES)
        validTimeRange = RANGES
    except:
        pass
    
    numRanges = validTimeRange.shape[1]
    if numRanges > 0:
        data['time_ranges'] = validTimeRange
    
    for thisEvent in eventNames:

        if EVENTNAME and EVENTNAME != thisEvent:
            continue
        
        # get files for this event only
        file_list_temp = [file for file in file_list if file['eventName'] == thisEvent]
        
        # extract header info
        fs = file_list_temp[0]['fs']
        eventName = file_list_temp[0]['eventName']
        dForm = file_list_temp[0]['dForm']
        
        chans = [f['chan'] for f in file_list_temp]
        hours = [f['hour'] for f in file_list_temp]
        max_chan = np.max(chans)
        min_chan = np.min(chans)
        max_hour = np.max(hours)
        hour_values = sorted(list(set(hours)))
        
        # preallocate data array
        if CHANNEL > 0:
            try:
                chans.index(CHANNEL)
                search_ch = CHANNEL
            except:
                print('Warning: Channel {0} not found in {1} event'.format(CHANNEL, eventName))
                continue
        else:
            search_ch = min_chan
        
        # determine total samples if there is chunking, and how many samples are in each file
        total_samples = 0
        npts = [np.uint64(0) for i in hour_values]
        for jjj in hour_values:
            ind1 = np.asarray(hours) == jjj
            ind2 = np.asarray(chans) == search_ch
            temp_num = np.where(ind1 & ind2)[0]
            if len(temp_num) < 1:
                print('Error: matching file not found for hour', jjj, 'channel', matching_ch)
                return None
            elif len(temp_num) > 1:
                print('Error: too many matches found for hour', jjj, 'channel', matching_ch)
                return None
            temp_num = temp_num[0]
            npts[jjj] = np.uint64(file_list_temp[temp_num]['npts'])
            total_samples = total_samples + npts[jjj]
        
        # if we are doing time filtering, determine which files we need to read
        # from and how many samples
        absoluteStartSample = np.zeros(numRanges, dtype=np.uint64)
        absoluteEndSample = np.zeros(numRanges, dtype=np.uint64)
        startHourFile = np.zeros(numRanges, dtype=np.uint64)
        endHourFile = np.zeros(numRanges, dtype=np.uint64)
        startHourSamplesToSkip = np.zeros(numRanges, dtype=np.uint64)
        endHourSamplesEnd = np.zeros(numRanges, dtype=np.uint64)
        
        for jj in range(numRanges):
            
            absoluteStartSample[jj] = np.max(np.ceil(validTimeRange[0,jj] * fs), 0) + 1
            if np.isinf(validTimeRange[1,jj]):
                absoluteEndSample[jj] = total_samples
            else:
                absoluteEndSample[jj] = np.minimum(np.max(np.floor(validTimeRange[1,jj] * fs), 0) + 1, total_samples)            
            curr_samples = 0
            for jjj in hour_values:
                if curr_samples <= absoluteStartSample[jj]:
                    startHourSamplesToSkip[jj] = absoluteStartSample[jj] - curr_samples - 1
                    startHourFile[jj] = jjj
                if curr_samples + npts[jjj] >= absoluteEndSample[jj]:
                    endHourSamplesEnd[jj] = absoluteEndSample[jj] - curr_samples
                    endHourFile[jj] = jjj
                    break
                curr_samples = curr_samples + npts[jjj]
        
        # now allocate it
        if CHANNEL > 0:
            channels = [CHANNEL]
        else:
            channels = sorted(list(set(chans)))

        varname = file_list_temp[0]['varName']
        data[varname] = structtype()
        data[varname].data = [[] for i in range(numRanges)]
        for jj in range(numRanges):
            data[varname]['data'][jj] = np.zeros((len(channels), absoluteEndSample[jj] - absoluteStartSample[jj] + np.uint64(1)), dtype=dForm)
        data[varname]['channels'] = channels
        
        # loop through the time ranges
        for ii in range(numRanges):
            
            # loop through the channels
            for chan in channels:
                chanIndex = np.uint64(0)
                
                ind2 = np.asarray(chans) == chan
                
                # loop through the chunks
                for kk in range(startHourFile[ii], endHourFile[ii] + np.uint64(1)):
                    ind1 = np.asarray(hours) == kk
                    file_num = np.where(ind1 & ind2)[0][0]
                    
                    # open file
                    with open(file_list_temp[file_num]['fullname'], 'rb') as f:

                        # skip first 40 bytes from header
                        f.seek(40, os.SEEK_SET)
                        
                        # read rest of file into data array as correct format
                        data[varname]['name'] = eventName
                        data[varname]['fs'] = fs
                        
                        if kk == startHourFile[ii]:
                            firstSample = startHourSamplesToSkip[ii]
                        else:
                            firstSample = 0
                        
                        if kk == endHourFile[ii]:
                            lastSample = endHourSamplesEnd[ii]
                        else:
                            lastSample = MAX_UINT64
                        
                        # skip ahead
                        if firstSample > 0:
                            f.seek(int(firstSample * file_list_temp[file_num]['itemsize']), os.SEEK_CUR)
                        if lastSample == MAX_UINT64:
                            ddd = np.frombuffer(f.read(), dtype=dForm)
                        else:
                            ddd = np.frombuffer(f.read(int((lastSample - firstSample)*file_list_temp[file_num]['itemsize'])), dtype=dForm)
                        readSize = np.uint64(len(ddd))

                        if len(channels) > 1:
                            data[varname].data[ii][chan-1, int(chanIndex):int(chanIndex + readSize)] = ddd
                        else:
                            data[varname].data[ii][0, int(chanIndex):int(chanIndex + readSize)] = ddd
                        chanIndex = chanIndex + readSize
                    
                    if VERBOSE:
                        print(file_list[file_num])
        
        if numRanges == 1:
            data[varname].data = data[varname].data[ii]
        if len(channels) == 1:
            data[varname].data = data[varname].data[ii]
            
    return data
    
if __name__ == '__main__':
    #BLOCK_PATH = 'F:\\OLD_G\\TDT\\Synapse\\Tanks\\IZimptest-170516-094356\\Subject1-170516-094521\\'
    #BLOCK_PATH = 'F:\\OLD_G\\TDT\\Synapse\\Tanks\\IZimptest-170516-094356\\Subject1-170516-094521\\IZimptest-170516-094356_Subject1-170516-094521_EEG2_Ch1.sev'
    #BLOCK_PATH = '\\\\rs4-01001\\data\\Experiment150-181203-154010\\Subject1-181206-122240'
    #BLOCK_PATH = 'C:\\TDT\\Synapse\\Tanks\\Experiment150-181203-154010\\Subject1-181206-122240'
    
    BLOCK_PATH = 'C:\\TDT\\Synapse\\Tanks\\Experiment150-181203-154010\\Subject1-181206-122240'
    #data = SEV2py(BLOCK_PATH, EVENTNAME='EEG1', FS=5000)
    #data = SEV2py(BLOCK_PATH, CHANNEL=5, T2=10)
    #data = SEV2py(BLOCK_PATH, T2=10, CHANNEL=5)
    start = time.time()
    #data = SEV2py(BLOCK_PATH, RANGES=np.array([[5, 3550],[10, 3650]]))    
    #data = SEV2py(BLOCK_PATH, T2=10)
    data = SEV2py(BLOCK_PATH, CHANNEL=5)
    print(time.time() - start, 'elapsed')
    
    #data = SEV2py(BLOCK_PATH, RANGES=np.array([[5, 3550],[10, 3650]]))
    
    #import matplotlib.pyplot as plt
    #plt.plot(data.Ch11.data[0][0,:])
    #plt.show()
    
    #data = SEV2py(BLOCK_PATH, RANGES=np.array([[5, 3550],[10, 3650]]))
    #data = SEV2py(BLOCK_PATH, RANGES=np.array([[5, 3550],[10, 3650]]))
    #data = SEV2py(DEVICE='RS4-01001',TANK='Experiment150-181203-154010', BLOCK='Subject1-181206-122240')