from core.datasets.image.ImageDataset import ImageDataset
from tensorflow.keras.datasets import mnist

import numpy as np


class MnistDataset(ImageDataset):

    def __init__(self, seed):
        super().__init__(
            name='mnist',
            num_classes=10,
            num_features=784,
            sample_shape=(28, 28, 1),
            seed=seed
        )

    def _get_x_y(self):
        (x_train, y_train), (x_test, y_test) = mnist.load_data()
        x_train = x_train.reshape(x_train.shape[0], self.num_features)
        x_test = x_test.reshape(x_test.shape[0], self.num_features)

        x = np.concatenate([x_train, x_test], axis=0)
        y = np.concatenate([y_train, y_test], axis=0)

        return x, y
