# Path: deep_architect/contrib/misc/datasets/loaders.py
# def load_mnist(flatten=False, one_hot=True, validation_frac=0.1):
#     from keras.datasets import mnist
#     (X_train, y_train), (X_test, y_test) = mnist.load_data()
#     X_train = X_train / 255.0
#     X_test = X_test / 255.0
#     if flatten:
#         X_train = X_train.reshape((X_train.shape[0], -1))
#         X_test = X_test.reshape((X_test.shape[0], -1))
#     if one_hot:
#         y_train = au.idx_to_onehot(y_train, 10)
#         y_test = au.idx_to_onehot(y_test, 10)
#     num_train = int((1.0 - validation_frac) * X_train.shape[0])
#     X_train, X_val = X_train[:num_train], X_train[num_train:]
#     y_train, y_val = y_train[:num_train], y_train[num_train:]
#     return (X_train, y_train, X_val, y_val, X_test, y_test)
# 
# Path: deep_architect/contrib/misc/datasets/dataset.py
# class InMemoryDataset:
#     """Wrapper around a dataset for iteration that allows cycling over the
#     dataset.
#     This functionality is especially useful for training. One can specify if
#     the data is to be shuffled at the end of each epoch. It is also possible
#     to specify a transformation function to applied to the batch before
#     being returned by next_batch.
#     """
# 
#     def __init__(self, X, y, shuffle_at_epoch_begin, batch_transform_fn=None):
#         if X.shape[0] != y.shape[0]:
#             assert ValueError("X and y the same number of examples.")
# 
#         self.X = X
#         self.y = y
#         self.shuffle_at_epoch_begin = shuffle_at_epoch_begin
#         self.batch_transform_fn = batch_transform_fn
#         self.iter_i = 0
# 
#     def get_num_examples(self):
#         return self.X.shape[0]
# 
#     def next_batch(self, batch_size):
#         """Returns the next batch in the dataset.
#         If there are fewer that batch_size examples until the end
#         of the epoch, next_batch returns only as many examples as there are
#         remaining in the epoch.
#         """
# 
#         n = self.X.shape[0]
#         i = self.iter_i
# 
#         # shuffling step.
#         if i == 0 and self.shuffle_at_epoch_begin:
#             inds = np.random.permutation(n)
#             self.X = self.X[inds]
#             self.y = self.y[inds]
# 
#         # getting the batch.
#         eff_batch_size = min(batch_size, n - i)
#         X_batch = self.X[i:i + eff_batch_size]
#         y_batch = self.y[i:i + eff_batch_size]
#         self.iter_i = (self.iter_i + eff_batch_size) % n
# 
#         # transform if a transform function was defined.
#         if self.batch_transform_fn != None:
#             X_batch_out, y_batch_out = self.batch_transform_fn(X_batch, y_batch)
#         else:
#             X_batch_out, y_batch_out = X_batch, y_batch
# 
#         return (X_batch_out, y_batch_out)


import tensorflow as tf
import numpy as np
import deep_architect.modules as mo
import deep_architect.hyperparameters as hp
import deep_architect.helpers.tensorflow_support as htf
import deep_architect.visualization as vi
import deep_architect.searchers.random as se
import deep_architect.core as co
from deep_architect.contrib.misc.datasets.loaders import load_mnist
from deep_architect.contrib.misc.datasets.dataset import InMemoryDataset

                    _, c = sess.run([optimizer, loss], feed_dict=train_feed)
                    avg_loss += c / num_batches

                if self.log_output_to_terminal and epoch % self.display_step == 0:
                    print("epoch:", '%d' % (epoch + 1), "loss:",
                          "{:.9f}".format(avg_loss))

            val_acc = self.compute_accuracy(sess, X_pl, y_pl, num_correct,
                                            self.val_dataset, eval_feed)
            print("validation accuracy: %0.4f" % val_acc)
            results = {
                'validation_accuracy': val_acc,
                'num_parameters': htf.get_num_trainable_parameters()
            }

        return results


def main():
    num_classes = 10
    num_samples = 4
    num_training_epochs = 2
    # NOTE: change to True for graph visualization
    show_graph = False

    # load data
    (X_train, y_train, X_val, y_val, X_test, y_test) = load_mnist(flatten=True)
    train_dataset = InMemoryDataset(X_train, y_train, True)
    val_dataset = InMemoryDataset(X_val, y_val, False)

