from abc import ABC, abstractmethod
from core.explanations.utils import group_samples_by_label
from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.utils import to_categorical

import numpy as np


class BaseDataset(ABC):

    def __init__(self, name, path, num_classes, num_features, sample_shape, seed):
        self.name = name
        self.path = path
        self.num_classes = num_classes
        self.num_features = num_features
        self.sample_shape = sample_shape
        self.seed = seed

        # Compute per-feature means from training data
        x, _ = self._get_x_y()
        self.feature_means = np.mean(x, axis=0)

    def get_data(self):
        x, y = self._get_x_y()
        x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=self.seed)

        x_train, y_train = self.__preprocess_data(x_train, y_train, self.num_classes)
        x_test, y_test = self.__preprocess_data(x_test, y_test, self.num_classes)

        if self.name == 'digits' or self.name == 'mnist':
            x_train = x_train.reshape((-1, *self.sample_shape))
            x_test = x_test.reshape((-1, *self.sample_shape))

        return (x_train, y_train), (x_test, y_test)

    def get_samples(self, label, num_samples):
        (_, _), (x_test, y_test) = self.get_data()

        samples_by_label = group_samples_by_label(x_test, y_test, self.num_classes)
        return samples_by_label[label][:num_samples]

    @abstractmethod
    def _get_x_y(self):
        pass

    @staticmethod
    def __preprocess_data(x, y, num_classes):
        # Normalize the data to [0, 1] range.
        scaler = MinMaxScaler()
        x = scaler.fit_transform(x.astype('float32'))

        # Convert labels to categorical format.
        y = to_categorical(y, num_classes=num_classes)

        return x, y

    @staticmethod
    def _save_figure(figure, path):
        figure.tight_layout()
        figure.savefig(path, bbox_inches='tight', pad_inches=0, dpi=600, format='pdf')
        plt.close(figure)
