import tensorflow
from tensorflow import keras
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, LSTM, Dropout, RepeatVector, TimeDistributed
import numpy as np
import os
import matplotlib.pyplot as plt
import sys


def create_model(X):
    model = Sequential()
    model.add(LSTM(10, input_shape=(X.shape[1], X.shape[2])))
    model.add(Dropout(rate=0.2))
    model.add(RepeatVector(X.shape[1]))
    model.add(LSTM(10, return_sequences=True))
    model.add(Dropout(rate=0.2))
    model.add(TimeDistributed(Dense(X.shape[2])))
    model.compile(optimizer='adam', loss='mse')
    model.summary()

    history = model.fit(X, X, epochs=100, batch_size=1, validation_split=0.1,
                        callbacks=[keras.callbacks.EarlyStopping(monitor='loss', patience=3, mode='min')],
                        shuffle=False)

    model.save('Working_Data/lstm_model')
    # model.predict(X[0:10, :])

    # plot the loss
    plt.plot(history.history['loss'])
    plt.plot(history.history['val_loss'])
    plt.title('model loss')
    plt.ylabel('loss')
    plt.xlabel('epoch')
    plt.legend(['train', 'test'], loc='upper left')
    plt.savefig("Working_Data/lstm_loss.png")
    plt.show()

    print("loss of the model is: ")
    print(history.history['loss'])


def create_sequences(data):
    Xs, ys = [], []
    time_steps = 10
    for i in range(len(data) - time_steps):
        Xs.append(data[i:(i + time_steps)].reshape(100*time_steps,4))
        ys.append(data[i + time_steps].reshape(100,4))

    return np.array(Xs), np.array(ys)


data = np.load(os.path.join("Working_Data/Normalized_Fixed_Dim_HBs_Idx" + str(1) + ".npy"))
data = data[0:1000, :, :]
# print(data[0:10].reshape(10000,4).shape)
X, y = create_sequences(data)
print(X.shape, y.shape)
# create_model(X)

model = keras.models.load_model('Working_Data/lstm_model')
model.predict(create_sequences(X[0:5, :, :])[0])




tensorflow.keras.models.load_model matplotlib.pyplot.legend tensorflow.keras.layers.Dropout matplotlib.pyplot.title tensorflow.keras.layers.Dense tensorflow.keras.callbacks.EarlyStopping matplotlib.pyplot.savefig matplotlib.pyplot.plot tensorflow.keras.layers.RepeatVector tensorflow.keras.layers.LSTM matplotlib.pyplot.xlabel numpy.array matplotlib.pyplot.show tensorflow.keras.models.Sequential matplotlib.pyplot.ylabel
[(60, 'tensorflow.keras.models.load_model', 'keras.models.load_model', (['"""Working_Data/lstm_model"""'], {}), False, 'from tensorflow import keras\n'), (12, 'tensorflow.keras.models.Sequential', 'Sequential', ([], {}), False, 'from tensorflow.keras.models import Sequential\n'), (30, 'matplotlib.pyplot.plot', 'plt.plot', (["history.history['loss']"], {}), True, 'import matplotlib.pyplot as plt\n'), (31, 'matplotlib.pyplot.plot', 'plt.plot', (["history.history['val_loss']"], {}), True, 'import matplotlib.pyplot as plt\n'), (32, 'matplotlib.pyplot.title', 'plt.title', (['"""model loss"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (33, 'matplotlib.pyplot.ylabel', 'plt.ylabel', (['"""loss"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (34, 'matplotlib.pyplot.xlabel', 'plt.xlabel', (['"""epoch"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (35, 'matplotlib.pyplot.legend', 'plt.legend', (["['train', 'test']"], {'loc': '"""upper left"""'}), True, 'import matplotlib.pyplot as plt\n'), (36, 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""Working_Data/lstm_loss.png"""'], {}), True, 'import matplotlib.pyplot as plt\n'), (37, 'matplotlib.pyplot.show', 'plt.show', ([], {}), True, 'import matplotlib.pyplot as plt\n'), (13, 'tensorflow.keras.layers.LSTM', 'LSTM', (['(10)'], {'input_shape': '(X.shape[1], X.shape[2])'}), False, 'from tensorflow.keras.layers import Dense, LSTM, Dropout, RepeatVector, TimeDistributed\n'), (14, 'tensorflow.keras.layers.Dropout', 'Dropout', ([], {'rate': '(0.2)'}), False, 'from tensorflow.keras.layers import Dense, LSTM, Dropout, RepeatVector, TimeDistributed\n'), (15, 'tensorflow.keras.layers.RepeatVector', 'RepeatVector', (['X.shape[1]'], {}), False, 'from tensorflow.keras.layers import Dense, LSTM, Dropout, RepeatVector, TimeDistributed\n'), (16, 'tensorflow.keras.layers.LSTM', 'LSTM', (['(10)'], {'return_sequences': '(True)'}), False, 'from tensorflow.keras.layers import Dense, LSTM, Dropout, RepeatVector, TimeDistributed\n'), (17, 'tensorflow.keras.layers.Dropout', 'Dropout', ([], {'rate': '(0.2)'}), False, 'from tensorflow.keras.layers import Dense, LSTM, Dropout, RepeatVector, TimeDistributed\n'), (50, 'numpy.array', 'np.array', (['Xs'], {}), True, 'import numpy as np\n'), (50, 'numpy.array', 'np.array', (['ys'], {}), True, 'import numpy as np\n'), (18, 'tensorflow.keras.layers.Dense', 'Dense', (['X.shape[2]'], {}), False, 'from tensorflow.keras.layers import Dense, LSTM, Dropout, RepeatVector, TimeDistributed\n'), (23, 'tensorflow.keras.callbacks.EarlyStopping', 'keras.callbacks.EarlyStopping', ([], {'monitor': '"""loss"""', 'patience': '(3)', 'mode': '"""min"""'}), False, 'from tensorflow import keras\n')]