import numpy as np
from configuration import conf
from keras.utils import np_utils
import keras.datasets as Datasets


class Load_data(object):

    def __init__(self):
        pass

    def load(self):
        task_type = conf.task_type
        dataset_name = conf.dataset_name
        assert task_type in ['Baseline', 'Sequential_split', 'Sequential_permute'], 'Task type is not valid'
        assert dataset_name in ['mnist', 'cifar10', 'cifar100', 'omniglot', 'timh',
                                'fashion_mnist'], 'Dataset is not available'

        data = self.load_dataset(dataset_name, conf.is_conv)

        if task_type == 'Baseline':
            return data
        elif task_type == 'Sequential_split':
            return self.split_data(data)
        elif task_type == 'Sequential_permute':
            return self.permute_data(data)

    def load_dataset(self, data_name, is_conv):
        if data_name in ['omniglot']:
            from utils.data_utils import load_omniglot
            (X_train, y_train), (X_test, y_test) = load_omniglot()
        elif data_name in ['timh']:
            from utils.data_utils import load_tihm
            (X_train, y_train), (X_test, y_test) = load_tihm()
        else:
            dataset_obj = getattr(Datasets, data_name)
            (X_train, y_train), (X_test, y_test) = dataset_obj.load_data()
        self.nb_classes = len(np.unique(y_train))
        conf.num_classes = self.nb_classes
        if data_name in ['cifar10', 'cifar100']:
            is_conv = True
        if not is_conv:
            # assert data_name not in ['cifar10','cifar100'], data_name + ' must be trained with is_conv = True'
            X_train = X_train.reshape(X_train.shape[0], -1)
            X_test = X_test.reshape(X_test.shape[0], -1)
        else:
            if data_name in ['mnist']:
                X_train = X_train.reshape(X_train.shape[0], 28, 28, 1)
                X_test = X_test.reshape(X_test.shape[0], 28, 28, 1)

            elif data_name in ['cifar10', 'cifar100']:
                X_train = X_train.reshape(X_train.shape[0], 32, 32, 3)
                X_test = X_test.reshape(X_test.shape[0], 32, 32, 3)

            elif data_name in ['omniglot']:
                X_train = X_train.reshape(X_train.shape[0], 105, 105, 1)
                X_test = X_test.reshape(X_test.shape[0], 105, 105, 1)

        if np.max(X_train) == 255.:
            print('Normalizing the training data ... ')
            X_train = X_train.astype('float32') / 255
            X_test = X_test.astype('float32') / 255

        data = {
            'X_train': X_train,
            'y_train': y_train,
            'X_test': X_test,
            'y_test': y_test,
        }
        return data

    def split_data(self, data):
        try:
            task_labels = conf.task_labels
        except:
            raise ValueError('Label is not provided ...')
        datasets = {}
        one_hot = conf.enable_one_hot
        for task_idx, labels in enumerate(task_labels):
            datasets[task_idx] = {}
            train_idx = np.in1d(data['y_train'], labels)
            datasets[task_idx]['X_train'] = data['X_train'][train_idx]

            test_idx = np.in1d(data['y_test'], labels)
            datasets[task_idx]['X_test'] = data['X_test'][test_idx]
            if conf.multi_head:
                if one_hot:
                    datasets[task_idx]['y_train'] = np_utils.to_categorical(data['y_train'][train_idx] - np.min(labels),
                                                                            len(labels))
                    datasets[task_idx]['y_test'] = np_utils.to_categorical(data['y_test'][test_idx] - np.min(labels),
                                                                           len(labels))
                else:
                    datasets[task_idx]['y_train'] = data['y_train'][train_idx] - np.min(labels)
                    datasets[task_idx]['y_test'] = data['y_test'][test_idx] - np.min(labels)
            else:
                datasets[task_idx]['y_train'] = np_utils.to_categorical(data['y_train'][train_idx],
                                                                        int(self.nb_classes)) if one_hot else \
                data['y_train'][train_idx]
                datasets[task_idx]['y_test'] = np_utils.to_categorical(data['y_test'][test_idx],
                                                                       int(self.nb_classes)) if one_hot else \
                data['y_test'][test_idx]

        # idx = np.in1d(data['y_test'],labels)
        # datasets[task_idx]['X_test'] = data['X_test'][idx]
        # datasets[task_idx]['y_test'] = np_utils.to_categorical(data['y_test'][idx], int(self.nb_classes)) if one_hot else data['y_test'][idx]

        return datasets

    def permute_data(self, data):
        num_tasks = conf.num_tasks
        permutations = []
        for i in range(num_tasks):
            idx = np.arange(data['X_train'].shape[1], dtype=int)
            if i > 0:
                np.random.shuffle(idx)
            permutations.append(idx)
        datasets = {}
        one_hot = conf.enable_one_hot
        for task_idx, perm in enumerate(permutations):
            datasets[task_idx] = {}

            datasets[task_idx]['X_train'] = data['X_train'][:, perm]
            datasets[task_idx]['X_test'] = data['X_test'][:, perm]
            datasets[task_idx]['y_train'] = np_utils.to_categorical(data['y_train'],
                                                                    int(self.nb_classes)) if one_hot else data[
                'y_train']
            datasets[task_idx]['y_test'] = np_utils.to_categorical(data['y_test'], int(self.nb_classes)) if one_hot else \
            data['y_test']

        return datasets

    def get_description(self, data):
        X_train, y_train, X_test, y_test = data

        print('X_train : ', X_train.shape)
        print('X_test : ', X_test.shape)

        print('y_train : ', np.unique(y_train))
        for l in np.unique(y_train):
            print('y_train == ', l, y_train[y_train == l].shape)
            print('y_test == ', l, y_test[y_test == l].shape)
