import h5py
import numpy as np
import tensorflow as tf
    
# this package
from astronn.data import fetch_notMNIST

"""
Explanation: Deep Learning
Assignment 2 -- Solutions
Previously in 1_notmnist.ipynb, we
1. Downloaded a test dataset with training, development, and testing subsets (based on the notMNIST dataset).
2. We visualized and explored the data.
3. We trained an out-of-box logistic regression model to classify the data and found that it performs very badly
The goal of this assignment is to progressively train deeper and more accurate models using TensorFlow.
End of explanation
"""


cache_file = fetch_notMNIST()

def randomize(data, labels):
    permutation = np.random.permutation(labels.shape[0])
    shuffled_data = data[permutation]
    shuffled_labels = labels[permutation]
    return shuffled_data, shuffled_labels

with h5py.File(cache_file, 'r') as f:
    train_dataset, train_labels = randomize(f['train']['images'][:], f['train']['labels'][:])
    valid_dataset, valid_labels = randomize(f['validate']['images'][:], f['validate']['labels'][:])
    test_dataset, test_labels = randomize(f['test']['images'][:], f['test']['labels'][:])
    
    print('Training set', train_dataset.shape, train_labels.shape)
    print('Validation set', valid_dataset.shape, valid_labels.shape)
    print('Test set', test_dataset.shape, test_labels.shape)
    
image_size = train_dataset.shape[-1]
num_labels = np.unique(train_labels).size

"""
Explanation: First load and randomize the data, as we did in 1_notmnist.ipynb
End of explanation
"""


def labels_to_1hot(labels):
    # this uses numpy array broadcasting to do this operation in a memory efficient way:
    #    see: http://docs.scipy.org/doc/numpy-1.10.1/user/basics.broadcasting.html
    one_hots = (np.arange(num_labels) == labels[:,np.newaxis]).astype(np.float32)
    return one_hots

def flatten_images(image_data):
    # -1 in the reshape tells numpy to infer the length along that axis
    # float32 takes up less memory
    return image_data.reshape((-1, image_size*image_size)).astype(np.float32)

train_labels = labels_to_1hot(train_labels)
valid_labels = labels_to_1hot(valid_labels)
test_labels = labels_to_1hot(test_labels)

train_dataset = flatten_images(train_dataset)
valid_dataset = flatten_images(valid_dataset)
test_dataset = flatten_images(test_dataset)

print('Training set', train_dataset.shape, train_labels.shape)
print('Validation set', valid_dataset.shape, valid_labels.shape)
print('Test set', test_dataset.shape, test_labels.shape)

"""
Explanation: Here we're going to reformat the data into a shape that's more adapted to the models we're going to train:
- we'll flatten the image data (so, instead of 28 by 28 images, length-784 arrays)
- we'll turn the labels into float 1-hot encodings (read about 1-hots in this excellent StackOverflow answer).
End of explanation
"""


# With gradient descent training, even this much data is prohibitive.
# Subset the training data for faster turnaround.
train_subset = 10000

# This is where we set up the graph and define variables that we will use,
# however *nothing actually gets executed here*! This only establishes and 
# defines variables and operations on the graph.
graph = tf.Graph()
with graph.as_default():

    # Input data.
    # Load the training, validation and test data into constants that are
    # attached to the graph.
    tf_train_dataset = tf.constant(train_dataset[:train_subset, :])
    tf_train_labels = tf.constant(train_labels[:train_subset])
    tf_valid_dataset = tf.constant(valid_dataset)
    tf_test_dataset = tf.constant(test_dataset)
    
    # Variables.
    # These are the parameters that we are going to be training. The weight
    # matrix will be initialized using random values following a (truncated)
    # normal distribution. The biases get initialized to zero.
    weights = tf.Variable(tf.truncated_normal([image_size * image_size, num_labels]))
    biases = tf.Variable(tf.zeros([num_labels]))
    
    # Training computation.
    # We multiply the inputs with the weight matrix, and add biases. We compute
    # the softmax and cross-entropy (it's one operation in TensorFlow, because
    # it's very common, and it can be optimized). We take the average of this
    # cross-entropy across all training examples: that's our loss.
    logits = tf.matmul(tf_train_dataset, weights) + biases
    loss = tf.reduce_mean(
      tf.nn.softmax_cross_entropy_with_logits(logits, tf_train_labels))
    
    # Optimizer.
    # We are going to find the minimum of this loss using gradient descent.
    optimizer = tf.train.GradientDescentOptimizer(0.5).minimize(loss) # 0.5 is the learning rate
    
    # Predictions for the training, validation, and test data.
    # These are not part of training, but merely here so that we can report
    # accuracy figures as we train.
    train_prediction = tf.nn.softmax(logits)
    valid_prediction = tf.nn.softmax(
      tf.matmul(tf_valid_dataset, weights) + biases)
    test_prediction = tf.nn.softmax(tf.matmul(tf_test_dataset, weights) + biases)

