import tensorflow as tf
import numpy as np

PAD = 0
EOS = 1
vocab_size = 672
output_size= 114
input_embedding_size = 1000
encoder_hidden_units = 20
decoder_hidden_units = 20
batch_size = 100



from pyreader import reader
model=reader()


train_graph = tf.Graph()
with train_graph.as_default():
    encoder_inputs = tf.placeholder(shape=(None, None), dtype=tf.int32, name='encoder_inputs')
    decoder_inputs = tf.placeholder(shape=(None, None), dtype=tf.int32, name='decoder_inputs')
    decoder_targets = tf.placeholder(shape=(None, None), dtype=tf.int32, name='decoder_targets')

    embeddings = tf.Variable(tf.random_uniform([vocab_size, input_embedding_size], -1.0, 1.0), dtype=tf.float32)
    encoder_inputs_embedded = tf.nn.embedding_lookup(embeddings, encoder_inputs)
    decoder_inputs_embedded = tf.nn.embedding_lookup(embeddings, decoder_inputs)

    encoder_cell = tf.contrib.rnn.LSTMCell(encoder_hidden_units)
    encoder_outputs, encoder_final_state = tf.nn.dynamic_rnn(
        encoder_cell, encoder_inputs_embedded,
        dtype=tf.float32, time_major=False,
    )
    decoder_cell = tf.contrib.rnn.LSTMCell(decoder_hidden_units)
    decoder_outputs, decoder_final_state = tf.nn.dynamic_rnn(
        decoder_cell, decoder_inputs_embedded,
        initial_state=encoder_final_state,
        dtype=tf.float32, time_major=False, scope="plain_decoder",
    )

    decoder_logits = tf.contrib.layers.linear(decoder_outputs, vocab_size)
    decoder_prediction = tf.argmax(decoder_logits, 2)
    stepwise_cross_entropy = tf.nn.softmax_cross_entropy_with_logits(
        labels=tf.one_hot(decoder_targets, depth=vocab_size, dtype=tf.float32),
        logits=decoder_logits,
    )
    loss = tf.reduce_mean(stepwise_cross_entropy)
    train_op = tf.train.AdamOptimizer().minimize(loss)
    saver=tf.train.Saver()

loss_track = []
epochs = 30000001


with tf.Session(graph=train_graph) as sess:
    sess.run(tf.global_variables_initializer())
#    ckpt = tf.train.get_checkpoint_state('tense/py')
#    saver.restore(sess, ckpt.model_checkpoint_path)
    for epoch in range(1,epochs):
        answers,pads,inputs = model.list_tags(batch_size=batch_size)

        decoder_targets_ = [(sequence) + [EOS] for sequence in answers]
        decoder_inputs_ = [[EOS] + (sequence) for sequence in answers]
        feed_dict = {encoder_inputs: inputs, decoder_inputs: decoder_inputs_,
                     decoder_targets: decoder_targets_,
                     }
        _, l = sess.run([train_op, loss], feed_dict)
        loss_track.append(l)
        if epoch == 1 or epoch % 100 == 0:
            print('saved to: ', saver.save(sess,'tense/py/py.ckpt',global_step=epoch))
            print('loss: {}'.format(sess.run(loss, feed_dict)))
        if epoch % 500 == 0:
            predict_ = sess.run(decoder_prediction, feed_dict)
           # for i, (inp, pred) in enumerate(zip(np.array(feed_dict[encoder_inputs]).T, predict_.T)):
           #     print('input > {}'.format(inp))
           #     print('predicted > {}'.format(pred))
           #     if i >= 20:
           #         break

