import tensorflow as tf
import numpy as np
from datetime import datetime

now = datetime.utcnow().strftime("%Y%m%d%H%M%S")
root_logdir = "tf_logs"
logdir = "{}/run-{}/".format(root_logdir, now)

# =============================================================================
# Cells of 150 recurrent neurons, plus a fully connected layer of 10 neurons
# =============================================================================
tf.reset_default_graph()
n_steps = 28
n_inputs = 28
n_neurons = 150
n_outputs = 10

learning_rate = 0.001

X = tf.placeholder(tf.float32, [None, n_steps, n_inputs])
y = tf.placeholder(tf.int32, [None])

basic_cell = tf.contrib.rnn.BasicRNNCell(num_units=n_neurons)
outputs, states = tf.nn.dynamic_rnn(basic_cell, X, dtype=tf.float32)

logits = tf.layers.dense(states, n_outputs)
xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y,
                                                          logits=logits)
loss = tf.reduce_mean(xentropy)
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
training_op = optimizer.minimize(loss)
correct = tf.nn.in_top_k(logits, y, 1)
accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))

xentropy_summary = tf.summary.scalar('X_entropy', xentropy)
file_writer = tf.summary.FileWriter(logdir, tf.get_default_graph())

init = tf.global_variables_initializer()

from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("/tmp/data")
X_test = mnist.test.images.reshape((-1, n_steps, n_inputs))
y_test = mnist.test.labels

#n_epochs = 100
#batch_size = 150
#
#with tf.Session() as sess:
#    init.run()
#    for epoch in range(n_epochs):
#        for iterastion in range(mnist.train.num_examples // batch_size):
#            X_batch, y_batch = mnist.train.next_batch(batch_size)
#            X_batch = X_batch.reshape((-1, n_steps, n_inputs))
#            summary_str = xentropy_summary.eval(feed_dict={X: X_batch, y: y_batch})
#            step = epoch*(mnist.train.num_examples // batch_size) + batch_size
#            file_writer.add_summary(summary_str, step)
#            sess.run(training_op, feed_dict={X: X_batch, y: y_batch})
#        acc_train = accuracy.eval(feed_dict={X: X_batch, y: y_batch})
#        acc_test = accuracy.eval(feed_dict={X: X_batch, y: y_batch})
#        print(epoch, "Train accuracy:", acc_train, "Test accuracy:", acc_test)
#    
#    file_writer.close()