import tensorflow as tf
import tensorflowvisu
import math
from tensorflow.examples.tutorials.mnist import input_data as mnist_data
print("Tensorflow version " + tf.__version__)
tf.set_random_seed(0)

# neural network with 6 layers of relu neurons and final softmax layer
#
# · · · · · · · · · ·       (input data, flattened pixels)       X [batch, 784]        # 784 = 28 * 28
# \x/x\x/x\x/x\x/x\x/    -- fully connected layer (softmax)      W [784, 10]     b[10]
#   · · · · · · · ·                                              Y [batch, 10]

# The model is:
#
# Y = softmax( X * W + b)
#              X: matrix for 100 grayscale images of 28x28 pixels, flattened (there are 100 images in a mini-batch)
#              W: weight matrix with 784 lines and 10 columns
#              b: bias vector with 10 dimensions
#              +: add with broadcasting: adds the vector to each line of the matrix (numpy)
#              softmax(matrix) applies softmax on each line
#              softmax(line) applies an exp to each value then divides by the norm of the resulting line
#              Y: output matrix with 100 lines and 10 columns

# Download images and labels into mnist.test (10K images+labels) and mnist.train (60K images+labels)
mnist = mnist_data.read_data_sets("data", one_hot=True, reshape=False, validation_size=0)

# input X: 28x28 grayscale images, the first dimension (None) will index the images in the mini-batch
X = tf.placeholder(tf.float32, [None, 28, 28, 1])
# correct answers will go here
Y_ = tf.placeholder(tf.float32, [None, 10])
# Variable (decaying) learning rate
learning_rate = tf.placeholder(tf.float32)

# Number of neurons in each layer - 7 total layers (Last layer has 10 softMax neurons)
A = 400
B = 200
C = 200
D = 100
E = 60
F = 30

# Sigmoid Layers
Wa = tf.Variable(tf.truncated_normal([784, A], stddev = 0.1))
Ba = tf.Variable(tf.zeros([A]))

Wb = tf.Variable(tf.truncated_normal([A, B], stddev = 0.1))
Bb = tf.Variable(tf.zeros([B]))

Wc = tf.Variable(tf.truncated_normal([B, C], stddev = 0.1))
Bc = tf.Variable(tf.zeros([C]))

Wd = tf.Variable(tf.truncated_normal([C, D], stddev = 0.1))
Bd = tf.Variable(tf.zeros([D]))

We = tf.Variable(tf.truncated_normal([D, E], stddev = 0.1))
Be = tf.Variable(tf.zeros([E]))

Wf = tf.Variable(tf.truncated_normal([E, F], stddev = 0.1))
Bf = tf.Variable(tf.zeros([F]))

# Last layer with softMax neurons
W = tf.Variable(tf.truncated_normal([F, 10], stddev = 0.1))
B = tf.Variable(tf.zeros([10]))

# flatten the images into a single line of pixels
# -1 in the shape definition means "the only possible dimension that will preserve the number of elements"
XX = tf.reshape(X, [-1, 784])

# The model
Ya = tf.nn.relu(tf.matmul(XX, Wa) + Ba)
Yb = tf.nn.relu(tf.matmul(Ya, Wb) + Bb)
Yc = tf.nn.relu(tf.matmul(Yb, Wc) + Bc)
Yd = tf.nn.relu(tf.matmul(Yc, Wd) + Bd)
Ye = tf.nn.relu(tf.matmul(Yd, We) + Be)
Yf = tf.nn.relu(tf.matmul(Ye, Wf) + Bf)
Ylogits = tf.matmul(Yf, W) + B
Y = tf.nn.softmax(Ylogits)

# loss function: cross-entropy = - sum( Y_i * log(Yi) )
#                           Y: the computed output vector
#                           Y_: the desired output vector

# cross-entropy
# log takes the log of each element, * multiplies the tensors element by element
# reduce_mean will add all the components in the tensor
# so here we end up with the total cross-entropy for all images in the batch
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(logits=Ylogits, labels=Y_)
cross_entropy = tf.reduce_mean(cross_entropy)*100

# accuracy of the trained model, between 0 (worst) and 1 (best)
correct_prediction = tf.equal(tf.argmax(Y, 1), tf.argmax(Y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

# training, learning rate is variable and decaying
train_step = tf.train.AdamOptimizer(learning_rate).minimize(cross_entropy)

# matplotlib visualisation
allweights = tf.concat([tf.reshape(Wa, [-1]),
                        tf.reshape(Wb, [-1]),
                        tf.reshape(Wc, [-1]),
                        tf.reshape(Wd, [-1]),
                        tf.reshape(We, [-1]),
                        tf.reshape(Wf, [-1])
                        ], 0)
allbiases = tf.concat([tf.reshape(Ba, [-1]),
                       tf.reshape(Bb, [-1]),
                       tf.reshape(Bc, [-1]),
                       tf.reshape(Bd, [-1]),
                       tf.reshape(Be, [-1]),
                       tf.reshape(Bf, [-1])
                       ], 0)
I = tensorflowvisu.tf_format_mnist_images(X, Y, Y_)  # assembles 10x10 images by default
It = tensorflowvisu.tf_format_mnist_images(X, Y, Y_, 1000, lines=25)  # 1000 images on 25 lines
datavis = tensorflowvisu.MnistDataVis()

# init
init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)


# You can call this function in a loop to train the model, 100 images at a time
def training_step(i, update_test_data, update_train_data):

    # training on batches of 100 images with 100 labels
    batch_X, batch_Y = mnist.train.next_batch(100)

    # Learning rate decay
    lr_min = 0.0001
    lr_max = 0.003
    decay_speed = 2000
    lr= lr_min + (lr_max - lr_min) * math.exp(-1/2000)

    # compute training values for visualisation
    if update_train_data:
        a, c, im, w, b = sess.run([accuracy, cross_entropy, I, allweights, allbiases], feed_dict={X: batch_X, Y_: batch_Y})
        datavis.append_training_curves_data(i, a, c)
        datavis.append_data_histograms(i, w, b)
        datavis.update_image1(im)
        print(str(i) + ": accuracy:" + str(a) + " loss: " + str(c))

    # compute test values for visualisation
    if update_test_data:
        a, c, im = sess.run([accuracy, cross_entropy, It], feed_dict={X: mnist.test.images, Y_: mnist.test.labels})
        datavis.append_test_curves_data(i, a, c)
        datavis.update_image2(im)
        print(str(i) + ": ********* epoch " + str(i*100//mnist.train.images.shape[0]+1) + " ********* test accuracy:" + str(a) + " test loss: " + str(c) + " learning rate: " + str(lr))

    # the backpropagation training step
    sess.run(train_step, feed_dict={X: batch_X, Y_: batch_Y, learning_rate: lr})


datavis.animate(training_step, iterations=5000+1, train_data_update_freq=10, test_data_update_freq=50, more_tests_at_start=True)

# to save the animation as a movie, add save_movie=True as an argument to datavis.animate
# to disable the visualisation use the following line instead of the datavis.animate line
# for i in range(2000+1): training_step(i, i % 50 == 0, i % 10 == 0)

print("max test accuracy: " + str(datavis.get_max_test_accuracy()))

# final max test accuracy = 0.9268 (10K iterations). Accuracy should peak above 0.92 in the first 2000 iterations.
