import numpy as np
import pandas as pd
import sys
from keras.models import Sequential
from keras.layers import LSTM, Dense, Dropout
import myenv

"""
Function to create sequences for learning lstm by input parameters
input: -- npArray -- scale indicator and close prices in numpy format
       --  periodB -- int to back period to learn 
       -- predictN -- number in days/interval to predict 
return: -- numpy array X -- history prices and technical indicators
        -- numpy array Y -- predict prices
"""
def createSequences(periodB, predictN, npArray):
    X, y = [], []
    for i in range(len(npArray)):
        end = i + periodB
        out_end = end + predictN
        if out_end > len(npArray):
            break
        arrayX, arrayY = npArray[i:end, :], npArray[end:out_end, 0]
        X.append(arrayX)
        y.append(arrayY)
    return np.array(X), np.array(y)

"""
Function to create layer for neural network. Function inspired from marcosan93 from https://github.com/marcosan93/Price-Forecaster
and kenas documentation.
input: -- layerNumber -- number of layer 
       -- nodes -- number of nodes 
       -- activation -- string activation function 
       -- drop -- number
       -- dRate -- float number
"""
def addLayer(layerNumber, nodes, activation, drop=None, dRate=.5):
    for x in range(1,layerNumber+1):
        model.add(LSTM(nodes, activation=activation, return_sequences=True))
        try:
            if x % drop == 0:
                model.add(Dropout(dRate))
        except:
            pass

# load prepare datasets in  pickle format
df = pd.read_pickle(myenv.PREPARED_DATASET) 

# periods looking back to learn
periodB  = myenv.PERIODB
# periods to predict
predictPeriod = myenv.PREDICT_PERIOD

data = df.shape[1]

X, y = createSequences(periodB, predictPeriod, df.to_numpy())

# Instatiating the model
model = Sequential()

activ = myenv.ACTIV_FUNCTION

# Data layer https://keras.io/api/layers/recurrent_layers/lstm/
model.add(LSTM(90, activation=activ, return_sequences=True, input_shape=(periodB, data)))

# Layers
addLayer(myenv.LAYERS, nodes=myenv.NODES, activation=activ, drop=1, dRate=.1)
model.add(LSTM(myenv.LAYERS_LTSM, activation=activ))

# Output predit layer
model.add(Dense(predictPeriod))

# Compute
model.summary()
model.compile(optimizer=myenv.OPTIMIZER, loss='mse', metrics=['accuracy'])
res = model.fit(X, y, epochs=myenv.EPOCHS, batch_size=32, validation_split=0.1)

model.save(myenv.MODEL_PREDICT_SAVE, save_format='tf')