"""
Explanation: We're first going to train a multinomial logistic regression using simple gradient descent.
TensorFlow works like this:
* First you describe the computation that you want to see performed: what the inputs, the variables, and the operations look like. These get created as nodes over a computation graph. This description is all contained within the block below:
  with graph.as_default():
      ...



Then you can run the operations on this graph as many times as you want by calling session.run(), providing it outputs to fetch from the graph that get returned. This runtime operation is all contained in the block below:
with tf.Session(graph=graph) as session:
      ...


Let's load all the data into TensorFlow and build the computation graph corresponding to our training:
End of explanation
"""


num_steps = 801

def accuracy(predictions, labels):
    # a simple comparision of the predicted labels vs. the truth to get the prediction accuracy
    return (100.0 * np.sum(np.argmax(predictions, axis=1) == np.argmax(labels, axis=1))
            / predictions.shape[0])

# define a Session with the graph we defined above
with tf.Session(graph=graph) as session:
    # This is a one-time operation which ensures the parameters get initialized as
    # we described in the graph: random weights for the matrix, zeros for the biases. 
    tf.initialize_all_variables().run()
    print('Initialized')
    
    for step in range(num_steps):
        # Run the computations. We tell run() that we want to run the optimizer,
        # and get the loss value and the training predictions returned as numpy arrays.
        # these three variables (optimizer, loss, and train_prediction) are defined and
        # attached to the graph in the previous cell.
        _, l, predictions = session.run([optimizer, loss, train_prediction])
        
        if (step % 100 == 0):
            print('Loss at step', step, ':', l)
            print('Training accuracy: {:.1f}%'.format(accuracy(predictions, train_labels[:train_subset, :])))
            
            # Calling .eval() on valid_prediction is basically like calling run(), but
            # just to get that one numpy array. Note that it recomputes all its graph
            # dependencies.
            print('Validation accuracy: {:.1f}%'.format(accuracy(valid_prediction.eval(), valid_labels)))
            
    print('Test accuracy: {:.1f}%'.format(accuracy(test_prediction.eval(), test_labels)))

"""
Explanation: Now that we have set up the graph, we will run some computations in a Session and iterate:
End of explanation
"""


batch_size = 128

graph = tf.Graph()
with graph.as_default():

    # Input data. For the training data, we use a placeholder that will be fed
    # at run time with a training minibatch.
    tf_train_dataset = tf.placeholder(tf.float32,
                                      shape=(batch_size, image_size * image_size))
    tf_train_labels = tf.placeholder(tf.float32, shape=(batch_size, num_labels))
    tf_valid_dataset = tf.constant(valid_dataset)
    tf_test_dataset = tf.constant(test_dataset)
    
    # Variables.
    weights = tf.Variable(
      tf.truncated_normal([image_size * image_size, num_labels]))
    biases = tf.Variable(tf.zeros([num_labels]))
    
    # Training computation.
    logits = tf.matmul(tf_train_dataset, weights) + biases
    loss = tf.reduce_mean(
      tf.nn.softmax_cross_entropy_with_logits(logits, tf_train_labels))
    
    # Optimizer.
    optimizer = tf.train.GradientDescentOptimizer(0.5).minimize(loss)
    
    # Predictions for the training, validation, and test data.
    train_prediction = tf.nn.softmax(logits)
    valid_prediction = tf.nn.softmax(
      tf.matmul(tf_valid_dataset, weights) + biases)
    test_prediction = tf.nn.softmax(tf.matmul(tf_test_dataset, weights) + biases)

"""
Explanation: The graph above holds the entire training set in memory as a constant() node on the graph. Below, we will slightly modify the graph to instead constain placeholder() nodes for the training data, and during the training we will pass in smaller (mini)batches of the training data. These placeholders will be fed the actual training data at every call of sesion.run().
End of explanation
"""


num_steps = 3001

