# The CIFAR-10 dataset consists of 60000 32x32 colour images in 10 classes, with 6000 images per class. There are 50000 training images and 10000 test images.

# The dataset is divided into five training batches and one test batch, each with 10000 images. The test batch contains exactly 1000 randomly-selected images from each class. The training batches contain the remaining images in random order, but some training batches may contain more images from one class than another. Between them, the training batches contain exactly 5000 images from each class.

import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf

CIFAR_DIR = 'cifar-10-batches-py/'

# Unpickle and store all the data in a list


def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        cifar_dict = pickle.load(fo, encoding='bytes')
    return cifar_dict


dirs = ['batches.meta', 'data_batch_1', 'data_batch_2',
        'data_batch_3', 'data_batch_4', 'data_batch_5', 'test_batch']
all_data = [0, 1, 2, 3, 4, 5, 6]
for i, direc in zip(all_data, dirs):
    all_data[i] = unpickle(CIFAR_DIR+direc)

batch_meta = all_data[0]
data_batch1 = all_data[1]
data_batch2 = all_data[2]
data_batch3 = all_data[3]
data_batch4 = all_data[4]
data_batch5 = all_data[5]
test_batch = all_data[6]
print(batch_meta)

# Loaded in this way, each of the batch files contains a dictionary with the following elements:

# data - - a 10000x3072 numpy array of uint8s. Each row of the array stores a 32x32 colour image. The first 1024 entries contain the red channel values, the next 1024 the green, and the final 1024 the blue. The image is stored in row-major order, so that the first 32 entries of the array are the red channel values of the first row of the image.
# labels - - a list of 10000 numbers in the range 0-9. The number at index i indicates the label of the ith image in the array data.

# Show the images

X = data_batch1[b"data"]
# Turn into an array of 32x32 pixels and 3 color channels
# reshape 10000 images, 3 color channels, 32x32
# transpose the images
X = X.reshape(10000, 3, 32, 32).transpose(0, 2, 3, 1).astype("uint8")
plt.imshow(X[34])
plt.show()

# HELPER FNS

# Create our one-hot encoder


def one_hot_encode(vector, vals=10):
    # 10 possible labels
    out = np.zeros((len(vector), vals))
    out[range(len(vector)), vector] = 1
    return out


# Helper functions from mnist project

def init_weights(shape):
    init_random_dist = tf.truncated_normal(shape, stddev=0.1)
    return tf.Variable(init_random_dist)


def init_bias(shape):
    init_bias_vals = tf.constant(0.1, shape=shape)
    return tf.Variable(init_bias_vals)


def conv2d(x, W):
    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding='SAME')


def max_pool_2by2(x):
    return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')


def convolutional_layer(input_x, shape):
    W = init_weights(shape)
    b = init_bias([shape[3]])
    return tf.nn.relu(conv2d(input_x, W) + b)


def normal_full_layer(input_layer, size):
    input_size = int(input_layer.get_shape()[1])
    W = init_weights([input_size, size])
    b = init_bias([size])
    return tf.matmul(input_layer, W) + b


class CifarHelper():
    def __init__(self):
        self.i = 0

        self.all_train_batches = [
            data_batch1, data_batch2, data_batch3, data_batch4, data_batch5]
        self.test_batch = [test_batch]

        self.training_images = None
        self.training_labels = None

        self.test_images = None
        self.test_labels = None

    def set_up_images(self):
        print("Setting Up Training Images and Labels")
        # Vertically stack the training images
        self.training_images = np.vstack(
            [d[b"data"] for d in self.all_train_batches])
        train_len = len(self.training_images)

        # Reshape and normalize training images`
        self.training_images = self.training_images.reshape(
            train_len, 3, 32, 32).transpose(0, 2, 3, 1)/255
        self.training_labels = one_hot_encode(
            np.hstack([d[b"labels"] for d in self.all_train_batches]), 10)

        print("Setting Up Test Images and Labels")

        # Vertically stack test images
        self.test_images = np.vstack([d[b"data"] for d in self.test_batch])
        test_len = len(self.test_images)

        # Reshape and normalize test images
        self.test_images = self.test_images.reshape(
            test_len, 3, 32, 32).transpose(0, 2, 3, 1)/255
        self.test_labels = one_hot_encode(
            np.hstack([d[b"labels"] for d in self.test_batch]), 10)

    def next_batch(self, batch_size=100):
        # The 100 dimension in the reshape call is set by an assumed batch size of 100
        x = self.training_images[self.i:self.i +
                                 batch_size].reshape(batch_size, 32, 32, 3)
        y = self.training_labels[self.i:self.i+batch_size]
        self.i = (self.i + batch_size) % len(self.training_images)
        return x, y


ch = CifarHelper()
ch.set_up_images()

batch = ch.next_batch(100)

# Creating the model

x = tf.placeholder(tf.float32, shape=[None, 32, 32, 3])
y_true = tf.placeholder(tf.float32, shape=[None, 10])
hold_prob = tf.placeholder(tf.float32)

# Layers
convo_1 = convolutional_layer(x, shape=[4, 4, 3, 32])
convo_1_pooling = max_pool_2by2(convo_1)

convo_2 = convolutional_layer(convo_1_pooling, shape=[4, 4, 32, 64])
convo_2_pooling = max_pool_2by2(convo_2)

# Create a flattened layer by reshaping the pool into [-1, 8*8*64]
convo_2_flat = tf.reshape(convo_2_pooling, [-1, 8*8*64])

# New full layer
full_layer_one = tf.nn.relu(normal_full_layer(convo_2_flat, 1024))

# Create the dropout layer
full_one_dropout = tf.nn.dropout(full_layer_one, keep_prob=hold_prob)

# Pass dropout layer to our pred
y_pred = normal_full_layer(full_one_dropout, 10)

# Loss fn
cross_entropy = tf.reduce_mean(
    tf.nn.softmax_cross_entropy_with_logits(labels=y_true, logits=y_pred))

# Optimizer
optimizer = tf.train.AdamOptimizer(learning_rate=0.001)
train = optimizer.minimize(cross_entropy)

init = tf.global_variables_initializer()

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())

    for i in range(5000):
        batch = ch.next_batch(100)
        sess.run(train, feed_dict={
                 x: batch[0], y_true: batch[1], hold_prob: 0.5})

        if i % 100 == 0:
            print('Currently on step {}'.format(i))
            print('Accuracy is:')
            # Test the Train Model
            matches = tf.equal(tf.argmax(y_pred, 1), tf.argmax(y_true, 1))

            acc = tf.reduce_mean(tf.cast(matches, tf.float32))

            print(sess.run(acc, feed_dict={
                  x: ch.test_images, y_true: ch.test_labels, hold_prob: 1.0}))
            print('\n')