# Path: deep_architect/contrib/misc/datasets/loaders.py
# def load_mnist(flatten=False, one_hot=True, validation_frac=0.1):
#     from keras.datasets import mnist
#     (X_train, y_train), (X_test, y_test) = mnist.load_data()
#     X_train = X_train / 255.0
#     X_test = X_test / 255.0
#     if flatten:
#         X_train = X_train.reshape((X_train.shape[0], -1))
#         X_test = X_test.reshape((X_test.shape[0], -1))
#     if one_hot:
#         y_train = au.idx_to_onehot(y_train, 10)
#         y_test = au.idx_to_onehot(y_test, 10)
#     num_train = int((1.0 - validation_frac) * X_train.shape[0])
#     X_train, X_val = X_train[:num_train], X_train[num_train:]
#     y_train, y_val = y_train[:num_train], y_train[num_train:]
#     return (X_train, y_train, X_val, y_val, X_test, y_test)
# 
# Path: deep_architect/contrib/misc/datasets/dataset.py
# class InMemoryDataset:
#     """Wrapper around a dataset for iteration that allows cycling over the
#     dataset.
#     This functionality is especially useful for training. One can specify if
#     the data is to be shuffled at the end of each epoch. It is also possible
#     to specify a transformation function to applied to the batch before
#     being returned by next_batch.
#     """
# 
#     def __init__(self, X, y, shuffle_at_epoch_begin, batch_transform_fn=None):
#         if X.shape[0] != y.shape[0]:
#             assert ValueError("X and y the same number of examples.")
# 
#         self.X = X
#         self.y = y
#         self.shuffle_at_epoch_begin = shuffle_at_epoch_begin
#         self.batch_transform_fn = batch_transform_fn
#         self.iter_i = 0
# 
#     def get_num_examples(self):
#         return self.X.shape[0]
# 
#     def next_batch(self, batch_size):
#         """Returns the next batch in the dataset.
#         If there are fewer that batch_size examples until the end
#         of the epoch, next_batch returns only as many examples as there are
#         remaining in the epoch.
#         """
# 
#         n = self.X.shape[0]
#         i = self.iter_i
# 
#         # shuffling step.
#         if i == 0 and self.shuffle_at_epoch_begin:
#             inds = np.random.permutation(n)
#             self.X = self.X[inds]
#             self.y = self.y[inds]
# 
#         # getting the batch.
#         eff_batch_size = min(batch_size, n - i)
#         X_batch = self.X[i:i + eff_batch_size]
#         y_batch = self.y[i:i + eff_batch_size]
#         self.iter_i = (self.iter_i + eff_batch_size) % n
# 
#         # transform if a transform function was defined.
#         if self.batch_transform_fn != None:
#             X_batch_out, y_batch_out = self.batch_transform_fn(X_batch, y_batch)
#         else:
#             X_batch_out, y_batch_out = X_batch, y_batch
# 
#         return (X_batch_out, y_batch_out)


# Path: examples/framework_starters/main_tensorflow.py
import tensorflow as tf
import numpy as np
import deep_architect.modules as mo
import deep_architect.hyperparameters as hp
import deep_architect.helpers.tensorflow_support as htf
import deep_architect.visualization as vi
import deep_architect.searchers.random as se
import deep_architect.core as co
from deep_architect.contrib.misc.datasets.loaders import load_mnist
from deep_architect.contrib.misc.datasets.dataset import InMemoryDataset


                    _, c = sess.run([optimizer, loss], feed_dict=train_feed)
                    avg_loss += c / num_batches

                if self.log_output_to_terminal and epoch % self.display_step == 0:
                    print("epoch:", '%d' % (epoch + 1), "loss:",
                          "{:.9f}".format(avg_loss))

            val_acc = self.compute_accuracy(sess, X_pl, y_pl, num_correct,
                                            self.val_dataset, eval_feed)
            print("validation accuracy: %0.4f" % val_acc)
            results = {
                'validation_accuracy': val_acc,
                'num_parameters': htf.get_num_trainable_parameters()
            }

        return results


def main():
    num_classes = 10
    num_samples = 4
    num_training_epochs = 2
    # NOTE: change to True for graph visualization
    show_graph = False

    # load data
    (X_train, y_train, X_val, y_val, X_test, y_test) = load_mnist(flatten=True)
    train_dataset = InMemoryDataset(X_train, y_train, True)
    val_dataset = InMemoryDataset(X_val, y_val, False)

    test_dataset = InMemoryDataset(X_test, y_test, False)