# Import libraries
import tensorflow as tf
import numpy as np
import time
import collections
import os

# Import MNIST data with TensorFlow
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets(os.path.join('datasets', 'mnist'), one_hot=True) # load data in local folder

train_data = mnist.train.images.astype(np.float32)
train_labels = mnist.train.labels

test_data = mnist.test.images.astype(np.float32)
test_labels = mnist.test.labels

print(train_data.shape)
print(train_labels.shape)
print(test_data.shape)
print(test_labels.shape)

"""
Explanation: A Network Tour of Data Science
&nbsp; &nbsp; &nbsp; Xavier Bresson, Winter 2016/17
Exercise 4 : Introduction to TensorFlow
End of explanation
"""


# computational graph inputs
batch_size = 100
d = train_data.shape[1]
nc = 10
x = tf.placeholder(tf.float32,[batch_size,d]); print('x=',x,x.get_shape())
y_label = tf.placeholder(tf.float32,[batch_size,nc]); print('y_label=',y_label,y_label.get_shape())

"""
Explanation: 1st Step: Construct Computational Graph
Question 1: Prepare the input variables (x,y_label) of the computational graph
Hint: You may use the function tf.placeholder()
End of explanation
"""


# computational graph variables
initial = tf.truncated_normal([d,nc], stddev=0.1); W = tf.Variable(initial); print('W=',W.get_shape())
b = tf.Variable(tf.zeros([nc],tf.float32)); print('b=',b.get_shape())

"""
Explanation: Question 2: Prepare the variables (W,b) of the computational graph
Hint: You may use the function tf.Variable(), tf.truncated_normal()
End of explanation
"""


# Construct CG / output value
y = tf.matmul(x, W); print('y1=',y,y.get_shape())
y += b; print('y2=',y,y.get_shape())
y = tf.nn.softmax(y); print('y3=',y,y.get_shape())

"""
Explanation: Question 3: Compute the classifier such that
$$
y=softmax(Wx +b)
$$
Hint: You may use the function tf.matmul(), tf.nn.softmax()
End of explanation
"""


# Loss
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_label * tf.log(y), 1))

"""
Explanation: Question 4: Construct the loss of the computational graph such that
$$
loss = cross\ entropy(y_{label},y) =  mean_{all\ data} \ \sum_{all\ classes} -\ y_{label}.\log(y)
$$
Hint: You may use the function tf.Variable(), tf.truncated_normal()
End of explanation
"""


reg_loss = tf.nn.l2_loss(W)
reg_loss += tf.nn.l2_loss(b)

"""
Explanation: Question 5: Construct the L2 regularization of (W,b) to the computational graph such that
$$
R(W) = \|W\|_2^2\
R(b) = \|b\|_2^2
$$
Hint: You may use the function tf.nn.l2_loss()
End of explanation
"""


reg_par = 1e-3
total_loss = cross_entropy + reg_par* reg_loss

"""
Explanation: Question 6: Form the total loss
$$
total\ loss = cross\ entropy(y_{label},y) + reg_par* (R(W) + R(b))
$$
End of explanation
"""


# Update CG variables / backward pass
train_step = tf.train.GradientDescentOptimizer(0.25).minimize(total_loss)

"""
Explanation: Question 7: Perform optimization of the total loss for learning weight variables of the computational graph
Hint: You may use the function tf.train.GradientDescentOptimizer(learning_rate).minimize(total_loss)
End of explanation
"""


# Accuracy
correct_prediction = tf.equal(tf.argmax(y,1), tf.argmax(y_label,1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

"""
Explanation: Question 8: Evaluate the accuracy
Hint: You may use the function tf.equal(tf.argmax(y,1), tf.argmax(y_label,1)) and tf.reduce_mean()
End of explanation
"""


# Create test set 
idx = np.random.permutation(test_data.shape[0]) # rand permutation
idx = idx[:batch_size]
test_x, test_y = test_data[idx,:], test_labels[idx]

n = train_data.shape[0]
indices = collections.deque()

# Running Computational Graph
init = tf.initialize_all_variables()
sess = tf.Session()
sess.run(init)
for i in range(50):
    
    # Batch extraction
    if len(indices) < batch_size:
        indices.extend(np.random.permutation(n)) # rand permutation
    idx = [indices.popleft() for i in range(batch_size)] # extract n_batch data
    batch_x, batch_y = train_data[idx,:], train_labels[idx]
    
    # Run CG for variable training
    _,acc_train,total_loss_o = sess.run([train_step,accuracy,total_loss], feed_dict={x: batch_x, y_label: batch_y})
    print('\nIteration i=',i,', train accuracy=',acc_train,', loss=',total_loss_o)
    
    # Run CG for testset
    acc_test = sess.run(accuracy, feed_dict={x: test_x, y_label: test_y})
    print('test accuracy=',acc_test)

"""
Explanation: 2nd Step: Run the Computational Graph with batches of training data
Check out the accuracy of test set
End of explanation
"""


