import numpy as np
import pandas as pd
import argparse

from pprint import pprint

#Allow relative imports to directories above cwd/
sys.path.insert(1, os.path.join(sys.path[0], '..'))
import twModel

import importlib
importlib.reload(twModel)

import matplotlib.pyplot as plt
%matplotlib inline

np.random.seed(42)

from sklearn.metrics import mean_squared_error

"""
Explanation: LSTM and the adding problem.
A basic baseline is to always predict the sum to have a value of 1 regardless of the inputs. This will
give the Mean Squared Error (MSE) around 0.1767. The goal is to train a model that achieves MSE
well below 0.1767.
End of explanation
"""


def create_data(args, T=10, n=10, split=0.8):
    x = np.zeros((n, 2, T))
    y = np.zeros((n,))
    idx_arr = np.arange(T)

    for i in range(n):
        idx = np.random.choice(idx_arr, 2, replace=False)
        x[i, 0, idx] = 1
        x[i, 1, :] = np.random.uniform(0, 1, size=T)
        y[i] = np.dot(x[i, 1], x[i, 0].T)

    x = np.swapaxes(x, 2, 1)

    row = round(n*split)

    return x[:row], y[:row], x[row:], y[row:]
X_train, y_train, X_val, y_val = create_data(None, 10, 1, split=1)
X_train[0]
y_train

"""
Explanation: Create data
End of explanation
"""


ltsm50 = pd.read_csv('./lstm_T50_bs128/LSTMModel_log.csv', sep=',')
ltsm100 = pd.read_csv('./lstm_T100_bs128/LSTMModel_log.csv', sep=',')
ltsm150 = pd.read_csv('./lstm_T150_bs128/LSTMModel_log.csv', sep=',')
dense50 = pd.read_csv('./dense_T50_bs254/DenseModel_log.csv', sep=',')
dense100 = pd.read_csv('./dense_T100_bs254/DenseModel_log.csv', sep=',')
dense150 = pd.read_csv('./dense_T150_bs254/DenseModel_log.csv', sep=',')

fig = plt.figure(facecolor='white', figsize=(20,10))
_ = plt.plot(ltsm50['val_loss'], label='ltsm50', ls='--')
_ = plt.plot(ltsm100['val_loss'], label='ltsm100', ls='--')
_ = plt.plot(ltsm150['val_loss'], label='ltsm150', ls='--')
_ = plt.plot(dense50['val_loss'], label='dense50')
_ = plt.plot(dense100['val_loss'], label='dense100')
_ = plt.plot(dense150['val_loss'], label='dense150')
_ = plt.plot([0, 16], [0.1767, 0.1767], 'k-', label='baseline')
_ = plt.xlim([0,20])
_ = plt.xlabel('epoch')
_ = plt.ylabel('mse')
_ = plt.legend()

"""
Explanation: Model architectures
Dense

LTSM

Results
Results for 50, 100 and 150 sequence length.
Dense model all converge and train well, whereas LSTM models only converge for sequence legnth 50. Dense models >50 tend to overfit so dropout and regularization might be suitable. However for proof of concept I did not delve into that.
End of explanation
"""


