import numpy as np
import os
import logging
from sklearn.model_selection import train_test_split

DATASET_ROOT_FOLDER = os.path.abspath('datasets')

class DataLoader:

    train = None
    validation = None
    test = None
    mode = None
    partial_dataset = None

    @staticmethod
    def load(train_path=None, validation_path=None, test_path=None, height=28, length=28, train_validation_split_point=10000):
        if train_path is not None:
            DataLoader.train = DataLoader.load_image_data_with_label_at_end(
                os.path.join(DATASET_ROOT_FOLDER, train_path), height=height, length=length)
        if validation_path is not None:
            DataLoader.validation = DataLoader.load_image_data_with_label_at_end(
                os.path.join(DATASET_ROOT_FOLDER, validation_path), height=height, length=length)
        elif train_validation_split_point is not None and train_validation_split_point > 0:
            if DataLoader.mode is None or DataLoader.partial_dataset is not None:
                train_validation_split_point = int(DataLoader.train['images'].shape[0] * 0.8)
            splited_train = {
                'images': DataLoader.train['images'][0:train_validation_split_point, :, :, :],
                'labels': DataLoader.train['labels'][0:train_validation_split_point]
            }
            splited_validation = {
                'images': DataLoader.train['images'][train_validation_split_point:, :, :, :],
                'labels': DataLoader.train['labels'][train_validation_split_point:]
            }
            DataLoader.train = splited_train
            DataLoader.validation = splited_validation
        if test_path is not None:
            DataLoader.test = DataLoader.load_image_data_with_label_at_end(os.path.join(DATASET_ROOT_FOLDER, test_path), height=height, length=length)

        logging.debug('Training data shape:{}'.format(str(DataLoader.train['images'].shape)))
        logging.debug('Validation data shape:{}'.format(str(DataLoader.validation['images'].shape)))
        logging.debug('Test data shape:{}'.format(str(DataLoader.test['images'].shape)))
        return DataLoader

    @staticmethod
    def get_training_data():
        """
        get training data
        :return: dict of (images, labels)
        :rtype: dict
        """
        images = DataLoader.train.images
        labels = DataLoader.train.labels

        return {
            'images': images,
            'labels': labels
        }

    @staticmethod
    def get_validation_data():
        """
        get validation data
        :return: dict of (images, labels)
        :rtype: dict
        """
        images = DataLoader.validation.images
        labels = DataLoader.validation.labels

        return {
            'images': images,
            'labels': labels
        }

    @staticmethod
    def get_test_data():
        """
            get test data
            :return: dict of (images, labels)
            :rtype: dict
            """
        images = DataLoader.test.images
        labels = DataLoader.test.labels

        return {
            'images': images,
            'labels': labels
        }

    @staticmethod
    def load_image_data_with_label_at_end(path, height, length):
        data = np.loadtxt(path)
        if DataLoader.mode is None:
            data = data[0:1000, :]
        elif DataLoader.partial_dataset is not None and DataLoader.partial_dataset > 0 and DataLoader.partial_dataset <1:
            # randomly pick partial dataset
            cut_point = int(data.shape[0] * DataLoader.partial_dataset)
            indices = np.random.permutation(data.shape[0])
            training_idx= indices[:cut_point]
            data = data[training_idx, :]
        images = data[:, 0:-1]
        labels = data[:, -1]
        images = np.reshape(images, [images.shape[0], height, length, 1], order='F')

        return {
            'images': images,
            'labels': labels
        }


[((6, 22, 6, 49), 'os.path.abspath', 'os.path.abspath', ({(6, 38, 6, 48): '"""datasets"""'}, {}), "('datasets')", False, 'import os\n'), ((92, 15, 92, 31), 'numpy.loadtxt', 'np.loadtxt', ({(92, 26, 92, 30): 'path'}, {}), '(path)', True, 'import numpy as np\n'), ((103, 17, 103, 84), 'numpy.reshape', 'np.reshape', (), '', True, 'import numpy as np\n'), ((20, 16, 20, 61), 'os.path.join', 'os.path.join', ({(20, 29, 20, 48): 'DATASET_ROOT_FOLDER', (20, 50, 20, 60): 'train_path'}, {}), '(DATASET_ROOT_FOLDER, train_path)', False, 'import os\n'), ((23, 16, 23, 66), 'os.path.join', 'os.path.join', ({(23, 29, 23, 48): 'DATASET_ROOT_FOLDER', (23, 50, 23, 65): 'validation_path'}, {}), '(DATASET_ROOT_FOLDER, validation_path)', False, 'import os\n'), ((38, 75, 38, 119), 'os.path.join', 'os.path.join', ({(38, 88, 38, 107): 'DATASET_ROOT_FOLDER', (38, 109, 38, 118): 'test_path'}, {}), '(DATASET_ROOT_FOLDER, test_path)', False, 'import os\n'), ((98, 22, 98, 58), 'numpy.random.permutation', 'np.random.permutation', ({(98, 44, 98, 57): 'data.shape[0]'}, {}), '(data.shape[0])', True, 'import numpy as np\n')]