import os

import keras
import numpy as np
from keras.utils import to_categorical


def read_labels_file(file_path: str) -> np.array:
    """Reads the labels from the csv file and returns it as an array

    :param file_path: path to the csv file containing the labels to the files in the dataset
    :return: labels as a numpy array
    """
    with open(file_path, 'r') as f:
        f.readline()
        vals = f.readlines()

    data = np.empty(shape=(len(vals), 1), dtype=np.float32)

    for val in vals:
        val_split = val.split(",")
        data[int(val_split[0])] = int(val_split[1])

    return data

def getmean(max_seq_len):
    mean = np.genfromtxt('mean.csv', delimiter=',')
    if len(mean) < max_seq_len:
        mean = np.array([mean,]*max_seq_len)
    return mean

def getmax(max_seq_len):
    max_arr = np.genfromtxt('max.csv', delimiter=',')
    if len(max_arr) < max_seq_len:
        max_arr = np.array([max_arr,]*max_seq_len)
        max_arr = max_arr.reshape(max_seq_len,102,1)
    return max_arr

def getstddev(max_seq_len):
    stddev = np.genfromtxt('mean.csv', delimiter=',')
    # f = open('stddev.csv', 'r')
    # for line in f:
    #     stddev.append(line)
    # f.close()
    if len(stddev) < max_seq_len:
        stddev = np.array([stddev, ] * max_seq_len)
    return stddev

class DataGenerator(keras.utils.Sequence):
    """Iterator can be used to iterate the dataset given for CS5242

    Attributes
    ----------
        max_seq_len : int
            max length of each sample in the dataset
        batch_size : int
            number of samples in each batch of training
        folder : str
            name of the folder containing the files, all the files are assumed to be in .npy format
        shuffle : bool
            this is to enable shuffle of the indices for each epoch
        file_ids : [int]
            array containing the names of all the files in the dataset excluding the .npy extension
        indexes : [int]
            copy of the file_ids array, but used for shuffling the indexes if shuffle enabled and used for slicing the batch
        labels : np.array(np.float32)
            contains the label for each file with indices corresponding to the file id from the file_ids or indexes list

    Usage
    -----
        train = DatasetGenerator('train/train', 'train_kaggle.csv')

    """

    def __init__(self, data_folder='', labels_file='', batch_size=32, max_seq_len=1000, shuffle=True):
        """Constructor to initialize the iterator

        :param data_folder: folder containing the data files; in our case the path to the folder containing the .npy files
        :param labels_file: path to the csv containing the labels to the training samples
        :param batch_size: number of samples in each batch
        :param max_seq_len: max len of each sample
        :param shuffle: whether to shuffle or not on each epoch
        """
        self.max_seq_len = max_seq_len
        self.batch_size = batch_size
        self.folder = data_folder
        self.shuffle = shuffle
        self.file_ids = np.array([int(file_name.split(".")[0]) for file_name in os.listdir(self.folder)])
        self.indexes = np.copy(self.file_ids)
        self.labels = read_labels_file(labels_file)
        self.on_epoch_end()



    def __len__(self) -> int:
        """Denotes the number of batches per epoch

        :return: batches per epoch
        """
        return int(np.floor(len(self.file_ids) / self.batch_size))

    def __getitem__(self, index: int) -> (np.array, np.array):
        """Returns each batch of samples

        :param index: based on the __len__ function this function is called with the index of the batch of samples to be sent
        :return: tuple of the X and the Y for the batch in the dataset
        """

        indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]

        Xin = np.empty(shape=(self.batch_size, 102, self.max_seq_len))
        X = np.empty(shape=(self.batch_size, self.max_seq_len, 102))
        y = np.zeros(shape=(self.batch_size, 1))
        # y_cat = np.zeros(shape=(self.batch_size, 2))
        mean = getmean(self.max_seq_len)
        stddev = getstddev(self.max_seq_len)
        max_arr = getmax(self.max_seq_len).reshape(self.max_seq_len, 102)

        for idx, id in enumerate(indexes):
            # get X data for each file
            data_x = np.load(os.path.join(self.folder, str(id)+".npy"))
            if len(data_x) < self.max_seq_len:
                X[idx] = (np.append(data_x, np.zeros(shape=(self.max_seq_len - len(data_x), 102)), axis=0)).reshape(self.max_seq_len, 102)
            else:
                X[idx] = data_x.reshape(self.max_seq_len, 102)
            # X[idx] = np.nan_to_num((X[idx]-mean)/stddev)
            # X[idx] = X[idx]/max_arr
            # X[idx] = X[idx].reshape(1000, 102)
            # Xin[idx] = np.transpose(X[idx], (1,0,2))
            # X[idx]= np.transpose(X[idx], (1,0))

            # add y for each file
            y[idx] = (self.labels[id])
            # y_cat += to_categorical(y)
        # X = np.squeeze(X)
        # X[idx] = X[idx].reshape(1000, 102, 1)
        # X = np.transpose(X, axes=(0,2,1,3))

        # print(y_cat)
        return X, y

    def on_epoch_end(self):
        """Called at the end of each epoch to shuffle the elements in the next batch"""
        if self.shuffle:
            np.random.shuffle(self.indexes)



