# Import
import tensorflow as tf
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
import matplotlib.pyplot as plt

# Import data
data = pd.read_csv('./dataset/data_stocks.csv')

# Drop date variable
data = data.drop(['DATE'], 1)

# Make data a np.array
data = data.values

# Training and val data
train_start = 0
train_end = int(np.floor(0.8*data.shape[0]))
val_start = train_end + 1
val_end = data.shape[0]
data_train = data[np.arange(train_start, train_end), :]
data_val = data[np.arange(val_start, val_end), :]

# Build X and y
X_train = data_train[:, 1:]
y_train = data_train[:, 0]
X_val = data_val[:, 1:]
y_val = data_val[:, 0]

y_train = y_train.reshape(y_train.shape[0], 1)
y_val = y_val.reshape(y_val.shape[0], 1)

# Scale data
scaler_X = MinMaxScaler(feature_range=(-1, 1))
scaler_y = MinMaxScaler(feature_range=(0, 1))
scaler_X.fit(data[:,1:])
scaler_y.fit(data[:,0].reshape(data[:,0].shape[0],1))
X_train = scaler_X.transform(X_train)
X_val = scaler_X.transform(X_val)
y_train = scaler_y.transform(y_train)
y_val = scaler_y.transform(y_val)

y_train = y_train.reshape(y_train.shape[0])

# Number of stocks in training data
n_stocks = X_train.shape[1]

# Neurons
n_neurons_1 = 1024
n_neurons_2 = 512
n_neurons_3 = 256
n_neurons_4 = 128

# Session
net = tf.InteractiveSession()

# Placeholder
X = tf.placeholder(dtype=tf.float32, shape=[None, n_stocks])
Y = tf.placeholder(dtype=tf.float32, shape=[None])

# Initializers
sigma = 1
weight_initializer = tf.variance_scaling_initializer(mode="fan_avg", distribution="uniform", scale=sigma)
bias_initializer = tf.zeros_initializer()

# Hidden weights
W_hidden_1 = tf.Variable(weight_initializer([n_stocks, n_neurons_1]))
bias_hidden_1 = tf.Variable(bias_initializer([n_neurons_1]))
W_hidden_2 = tf.Variable(weight_initializer([n_neurons_1, n_neurons_2]))
bias_hidden_2 = tf.Variable(bias_initializer([n_neurons_2]))
W_hidden_3 = tf.Variable(weight_initializer([n_neurons_2, n_neurons_3]))
bias_hidden_3 = tf.Variable(bias_initializer([n_neurons_3]))
W_hidden_4 = tf.Variable(weight_initializer([n_neurons_3, n_neurons_4]))
bias_hidden_4 = tf.Variable(bias_initializer([n_neurons_4]))

# Output weights
W_out = tf.Variable(weight_initializer([n_neurons_4, 1]))
bias_out = tf.Variable(bias_initializer([1]))

# Hidden layer
hidden_1 = tf.nn.relu(tf.add(tf.matmul(X, W_hidden_1), bias_hidden_1))
hidden_2 = tf.nn.relu(tf.add(tf.matmul(hidden_1, W_hidden_2), bias_hidden_2))
hidden_3 = tf.nn.relu(tf.add(tf.matmul(hidden_2, W_hidden_3), bias_hidden_3))
hidden_4 = tf.nn.relu(tf.add(tf.matmul(hidden_3, W_hidden_4), bias_hidden_4))

# Output layer (transpose!)
out = tf.transpose(tf.add(tf.matmul(hidden_4, W_out), bias_out))

# Cost function
mse = tf.reduce_mean(tf.squared_difference(out, Y))

# Optimizer
opt = tf.train.AdamOptimizer().minimize(mse)

# Init
net.run(tf.global_variables_initializer())


# Fit neural net
batch_size = 256
mse_train = []
mse_val = []

# Run
epochs = 5
X_axis = [i for i in range(y_val.shape[0])]


print(X_train.shape)
print(y_train.shape)
print(len(y_train) // batch_size)

for e in range(epochs):

    # Shuffle training data
    shuffle_indices = np.random.permutation(np.arange(len(y_train)))
    X_train = X_train[shuffle_indices]
    y_train = y_train[shuffle_indices]

    # Minibatch training
    for i in range(0, len(y_train) // batch_size):
        start = i * batch_size
        batch_x = X_train[start:start + batch_size]
        batch_y = y_train[start:start + batch_size]
        # Run optimizer with batch
        net.run(opt, feed_dict={X: batch_x, Y: batch_y})

        # Show progress
        if np.mod(i, 50) == 0:
            # MSE train and val
            mse_train.append(net.run(mse, feed_dict={X: X_train, Y: y_train}))
            print('MSE Train: ', mse_train[-1])
            mse_val.append(net.run(mse, feed_dict={X: X_val, Y: y_val.reshape((y_val.shape[0]))}))
            print('MSE val: ', mse_val[-1])
            # Prediction
            pred = net.run(out, feed_dict={X: X_val})
            
            
            plt.figure()
            plt.plot(X_axis, scaler_y.inverse_transform(y_val), 'r')
            plt.plot(X_axis, scaler_y.inverse_transform(pred.T), 'b')
            plt.title('Epoch ' + str(e) + ', Batch ' + str(i))
            plt.show()
