'''Example script to predict sequence using stateful rnns.
At least 10 epochs are required before the generated text
starts sounding coherent.
'''

import numpy as np
import matplotlib.pyplot as mpl
from keras.models import Sequential
from keras.layers.core import Dense
from keras.layers.recurrent import LSTM


# since we are using stateful rnn tsteps can be set to 1
tsteps = 1
batch_size = 25
epochs = 25
# number of elements ahead that are used to make the prediction
lahead = 1


def gen_cosine_amp(amp=100, period=25, x0=0, xn=50000, step=1, k=0.0001):
    """
    Generates an absolute cosine time series with the amplitude exponentially
    decreasing

    Keyword arguments:
    amp -- amplitude of the cosine function
    period -- period of the cosine function
    x0 -- initial x of the time series
    xn -- final x of the time series
    step -- step of the time series discretization
    k -- exponential rate
    """
    cos = np.zeros(((xn - x0) * step, 1, 1))
    for i in range(len(cos)):
        idx = x0 + i * step
        cos[i, 0, 0] = amp * np.cos(idx / (2 * np.pi * period))
        cos[i, 0, 0] = cos[i, 0, 0] * np.exp(-k * idx)
    return cos


print('Creating Data')
cos = gen_cosine_amp()
print('Input shape:')
print(cos.shape)
print('Calculating expected predicted_out')
expected_out = np.zeros((len(cos), 1))
for i in range(len(cos) - lahead):
    expected_out[i, 0] = np.mean(cos[i + 1:i + lahead + 1])

print('Output shape')
print(expected_out.shape)

print('Creating Model')
model = Sequential()
model.add(
    LSTM(
        50,
        batch_input_shape=(
            batch_size,
            tsteps,
            1),
        return_sequences=True,
        stateful=True))
model.add(
    LSTM(
        50,
        batch_input_shape=(
            batch_size,
            tsteps,
            1),
        return_sequences=False,
        stateful=True))
model.add(Dense(1))
model.compile(loss='rmse', optimizer='rmsprop')

print('Training')
for i in range(epochs):
    model.fit(
        cos,
        expected_out,
        batch_size=batch_size,
        verbose=1,
        nb_epoch=1)
    model.reset_states()

print('Predicting')
predicted_out = model.predict(cos, batch_size=batch_size)

print('Ploting Results')
mpl.subplot(2, 1, 1)
mpl.plot(expected_out)
mpl.title('Expected')
mpl.subplot(2, 1, 2)
mpl.plot(predicted_out)
mpl.title('Predicted')
mpl.show()

Add example of stateful LSTM sequence prediction
Add example of stateful LSTM sequence prediction