if __name__ == '__main__':
    train = DataGenerator('train/train', 'train_kaggle.csv')
    # print(train[0][1].shape)
#
# def compute_mean_and_stddev():
#     sum_x = np.zeros(102)
#     sum_x2 = np.zeros(102)
#     mean = np.zeros(102)
#     stddev = np.zeros(102)
#     length = 0
#     for id in range(len(os.listdir('train/train'))):
#         data_x = np.load(os.path.join('train/train', str(id) + ".npy"))
#         length = length + len(data_x)
#         for tuple in range(len(data_x)):
#             sum_x = np.add(sum_x, data_x[tuple])
#             sum_x2 = np.add(sum_x2, np.square(tuple))
#         data_x =[]
#
#     mean = np.divide(sum_x, length)
#     stddev = np.sqrt(np.divide(sum_x2, length) - np.square(mean))
#     return mean, stddev
    # print(mean)
    # print(stddev)
    # for j in range (102):
    #     sum_x = sum_x +


import os

import keras
import numpy as np
from keras.utils import to_categorical


def read_labels_file(file_path: str) -> np.array:
    """Reads the labels from the csv file and returns it as an array

    :param file_path: path to the csv file containing the labels to the files in the dataset
    :return: labels as a numpy array
    """
    with open(file_path, 'r') as f:
        f.readline()
        vals = f.readlines()

    data = np.empty(shape=(len(vals), 1), dtype=np.float32)

    for val in vals:
        val_split = val.split(",")
        data[int(val_split[0])] = int(val_split[1])

    return data

def getmean(max_seq_len):
    mean = np.genfromtxt('mean.csv', delimiter=',')
    if len(mean) < max_seq_len:
        mean = np.array([mean,]*max_seq_len)
    return mean

def getmax(max_seq_len):
    max_arr = np.genfromtxt('max.csv', delimiter=',')
    if len(max_arr) < max_seq_len:
        max_arr = np.array([max_arr,]*max_seq_len)
        max_arr = max_arr.reshape(max_seq_len,102,1)
    return max_arr

def getstddev(max_seq_len):
    stddev = np.genfromtxt('mean.csv', delimiter=',')
    # f = open('stddev.csv', 'r')
    # for line in f:
    #     stddev.append(line)
    # f.close()
    if len(stddev) < max_seq_len:
        stddev = np.array([stddev, ] * max_seq_len)
    return stddev

class DataGenerator(keras.utils.Sequence):
    """Iterator can be used to iterate the dataset given for CS5242

    Attributes
    ----------
        max_seq_len : int
            max length of each sample in the dataset
        batch_size : int
            number of samples in each batch of training
        folder : str
            name of the folder containing the files, all the files are assumed to be in .npy format
        shuffle : bool
            this is to enable shuffle of the indices for each epoch
        file_ids : [int]
            array containing the names of all the files in the dataset excluding the .npy extension
        indexes : [int]
            copy of the file_ids array, but used for shuffling the indexes if shuffle enabled and used for slicing the batch
        labels : np.array(np.float32)
            contains the label for each file with indices corresponding to the file id from the file_ids or indexes list

    Usage
    -----
        train = DatasetGenerator('train/train', 'train_kaggle.csv')

    """

    def __init__(self, data_folder='', labels_file='', batch_size=32, max_seq_len=1000, shuffle=True):
        """Constructor to initialize the iterator

        :param data_folder: folder containing the data files; in our case the path to the folder containing the .npy files
        :param labels_file: path to the csv containing the labels to the training samples
        :param batch_size: number of samples in each batch
        :param max_seq_len: max len of each sample
        :param shuffle: whether to shuffle or not on each epoch
        """
        self.max_seq_len = max_seq_len
        self.batch_size = batch_size
        self.folder = data_folder
        self.shuffle = shuffle
        self.file_ids = np.array([int(file_name.split(".")[0]) for file_name in os.listdir(self.folder)])
        self.indexes = np.copy(self.file_ids)
        self.labels = read_labels_file(labels_file)
        self.on_epoch_end()



    def __len__(self) -> int:
        """Denotes the number of batches per epoch

        :return: batches per epoch
        """
        return int(np.floor(len(self.file_ids) / self.batch_size))

    def __getitem__(self, index: int) -> (np.array, np.array):
        """Returns each batch of samples

        :param index: based on the __len__ function this function is called with the index of the batch of samples to be sent
        :return: tuple of the X and the Y for the batch in the dataset
        """

        indexes = self.indexes[index * self.batch_size:(index + 1) * self.batch_size]

        Xin = np.empty(shape=(self.batch_size, 102, self.max_seq_len))
        X = np.empty(shape=(self.batch_size, self.max_seq_len, 102))
        y = np.zeros(shape=(self.batch_size, 1))
        # y_cat = np.zeros(shape=(self.batch_size, 2))
        mean = getmean(self.max_seq_len)
        stddev = getstddev(self.max_seq_len)
        max_arr = getmax(self.max_seq_len).reshape(self.max_seq_len, 102)

        for idx, id in enumerate(indexes):
            # get X data for each file
            data_x = np.load(os.path.join(self.folder, str(id)+".npy"))
            if len(data_x) < self.max_seq_len:
                X[idx] = (np.append(data_x, np.zeros(shape=(self.max_seq_len - len(data_x), 102)), axis=0)).reshape(self.max_seq_len, 102)
            else:
                X[idx] = data_x.reshape(self.max_seq_len, 102)
            # X[idx] = np.nan_to_num((X[idx]-mean)/stddev)
            # X[idx] = X[idx]/max_arr
            # X[idx] = X[idx].reshape(1000, 102)
            # Xin[idx] = np.transpose(X[idx], (1,0,2))
            # X[idx]= np.transpose(X[idx], (1,0))

            # add y for each file
            y[idx] = (self.labels[id])
            # y_cat += to_categorical(y)
        # X = np.squeeze(X)
        # X[idx] = X[idx].reshape(1000, 102, 1)
        # X = np.transpose(X, axes=(0,2,1,3))

        # print(y_cat)
        return X, y

    def on_epoch_end(self):
        """Called at the end of each epoch to shuffle the elements in the next batch"""
        if self.shuffle:
            np.random.shuffle(self.indexes)



