import numpy as np
np.random.seed(42)
import tensorflow as tf
tf.set_random_seed(42)

"""
Explanation: Deep Convolutional Neural Network in TensorFlow
In this notebook, we convert our LeNet-5-inspired, MNIST-classifying, deep convolutional network from Keras to TensorFlow (compare them side by side) following Aymeric Damien's style.
Load dependencies
End of explanation
"""


from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

"""
Explanation: Load data
End of explanation
"""


epochs = 20
batch_size = 128
display_progress = 40 # after this many batches, output progress to screen
wt_init = tf.contrib.layers.xavier_initializer() # weight initializer

"""
Explanation: Set neural network hyperparameters
End of explanation
"""


# input layer: 
n_input = 784

# first convolutional layer: 
n_conv_1 = 32
k_conv_1 = 3 # k_size

# second convolutional layer: 
n_conv_2 = 64
k_conv_2 = 3

# max pooling layer:
pool_size = 2
mp_layer_dropout = 0.25

# dense layer: 
n_dense = 128
dense_layer_dropout = 0.5

# output layer: 
n_classes = 10

"""
Explanation: Set parameters for each layer
End of explanation
"""


x = tf.placeholder(tf.float32, [None, n_input])
y = tf.placeholder(tf.float32, [None, n_classes])

"""
Explanation: Define placeholder Tensors for inputs and labels
End of explanation
"""


# dense layer with ReLU activation:
def dense(x, W, b):
    z = tf.add(tf.matmul(x, W), b)
    a = tf.nn.relu(z)
    return a

# convolutional layer with ReLU activation:
def conv2d(x, W, b, stride_length=1):
    xW = tf.nn.conv2d(x, W, strides=[1, stride_length, stride_length, 1], padding='SAME')
    z = tf.nn.bias_add(xW, b)
    a = tf.nn.relu(z)
    return a

# max-pooling layer: 
def maxpooling2d(x, p_size):
    return tf.nn.max_pool(x, 
                          ksize=[1, p_size, p_size, 1], 
                          strides=[1, p_size, p_size, 1], 
                          padding='SAME')

"""
Explanation: Define types of layers
End of explanation
"""


def network(x, weights, biases, n_in, mp_psize, mp_dropout, dense_dropout):

    # reshape linear MNIST pixel input into square image: 
    square_dimensions = int(np.sqrt(n_in))
    square_x = tf.reshape(x, shape=[-1, square_dimensions, square_dimensions, 1])
    
    # convolutional and max-pooling layers:
    conv_1 = conv2d(square_x, weights['W_c1'], biases['b_c1'])
    conv_2 = conv2d(conv_1, weights['W_c2'], biases['b_c2'])
    pool_1 = maxpooling2d(conv_2, mp_psize)
    pool_1 = tf.nn.dropout(pool_1, 1-mp_dropout)
    
    # dense layer: 
    flat = tf.reshape(pool_1, [-1, weights['W_d1'].get_shape().as_list()[0]])
    dense_1 = dense(flat, weights['W_d1'], biases['b_d1'])
    dense_1 = tf.nn.dropout(dense_1, 1-dense_dropout)
    
    # output layer: 
    out_layer_z = tf.add(tf.matmul(dense_1, weights['W_out']), biases['b_out'])
    
    return out_layer_z

"""
Explanation: Design neural network architecture
End of explanation
"""


bias_dict = {
    'b_c1': tf.Variable(tf.zeros([n_conv_1])),
    'b_c2': tf.Variable(tf.zeros([n_conv_2])),
    'b_d1': tf.Variable(tf.zeros([n_dense])),
    'b_out': tf.Variable(tf.zeros([n_classes]))
}

# calculate number of inputs to dense layer: 
full_square_length = np.sqrt(n_input)
pooled_square_length = int(full_square_length / pool_size)
dense_inputs = pooled_square_length**2 * n_conv_2

weight_dict = {
    'W_c1': tf.get_variable('W_c1', 
                            [k_conv_1, k_conv_1, 1, n_conv_1], initializer=wt_init),
    'W_c2': tf.get_variable('W_c2', 
                            [k_conv_2, k_conv_2, n_conv_1, n_conv_2], initializer=wt_init),
    'W_d1': tf.get_variable('W_d1', 
                            [dense_inputs, n_dense], initializer=wt_init),
    'W_out': tf.get_variable('W_out', 
                             [n_dense, n_classes], initializer=wt_init)
}

"""
Explanation: Define dictionaries for storing weights and biases for each layer -- and initialize
End of explanation
"""


predictions = network(x, weight_dict, bias_dict, n_input, 
                      pool_size, mp_layer_dropout, dense_layer_dropout)

"""
Explanation: Build model
End of explanation
"""


cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=predictions, labels=y))
optimizer = tf.train.AdamOptimizer().minimize(cost)

"""
Explanation: Define model's loss and its optimizer
End of explanation
"""


correct_prediction = tf.equal(tf.argmax(predictions, 1), tf.argmax(y, 1))
accuracy_pct = tf.reduce_mean(tf.cast(correct_prediction, tf.float32)) * 100

"""
Explanation: Define evaluation metrics
End of explanation
"""


initializer_op = tf.global_variables_initializer()

"""
Explanation: Create op for variable initialization
End of explanation
"""


with tf.Session() as session:
    session.run(initializer_op)
    
    print("Training for", epochs, "epochs.")
    
    # loop over epochs: 
    for epoch in range(epochs):
        
        avg_cost = 0.0 # track cost to monitor performance during training
        avg_accuracy_pct = 0.0
        
        # loop over all batches of the epoch:
        n_batches = int(mnist.train.num_examples / batch_size)
        for i in range(n_batches):

            # to reassure you something's happening! 
            if i % display_progress == 0:
                print("Step ", i+1, " of ", n_batches, " in epoch ", epoch+1, ".", sep='')
            
            batch_x, batch_y = mnist.train.next_batch(batch_size)
            
            # feed batch data to run optimization and fetching cost and accuracy: 
            _, batch_cost, batch_acc = session.run([optimizer, cost, accuracy_pct], 
                                                   feed_dict={x: batch_x, y: batch_y})
            
            # accumulate mean loss and accuracy over epoch: 
            avg_cost += batch_cost / n_batches
            avg_accuracy_pct += batch_acc / n_batches
            
        # output logs at end of each epoch of training:
        print("Epoch ", '%03d' % (epoch+1), 
              ": cost = ", '{:.3f}'.format(avg_cost), 
              ", accuracy = ", '{:.2f}'.format(avg_accuracy_pct), "%", 
              sep='')
    
    print("Training Complete. Testing Model.\n")
    
    test_cost = cost.eval({x: mnist.test.images, y: mnist.test.labels})
    test_accuracy_pct = accuracy_pct.eval({x: mnist.test.images, y: mnist.test.labels})
    
    print("Test Cost:", '{:.3f}'.format(test_cost))
    print("Test Accuracy: ", '{:.2f}'.format(test_accuracy_pct), "%", sep='')

"""
Explanation: Train the network in a session (identical to intermediate_net_in_tensorflow.ipynb except addition of display_progress)
End of explanation
"""


