# -*- coding: utf-8 -*-

import tflearn
import tflearn.datasets.mnist as mnist

mnistdata = mnist.read_data_sets("mnist/", one_hot=True)
# 55000 images
X     = mnistdata.train.images
Y     = mnistdata.train.labels
# 10000 images
Xtest = mnistdata.test.images
Ytest = mnistdata.test.labels
# 5000 images
Xval  = mnistdata.validation.images
Yval  = mnistdata.validation.labels

# Building deep neural network
input_layer  = tflearn.input_data(shape=[None, 784])
hidden_layer = tflearn.fully_connected(input_layer, 64)
output_layer = tflearn.fully_connected(hidden_layer, 10, activation='softmax')
net = tflearn.regression(output_layer)
# Training
model = tflearn.DNN(net, tensorboard_verbose=0)
model.fit(X, Y, n_epoch=20, validation_set=(Xtest, Ytest), show_metric=True, run_id="mnist")


# checking
predictions = model.predict_label(Xval)

# Visualizing the data
import matplotlib.pyplot as plt
import random
import time
for i in range(1,10):
  n = random.randint(1,5000)
  label = Yval[n].argmax(axis=0)
  image = Xval[n].reshape([28,28])
  pred = predictions[n][0]
  if (pred == label):
    rec = "Correct"
  else:
    rec = "Incorrect"
  plt.title('%s: index: %d, label: %d, prediction: %d' % (rec, n, label, pred))
  plt.imshow(image, cmap='gray_r')
  plt.show(block=False)
  time.sleep(2)
  plt.close('all')