with tf.Session(graph=graph) as session:
    tf.initialize_all_variables().run()
    print("Initialized")
    for step in range(num_steps):
        # Pick an offset within the training data, which has been randomized.
        # Note: we could use better randomization across epochs.
        offset = (step * batch_size) % (train_labels.shape[0] - batch_size)
        
        # Generate a minibatch.
        batch_data = train_dataset[offset:(offset + batch_size), :]
        batch_labels = train_labels[offset:(offset + batch_size), :]
        
        # Prepare a dictionary telling the session where to feed the minibatch.
        # The key of the dictionary is the placeholder node of the graph to be fed,
        # and the value is the numpy array to feed to it.
        feed_dict = {tf_train_dataset : batch_data, tf_train_labels : batch_labels}
        _, l, predictions = session.run([optimizer, loss, train_prediction], feed_dict=feed_dict)
        
        if (step % 500 == 0):
            print("Minibatch loss at step", step, ":", l)
            print("Minibatch accuracy: {:.1f}%".format(accuracy(predictions, batch_labels)))
            print("Validation accuracy: {:.1f}%".format(accuracy(valid_prediction.eval(), valid_labels)))
    print("Test accuracy: {:.1f}%".format(accuracy(test_prediction.eval(), test_labels)))

"""
Explanation: Let's run it:
End of explanation
"""


n_hidden = 1024
batch_size = 128

def model(X, weights, biases):
    # Hidden layer with RELU activation
    layer_1 = tf.nn.relu(tf.add(tf.matmul(X, weights['layer1']), biases['layer1'])) 
    logits = tf.add(tf.matmul(layer_1, weights['out']), biases['out'])
    return logits

graph = tf.Graph()
with graph.as_default():

    # Input data. For the training data, we use a placeholder that will be fed
    # at run time with a training minibatch.
    tf_train_dataset = tf.placeholder(tf.float32,
                                      shape=(batch_size, image_size * image_size))
    tf_train_labels = tf.placeholder(tf.float32, shape=(batch_size, num_labels))
    tf_valid_dataset = tf.constant(valid_dataset)
    tf_test_dataset = tf.constant(test_dataset)
    
    # Variables.
    # Similar to the above SGD model, we will need weights and biases, but now we
    # are building a neural network with one hidden layer. Meaning, we will need 
    # two sets of weights and biases: one set for the hidden layer, one set for 
    # the output layer. 
    weights = {
        "layer1": tf.Variable(tf.truncated_normal([image_size * image_size, n_hidden])),
        "out": tf.Variable(tf.truncated_normal([n_hidden, num_labels]))
    }
    biases = {
        "layer1": tf.Variable(tf.zeros([n_hidden])),
        "out": tf.Variable(tf.zeros([num_labels]))
    }
    
    # Training computation.
    logits = model(tf_train_dataset, weights, biases)
    loss = tf.reduce_mean(
      tf.nn.softmax_cross_entropy_with_logits(logits, tf_train_labels))

    # Optimizer.
    optimizer = tf.train.GradientDescentOptimizer(0.5).minimize(loss)
    
    # Predictions for the training, validation, and test data.
    train_prediction = tf.nn.softmax(logits)
    valid_prediction = model(tf_valid_dataset, weights, biases)
    test_prediction = model(tf_test_dataset, weights, biases) 

"""
Explanation: Problem 1
Turn the logistic regression example with SGD into a 1-hidden layer neural network with rectified linear units (nn.relu()) and 1024 hidden nodes. This model should improve the validation / test accuracy compared to the SGD model above.
End of explanation
"""


num_steps = 1001

with tf.Session(graph=graph) as session:
    tf.initialize_all_variables().run()
    print("Initialized")
    
    for step in range(num_steps):
        # Pick an offset within the training data, which has been randomized.
        # Note: we could use better randomization across epochs.
        offset = (step * batch_size) % (train_labels.shape[0] - batch_size)
        
        # Generate a minibatch.
        batch_data = train_dataset[offset:(offset + batch_size), :]
        batch_labels = train_labels[offset:(offset + batch_size), :]
        
        # Prepare a dictionary telling the session where to feed the minibatch.
        # The key of the dictionary is the placeholder node of the graph to be fed,
        # and the value is the numpy array to feed to it.
        feed_dict = {tf_train_dataset : batch_data, tf_train_labels : batch_labels}
        _, l, predictions = session.run(
          [optimizer, loss, train_prediction], feed_dict=feed_dict)
        
        if (step % 100 == 0):
            print("Minibatch loss at step", step, ":", l)
            print("Minibatch accuracy: {:.1f}%".format(accuracy(predictions, batch_labels)))
            print("Validation accuracy: {:.1f}%".format(accuracy(valid_prediction.eval(), valid_labels)))
    print("Test accuracy: {:.1f}%".format(accuracy(test_prediction.eval(), test_labels)))

"""
Explanation: Note: you should be able to run the graph exactly the same way as the SGD model above:
End of explanation
"""


