import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np

tf.reset_default_graph()

samples = 1000
times = [1e-2*i for i in range(samples+1)]
sin = np.sin(times[:-1])
sin_next = np.sin(times[1:])

time_step = 10
sin = np.reshape(sin, [-1, time_step, 1])
sin_next = np.reshape(sin_next, [-1, 1])

signal = tf.placeholder(tf.float32,
                        shape=[None, time_step, 1])
signal_next = tf.placeholder(tf.float32,
                            [None, 1])

unstacked_signal = tf.unstack(signal, axis=1)
# print('inputs: ')
# for t in unstacked_signal:
#     print(t)

state_size = 30
rnn_cell = tf.nn.rnn_cell.BasicRNNCell(num_units=state_size)
states, state = tf.nn.static_rnn(cell=rnn_cell,
                                 inputs=unstacked_signal,
                                 dtype=tf.float32)

# print('outputs: ')
# for s in states:
#     print(s)
# print(state.shape)

states = tf.stack(states, axis=1)
reshaped_states = tf.reshape(states, [-1, state_size])
print(reshaped_states.shape)

output = tf.layers.dense(reshaped_states, 1, use_bias=False)
print(output.shape)

loss = tf.losses.mean_squared_error(signal_next, output)
train_op = tf.train.GradientDescentOptimizer(1e-2).minimize(loss)

accuracy = tf.contrib.metrics.streaming_pearson_correlation(output, signal_next)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    sess.run(tf.local_variables_initializer())
    for i in range(3000):
        _, _loss, _acc = sess.run([train_op, loss, accuracy],
                            feed_dict={signal: sin, signal_next: sin_next})
        if i%100 == 0:
            print('step: {}, loss: {}, acc: {}'.format(i, _loss, _acc[0]))

    _pred = sess.run(output, feed_dict={signal: sin})
    plt.figure(1)
    plt.plot(sin_next)
    plt.figure(2)
    plt.plot(_pred)
    plt.show()