import numpy as np
import os
from spkeras.datasets.load_event import loadaerdat

path_to_dataset = './cifar10_dvs/'   #Path to the directory of dataset
split_ratio = 0.9   # 90% samples for training

N = 1                  # Frame number of single sample
split_num = 1          # Split one sample into samples
H = 128     
W = 128
t_start = 0.0e6        # start time 
t_tot = t_start+1.3e6*split_num  # compute total time of spike train


file_num = 0
y_train = []
x_train = []
y_test = []
x_test = []


for root,dirs,files in os.walk(path_to_dataset):
    root = str(root)
    num = 0
    for names in dirs:
        class_names.append(names)
        
for root,dirs,files in os.walk(path_to_dataset):
    root = str(root)
    _num = 0
    sample_num = len(files)
    sample_array = np.arange(sample_num)
    sample_for_train = int(sample_num*split_ratio)
    #np.random.shuffle(sample_array)     #shuffle train and test
    for names in files:
        names = str(names)
        if names.endswith(".aedat"):
            path = root+'/'+names 
            if sample_array[_num] < sample_for_train:
                _train = True
            else:
                _train = False
            _num += 1
            file_num +=1 
            print('\r file number: ',file_num,end='')

            data = loadaerdat(path)

            data_t = data['t']-data['t'][0]
            loc_tot = np.where(np.logical_and(data_t > t_start, data_t < t_tot))
            data_t = data_t[loc_tot[0]]-data_t[loc_tot[0]][0]
            slot = (data_t[-1])/split_num

            if _train:
                for n in range(split_num):
                    __p = n + 1
                    for _n in range (N):
                        _p=_n+1
                        #_frame = np.zeros(shape=[128,128,N*2])
                        frame = np.zeros(shape=[2, H * W])
                        loc = np.where(np.logical_and(data_t < __p*slot-(N-_p)/N*slot, 
                                                          data_t >= (__p-1)*slot+_n/N*slot))
                        x = data['x'][loc[0]].astype(int)  # avoid overflow
                        y = data['y'][loc[0]].astype(int)
                        p = data['p'][loc[0]]
                        mask = []
                        mask.append(p == 0)
                        mask.append(np.logical_not(mask[0]))
                        for c in range(2):
                            position = y[mask[c]] * W + x[mask[c]]
                            events_number_per_pos = np.bincount(position)
                            frame[c][np.arange(events_number_per_pos.size)] += events_number_per_pos

                        if _n == 0:
                            frame = frame.reshape((2, H, W))
                            frame = np.moveaxis(frame, 0, -1)
                            frame = np.expand_dims(frame,axis=0)
                            _frame = frame
                        else:
                            frame = frame.reshape((2, H, W))
                            frame = np.moveaxis(frame, 0, -1)
                            frame = np.expand_dims(frame,axis=0)
                            _frame = np.append(_frame,frame,axis=0)

                    x_train.append(_frame)
                    num=0
                    for subname in class_names:
                        _path = path.split("/")
                        if subname == _path[-2]:
                            y_train.append(num)
                        num += 1                
            else:
                for n in range(split_num):
                    __p = n + 1
                    for _n in range (N):
                        _p=_n+1
                        #_frame = np.zeros(shape=[128,128,N*2])
                        frame = np.zeros(shape=[2, H * W])
                        loc = np.where(np.logical_and(data_t < __p*slot-(N-_p)/N*slot, 
                                                          data_t >= (__p-1)*slot+_n/N*slot))

                        x = data['x'][loc[0]].astype(int)  # avoid overflow
                        y = data['y'][loc[0]].astype(int)
                        p = data['p'][loc[0]]
                        mask = []
                        mask.append(p == 0)
                        mask.append(np.logical_not(mask[0]))
                        for c in range(2):
                            position = y[mask[c]] * W + x[mask[c]]
                            events_number_per_pos = np.bincount(position)
                            frame[c][np.arange(events_number_per_pos.size)] += events_number_per_pos
                        if _n == 0:
                            frame = frame.reshape((2, H, W))
                            frame = np.moveaxis(frame, 0, -1)
                            frame = np.expand_dims(frame,axis=0)
                            _frame = frame
                        else:
                            frame = frame.reshape((2, H, W))
                            frame = np.moveaxis(frame, 0, -1)
                            frame = np.expand_dims(frame,axis=0)
                            _frame = np.append(_frame,frame,axis=0)

                    x_test.append(_frame)
                    num=0
                    
                    for subname in class_names:
                        _path = path.split("/")
                        if subname == _path[-2]:
                            y_test.append(num)
                            continue
                        num += 1      
                        
x_test = np.array(x_test)
y_test = np.array(y_test)
x_train = np.array(x_train)
y_train = np.array(y_train)
loc = np.arange(x_train.shape[0])
np.random.shuffle(loc)
x_train = x_train[loc]
y_train = y_train[loc]
np.savez('./T_cifar10_dvs'+'_'+str(int((t_tot-t_start)*1e-3))+'ms_S'+str(split_num)+'_N'+str(N)+'.npz',
         x_train=x_train.astype(int),
         y_train=y_train.astype(int),
         x_test=x_test.astype(int),
         y_test=y_test.astype(int)
         )        
    
    
