import numpy as np

import os

from glob import glob

from natsort import natsorted

import tensorflow as tf

from functools import partial

data_cls = natsorted(glob('../../data/b2i_data/images/train/*'))

cls2idx  = {key.split(os.path.sep)[-1]:idx for idx, key in enumerate(data_cls, start=0)}

idx2cls  = {value:key for key, value in cls2idx.items()}

# def map_func(path):

#   eeg, class_name, subject_name = np.load(path, allow_pickle=True)

#   eeg = np.transpose(eeg, [1, 0])

#   class_idx  = cls2idx[class_name]

#   return tf.cast(eeg, dtype=tf.float32), tf.cast(class_idx, dtype=tf.int32)

# def load_complete_data(data_path, batch_size=16):

#   dataset = tf.data.Dataset.list_files(data_path)

#   dataset = dataset.map(lambda X: tf.numpy_function( partial(map_func), [X], [tf.float32, tf.int32,],), num_parallel_calls=tf.data.experimental.AUTOTUNE)

#   dataset = dataset.shuffle(buffer_size=2*batch_size).batch(batch_size, drop_remainder=False).prefetch(tf.data.experimental.AUTOTUNE)

#   return dataset

def preprocess_data(X, Y):

    X = tf.squeeze(X, axis=-1)

    max_val = tf.reduce_max(X)/2.0

    X = (X - max_val) / max_val

    X = tf.transpose(X, [1, 0])

    X = tf.cast(X, dtype=tf.float32)

    Y = tf.argmax(Y)

    return X, Y

def load_complete_data(X, Y, batch_size=16):    

    dataset = tf.data.Dataset.from_tensor_slices((X, Y)).map(preprocess_data).shuffle(buffer_size=2*batch_size).batch(batch_size, drop_remainder=False).prefetch(tf.data.experimental.AUTOTUNE)

    return dataset