# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/data_generator.ipynb (unless otherwise specified).

__all__ = ['MNISTDataGenerator']

# Cell
from fastcore.test import *
from nbdev.showdoc import *
import numpy as np
import jovsatools
import tensorflow as tf

# Cell

class MNISTDataGenerator(object):
    """Generates a multi-calss regression/classification data based on MNIST

        x = mnist data flattened
        y = regression/classification targets based on self.initialize()

        Arguments:
            additional_y: Integer, used to specify the additional targets
            seed = Integer, used to specify np.random.seed()

        Retruns:
            train_x: numpy.ndarray, MNIST train data flattened
            train_y: numpy.ndarray, all the targets; MNIST label + additional_y
            test_x: numpy.ndarray, MNIST test data flattened
            test_y: numpy.ndarray, all the targets; MNIST label + additional_y
    """
    def __init__(self, additional_y=0, seed=1123):
        self.additional_y = additional_y
        self.seed = np.random.seed(seed=seed)

        # needed to store additional generation functions
        self._func_map = {}
        self._initialize(additional_y)

        assert list(filter(lambda x: x is None, [
            self.mnist_train_x,
            self.mnist_train_y,
            self.mnist_test_x,
            self.mnist_test_y
        ])) == []

        self.train_n = len(self.mnist_train_x)
        self.test_n = len(self.mnist_test_x)

    def _initialize(self, additional_y):
        """prepare functions to approximate """
        epsilon = 0.000123
        C = 102 # emperical value from analyzing MNIST data

        if additional_y >= 1:
            # classification (0, 1) target
            self._func_map[0] = lambda x: int(np.random.random()<0.9) if 2*(x[10]**3)-2*x[3]+15 > 8.9 else 0

        if additional_y >= 2:
            # unbounded regression target
            self._func_map[1] = lambda x: np.log(np.sum(x, axis=0)+epsilon)

        if additional_y >= 3:
            # bounded (0, 1) regression target
            self._func_map[2] = lambda x: np.mean(x, axis=0)/C

        # regular MNIST train/test data
        train, test = tf.keras.datasets.mnist.load_data()
        self.mnist_train_x, self.mnist_train_y = train[0], train[1]
        self.mnist_test_x, self.mnist_test_y = test[0], test[1]

    def __call__(self, train_n):
        """data generation """
        assert train_n <= self.train_n
        train_x, train_y = [], []
        test_x, test_y = [], []

        #TODO (jovsa): consolidate the two identical loop
        for i in range(train_n):
            x = self.mnist_train_x[i].flatten().reshape(-1, 1)
            y = [self.mnist_train_y[i]]
            for f in range(self.additional_y):
                y.append(self._func_map[f](x))
            train_x.append(x)
            train_y.append(y)

        # prepating test data
        for i in range(self.test_n):
            x = self.mnist_test_x[i].flatten().reshape(-1, 1)
            y = [self.mnist_test_y[i]]
            for f in range(self.additional_y):
                y.append(self._func_map[f](x))
            test_x.append(x)
            test_y.append(y)

        return (np.asarray(train_x), np.asarray(train_y), np.asarray(test_x), np.asarray(test_y))