import tensorflow as tf
import numpy as np
import time

def network(input_list):
    input,init_hidden_c,init_hidden_m = input_list
    cell = tf.nn.rnn_cell.BasicLSTMCell(256, state_is_tuple=True)
    init_hidden = tf.nn.rnn_cell.LSTMStateTuple(init_hidden_c, init_hidden_m)
    states, hidden_cm = tf.nn.dynamic_rnn(cell, input, dtype=tf.float32, initial_state=init_hidden)
    net = [v for v in tf.trainable_variables()]
    return states, hidden_cm, net

def action(x, h_c, h_m):
    t0 = time.time()
    outputs, output_h = sess.run([rnn_states[:,-1:,:], rnn_hidden_cm], feed_dict={
    rnn_input:x,
    rnn_init_hidden_c: h_c,
    rnn_init_hidden_m: h_m
})
    dt = time.time() - t0
    return outputs, output_h, dt

rnn_input = tf.placeholder("float", [None, None, 512])
rnn_init_hidden_c = tf.placeholder("float", [None,256])
rnn_init_hidden_m = tf.placeholder("float", [None,256])
rnn_input_list = [rnn_input, rnn_init_hidden_c, rnn_init_hidden_m]
rnn_states, rnn_hidden_cm, rnn_net = network(rnn_input_list)

feed_input = np.random.uniform(low=-1.,high=1.,size=(1,1,512))
feed_init_hidden_c = np.zeros(shape=(1,256))
feed_init_hidden_m = np.zeros(shape=(1,256))

sess = tf.Session()
sess.run(tf.global_variables_initializer())
for i in range(10000):
    _, output_hidden_cm, deltat = action(feed_input, feed_init_hidden_c, feed_init_hidden_m)
    if i % 10 == 0:
        print 'Running time: ' + str(deltat)
    (feed_init_hidden_c, feed_init_hidden_m) = output_hidden_cm
    feed_input = np.random.uniform(low=-1.,high=1.,size=(1,1,512))