import tensorflow as tf
from tensorflow.contrib.learn.python.learn.estimators.estimator import SKCompat
import numpy as np
import matplotlib.pyplot as plt
import sys


HIDDEN_SIZE = 30
NUM_HIDDEN_LAYERS = 2
BATCH_SIZE = 64
TRAINING_STEPS = 1000 # 10000 will highly simulate

TRAINING_EXAMPLES = 10000
TESTING_EXAMPLES = 1000
SAMPLE_GAP = 0.01
TIMESTEPS = 10


def lstm_cell():
    return tf.nn.rnn_cell.LSTMCell(HIDDEN_SIZE, state_is_tuple=True)
                                   
def lstm_model(x, y):
    cells = tf.nn.rnn_cell.MultiRNNCell([lstm_cell() for _ in range(NUM_HIDDEN_LAYERS)])
    
    output, _ = tf.nn.dynamic_rnn(cells, x, dtype=tf.float32)
    
    output = tf.reshape(output, [-1, HIDDEN_SIZE])
    predictions = tf.contrib.layers.fully_connected(output, 1, None)
    predictions = tf.reshape(predictions, [-1])
    labels = tf.reshape(y, [-1])
    loss = tf.losses.mean_squared_error(predictions, labels)
    op = tf.contrib.layers.optimize_loss(loss,
                                         tf.train.get_global_step(),
                                         optimizer='Adagrad',
                                         learning_rate=0.1)
                                         
    return predictions, loss, op


# get train and test datasets
def generate_data(seq):
    x = []
    y = []
    for i in range(len(seq) - TIMESTEPS - 1):
        x.append([seq[i:i+TIMESTEPS]])
        y.append([seq[i+TIMESTEPS]])
    return np.array(x, dtype=np.float32), np.array(y, dtype=np.float32)

test_start = TRAINING_EXAMPLES * SAMPLE_GAP # 100
test_end = (TRAINING_EXAMPLES + TESTING_EXAMPLES) * SAMPLE_GAP # 110
train_x, train_y = generate_data(np.sin(np.linspace(0, test_start, TRAINING_EXAMPLES, dtype=np.float32))) # [9989, 1, 10] / [9989, 1]
test_x, test_y = generate_data(np.sin(np.linspace(test_start, test_end, TESTING_EXAMPLES, dtype=np.float32))) # [989, 1, 10] / [989, 1]


regressor = SKCompat(tf.contrib.learn.Estimator(model_fn=lstm_model))
regressor.fit(train_x, train_y, batch_size=BATCH_SIZE, steps=TRAINING_STEPS)
predicted = [[pred] for pred in regressor.predict(test_x)]
rmse = np.sqrt(((predicted - test_y) ** 2).mean(axis=0))
print("Mean Square Error is:%f" % rmse[0])


plot_predicted, = plt.plot(predicted, label='predicted')
plot_test, = plt.plot(test_y, label='real_sin')
plt.legend([plot_predicted, plot_test],['predicted', 'real_sin'])
plt.show()




'''
# problems:
tf.contrib.layers.fully_connected(input_tensor, output_size, act)
@act should specify None, or simulation is failed
tf.contrib.learn.Estimator(model_dir) makes some difference, although I don't know why
'''