#import statements
from keras.layers.core import Dense, Activation, Dropout
from keras.layers.recurrent import LSTM
from keras.models import Sequential
import lstm, time
import pandas as pd
import matplotlib as plt
from sklearn.preprocessing import StandardScaler
import numpy as np

#read the annotated file 
spx_data = pd.read_csv('data/spx_data.csv')

"""
Explanation: Predicting stock prices with LSTM networks
By Khaled Tinubu
End of explanation
"""


spx_data.head()

"""
Explanation: This is what our data looks like
We have 6 features: The opening price, the daily high, low, the adjusted closing price, the volume and finally the actual closing price 
We will first preprocess our data and then build a LSTM network in order to model the fluctuations of the S&P 500 as a time series
Here's a look at what the data looks like:
End of explanation
"""


# reorder data frame to put the closing price at the end of the arrays
spx_data = spx_data[['Open', 'High', 'Low', 'Adj Close', 'Volume', 'Close']]
spx_data_close =  spx_data[['Close']]
spx_data

plt.pyplot.plot(spx_data_close)
plt.pyplot.show()
#this is a plot of the closing prices over time

#normalize the data
normalized_spx_data = StandardScaler().fit_transform(spx_data)



## N.B the function below is a modifed form of a helped function written by Jakob Aungier 
## (full citation in lstm.py)

def prep_data(data, seq_len):
    """
    Returns the input data and targets for testing and validation sets 
    data -> the full corpus of data
    
    seq_len -> the length of our sequence

    """
    sequence_length = seq_len + 1
    result = []
    for index in range(len(data) - sequence_length):
        result.append(data[index: index + sequence_length])
        
    result = np.array(result)
    
    # 90% train, 10% validation partition
    partition = round(0.9 * result.shape[0])
    train = result[:int(partition), :]
    np.random.shuffle(train)
    x_train = train[:, :-1]
    y_train = train[:, -1, -1]
    x_test = result[int(partition):, :-1]
    y_test = result[int(partition):, -1, -1]

    x_train = np.reshape(x_train, (x_train.shape[0], x_train.shape[1], 6))
    x_test = np.reshape(x_test, (x_test.shape[0], x_test.shape[1], 6))  

    return [x_train, y_train, x_test, y_test]

X_train, y_train, X_test, y_test  = prep_data(normalized_spx_data, 20)

#building Model
model = Sequential()

model.add(LSTM(input_dim=6,
               output_dim=50,
               return_sequences=True))

model.add(Dropout(0.2))

model.add(LSTM(100, return_sequences=False))
model.add(Dropout(0.2))

model.add(Dense(
            output_dim=1))
model.add(Activation('linear'))

start = time.time()
model.compile(loss='mse', optimizer='rmsprop')

print('compilation time:', time.time() - start)

"""
Explanation: now we can proceed with the analysis
End of explanation
"""


#Train the model
model.fit(
    X_train,
    y_train,
    batch_size=512,
    nb_epoch=200,
    validation_split=0.05)

"""
Explanation: TRAINING TIME!
End of explanation
"""


