import numpy as np
import tensorflow as tf
from tensorflow.contrib import rnn
from pre_processing import *

X_test, y_test, X_train, y_train = get_all_data()
print(X_train.shape)
print(X_test.shape)
print(y_train.shape)
print(y_test.shape)

# TF input
x_in = tf.placeholder(tf.float32, [None, 22, 1000], name="input_x")
y_real = tf.placeholder(tf.int32, [None], name="real_y")
y_temp = tf.one_hot(y_real, 4)

x = tf.unstack(x_in, 1000, axis=2)
lstm_cell = rnn.BasicLSTMCell(100)
outputs, states = rnn.static_rnn(lstm_cell, x, dtype=tf.float32)

logits = tf.layers.dense(inputs=outputs[-1], units=4)

# calculate accuracy
prediction = tf.argmax(logits, 1, output_type=tf.int32)
my_acc = tf.reduce_mean(tf.cast(tf.equal(y_real, prediction), tf.float32))
test_accu_summ = tf.summary.scalar("test_accu", my_acc)
train_accu_summ = tf.summary.scalar("train_accu", my_acc)

loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(labels=y_temp, logits=logits))
test_loss_sum = tf.summary.scalar("test_loss", loss)
train_loss_sum = tf.summary.scalar("train_loss", loss)

optimizer = tf.train.AdamOptimizer().minimize(loss)


writer = tf.summary.FileWriter("./base_rnn/1")
init = tf.global_variables_initializer()
with tf.Session() as sess:
    sess.run(init)
    sess.run(tf.local_variables_initializer())
    writer.add_graph(sess.graph)

    for k in range(20):
        t_size = np.array(range(X_train.shape[0]))
        np.random.shuffle(t_size)
        test_mask = t_size[:100]
        # print("*****k is:", k)
        for i in range(30):
            inputs = {x_in: X_train[test_mask], y_real: y_train[test_mask]}
            cost, _, summ_1, summ_2 = sess.run([loss, optimizer, train_accu_summ, train_loss_sum], feed_dict=inputs)
            writer.add_summary(summ_1, (i + k * 30))
            writer.add_summary(summ_2, (i + k * 30))

            # print("training loss is:", cost)
            inputs = {x_in: X_test, y_real: y_test}
            accu, cost, summ_1, summ_2 = sess.run([my_acc, loss, test_accu_summ, test_loss_sum], feed_dict=inputs)
            writer.add_summary(summ_1, (i + k * 30))
            writer.add_summary(summ_2, (i + k * 30))
            # print("----test accuracy is:", accu)
            # print("test cost is:", cost)