if __name__ == '__main__':
    train = DataGenerator('train/train', 'train_kaggle.csv')
    # print(train[0][1].shape)
#
# def compute_mean_and_stddev():
#     sum_x = np.zeros(102)
#     sum_x2 = np.zeros(102)
#     mean = np.zeros(102)
#     stddev = np.zeros(102)
#     length = 0
#     for id in range(len(os.listdir('train/train'))):
#         data_x = np.load(os.path.join('train/train', str(id) + ".npy"))
#         length = length + len(data_x)
#         for tuple in range(len(data_x)):
#             sum_x = np.add(sum_x, data_x[tuple])
#             sum_x2 = np.add(sum_x2, np.square(tuple))
#         data_x =[]
#
#     mean = np.divide(sum_x, length)
#     stddev = np.sqrt(np.divide(sum_x2, length) - np.square(mean))
#     return mean, stddev
    # print(mean)
    # print(stddev)
    # for j in range (102):
    #     sum_x = sum_x +


Reads the labels from the csv file and returns it as an array      :param file_path: path to the csv file containing the labels to the files in the dataset     :return: labels as a numpy array # f = open('stddev.csv', 'r') # for line in f: #     stddev.append(line) # f.close() Iterator can be used to iterate the dataset given for CS5242      Attributes     ----------         max_seq_len : int             max length of each sample in the dataset         batch_size : int             number of samples in each batch of training         folder : str             name of the folder containing the files, all the files are assumed to be in .npy format         shuffle : bool             this is to enable shuffle of the indices for each epoch         file_ids : [int]             array containing the names of all the files in the dataset excluding the .npy extension         indexes : [int]             copy of the file_ids array, but used for shuffling the indexes if shuffle enabled and used for slicing the batch         labels : np.array(np.float32)             contains the label for each file with indices corresponding to the file id from the file_ids or indexes list      Usage     -----         train = DatasetGenerator('train/train', 'train_kaggle.csv') Constructor to initialize the iterator          :param data_folder: folder containing the data files; in our case the path to the folder containing the .npy files         :param labels_file: path to the csv containing the labels to the training samples         :param batch_size: number of samples in each batch         :param max_seq_len: max len of each sample         :param shuffle: whether to shuffle or not on each epoch Denotes the number of batches per epoch          :return: batches per epoch Returns each batch of samples          :param index: based on the __len__ function this function is called with the index of the batch of samples to be sent         :return: tuple of the X and the Y for the batch in the dataset # y_cat = np.zeros(shape=(self.batch_size, 2)) # get X data for each file # X[idx] = np.nan_to_num((X[idx]-mean)/stddev) # X[idx] = X[idx]/max_arr # X[idx] = X[idx].reshape(1000, 102) # Xin[idx] = np.transpose(X[idx], (1,0,2)) # X[idx]= np.transpose(X[idx], (1,0)) # add y for each file # y_cat += to_categorical(y) # X = np.squeeze(X) # X[idx] = X[idx].reshape(1000, 102, 1) # X = np.transpose(X, axes=(0,2,1,3)) # print(y_cat) Called at the end of each epoch to shuffle the elements in the next batch # print(train[0][1].shape) # # def compute_mean_and_stddev(): #     sum_x = np.zeros(102) #     sum_x2 = np.zeros(102) #     mean = np.zeros(102) #     stddev = np.zeros(102) #     length = 0 #     for id in range(len(os.listdir('train/train'))): #         data_x = np.load(os.path.join('train/train', str(id) + ".npy")) #         length = length + len(data_x) #         for tuple in range(len(data_x)): #             sum_x = np.add(sum_x, data_x[tuple]) #             sum_x2 = np.add(sum_x2, np.square(tuple)) #         data_x =[] # #     mean = np.divide(sum_x, length) #     stddev = np.sqrt(np.divide(sum_x2, length) - np.square(mean)) #     return mean, stddev # print(mean) # print(stddev) # for j in range (102): #     sum_x = sum_x +