import math
import matplotlib
import numpy as np
from numpy.core.numerictypes import ScalarType
import pandas as pd
import seaborn as sns
import time

from datetime import date, datetime, time, timedelta
from matplotlib import pyplot as plt
from pylab import rcParams
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score
from tqdm import tqdm_notebook

# %matplotlib inline

#### Input params ##################
stk_path = "data_final.csv"
test_size = 0.2                 # proportion of dataset to be used as test set
cv_size = 0.2                   # proportion of dataset to be used as cross-validation set
Nmax = 30                       # for feature at day t, we use lags from t-1, t-2, ..., t-N as features
                                # Nmax is the maximum N we are going to test
fontsize = 14
ticklabelsize = 14
####################################

def get_preds_lin_reg(df, target_col, N, pred_min, offset):
    """
    Given a dataframe, get prediction at timestep t using values from t-1, t-2, ..., t-N.
    Inputs
        df         : dataframe with the values you want to predict. Can be of any length.
        target_col : name of the column you want to predict e.g. 'adj_close'
        N          : get prediction at timestep t using values from t-1, t-2, ..., t-N
        pred_min   : all predictions should be >= pred_min
        offset     : for df we only do predictions for df[offset:]. e.g. offset can be size of training set
    Outputs
        pred_list  : the predictions for target_col. np.array of length len(df)-offset.
    """
    # Create linear regression object
    regr = LinearRegression(fit_intercept=True)

    pred_list = []

    for i in range(offset, len(df['daily'])):
        X_train = np.array(range(len(df['daily'][i-N:i]))) # e.g. [0 1 2 3 4]
        y_train = np.array(df['daily'][i-N:i]) # e.g. [2944 3088 3226 3335 3436]
        X_train = X_train.reshape(-1, 1)     # e.g X_train = 
                                             # [[0]
                                             #  [1]
                                             #  [2]
                                             #  [3]
                                             #  [4]]
        # X_train = np.c_[np.ones(N), X_train]              # add a column
        y_train = y_train.reshape(-1, 1)
    #     print X_train.shape
    #     print y_train.shape
    #     print 'X_train = \n' + str(X_train)
    #     print 'y_train = \n' + str(y_train)
        regr.fit(X_train, y_train)            # Train the model
        pred = regr.predict(np.array(N).reshape(1,-1))
    
        pred_list.append(pred[0][0])  # Predict the footfall using the model
    
    # If the values are < pred_min, set it to be pred_min
    pred_list = np.array(pred_list)
    pred_list[pred_list < pred_min] = pred_min
        
    return pred_list

def get_mape(y_true, y_pred): 
    """
    Compute mean absolute percentage error (MAPE)
    """
    y_true, y_pred = np.array(y_true), np.array(y_pred)
    return np.mean(np.abs((y_true - y_pred) / y_true)) * 100

df = pd.read_csv(stk_path, sep = ",")

# Convert Date column to datetime
# Change all column headings to be lower case, and remove spacing
df.columns = [str(x).lower().replace(' ', '_') for x in df.columns]

df.head()

# Plot adjusted close over time
rcParams['figure.figsize'] = 10, 8 # width 10, height 8

ax = df.plot(x='day', y='daily', style='b-', grid=True)
ax.set_xlabel("Day")
ax.set_ylabel("Price")

# Get sizes of each of the datasets
num_cv = int(cv_size*len(df))
num_test = int(test_size*len(df))
num_train = len(df) - num_cv - num_test
print("num_train = " + str(num_train))
print("num_cv = " + str(num_cv))
print("num_test = " + str(num_test))

# Split into train, cv, and test
train = df[:num_train].copy()
cv = df[num_train:num_train+num_cv].copy()
train_cv = df[:num_train+num_cv].copy()
test = df[num_train+num_cv:].copy()
print("train.shape = " + str(train.shape))
print("cv.shape = " + str(cv.shape))
print("train_cv.shape = " + str(train_cv.shape))
print("test.shape = " + str(test.shape))

# Plot adjusted close over time
rcParams['figure.figsize'] = 10, 8 # width 10, height 8

ax = train.plot(x='day', y='daily', style='b-', grid=True)
ax = cv.plot(x='day', y='daily', style='y-', grid=True, ax=ax)
ax = test.plot(x='day', y='daily', style='g-', grid=True, ax=ax)
ax.legend(['train', 'dev', 'test'])
ax.set_xlabel("day")
ax.set_ylabel("price")

RMSE = []
R2 = []
mape = []
for N in range(1, Nmax+1): # N is no. of samples to use to predict the next value
    est_list = get_preds_lin_reg(train_cv, 'daily', N, 0, num_train)
    
    cv.loc[:, 'est' + '_N' + str(N)] = est_list
    RMSE.append(math.sqrt(mean_squared_error(est_list, cv['daily'])))
    R2.append(r2_score(cv['daily'], est_list))
    mape.append(get_mape(cv['daily'], est_list))
print('RMSE = ' + str(RMSE))
print('R2 = ' + str(R2))
print('MAPE = ' + str(mape))
cv.head()

# Plot RMSE versus N
matplotlib.rcParams.update({'font.size': 14})
plt.figure(figsize=(12, 8), dpi=80)
plt.plot(range(1, Nmax+1), RMSE, 'x-')
plt.grid()
plt.xlabel('N')
plt.ylabel('RMSE')
plt.xlim([2, 30])

# Plot R2 versus N. Note for R2 larger better. 
matplotlib.rcParams.update({'font.size': 14})
plt.figure(figsize=(12, 8), dpi=80)
plt.plot(range(1, Nmax+1), R2, 'x-')
plt.grid()
plt.xlabel('N')
plt.ylabel('R2')

# Plot MAPE versus N. Note for MAPE smaller better. 
plt.figure(figsize=(12, 8), dpi=80)
plt.plot(range(1, Nmax+1), mape, 'x-')
plt.grid()
plt.xlabel('N')
plt.ylabel('MAPE')

# Set optimum N
N_opt = 5

# Specify the day you are interested in
day = 200

# Specify the maximum N you want to plot (If Nmax2 is too large it gets very cluttered) 
Nmax2 = 5

df_temp = cv[cv['day'] <= day]
plt.figure(figsize=(12, 8), dpi=80)
plt.plot(range(1,Nmax2+2), df_temp[-Nmax2-1:]['daily'], 'bx-')
plt.plot(Nmax2+1, df_temp[-1:]['daily'], 'ys-')
legend_list = ['daily', 'actual_value']

# Plot the linear regression lines and the predictions
color_list = ['r', 'g', 'k', 'y', 'm', 'c', '0.75']
marker_list = ['x', 'x', 'x', 'x', 'x', 'x', 'x']
regr = LinearRegression(fit_intercept=True) # Create linear regression object
for N in range(5, Nmax2+1):
    # Plot the linear regression lines
    X_train = np.array(range(len(df_temp['daily'][-N-1:-1]))) # e.g. [0 1 2 3 4]
    y_train = np.array(df_temp['daily'][-N-1:-1]) # e.g. [2944 3088 3226 3335 3436]
    X_train = X_train.reshape(-1, 1)     
    y_train = y_train.reshape(-1, 1)
    regr.fit(X_train, y_train)            # Train the model
    y_est = regr.predict(X_train)         # Get linear regression line
    plt.plot(range(Nmax2+1-N,Nmax2+2), 
             np.concatenate((y_est, np.array(df_temp['est_N'+str(N)][-1:]).reshape(-1,1))),
             color=color_list[N%len(color_list)], 
             marker=marker_list[N%len(marker_list)])
    legend_list.append('est_N'+str(N)+'_lr')
    
    # Plot the predictions
    plt.plot(Nmax2+1, 
             df_temp['est_N'+str(N)][-1:], 
             color=color_list[N%len(color_list)], 
             marker='o')
    legend_list.append('est_N'+str(N))
    

plt.grid()
plt.xlabel('timestep')
plt.ylabel('USD')
plt.legend(legend_list, bbox_to_anchor=(1.05, 1))
matplotlib.rcParams.update({'font.size': fontsize})

# Plot adjusted close over time
rcParams['figure.figsize'] = 10, 8 # width 10, height 8

ax = train.plot(x='day', y='daily', style='b-', grid=True)
ax = cv.plot(x='day', y='daily', style='y-', grid=True, ax=ax)
ax = test.plot(x='day', y='daily', style='g-', grid=True, ax=ax)
ax = cv.plot(x='day', y='est_N1', style='r-', grid=True, ax=ax)
ax = cv.plot(x='day', y='est_N5', style='m-', grid=True, ax=ax)
ax.legend(['train', 'dev', 'test', 'predictions with N=1', 'predictions with N=5'])
ax.set_xlabel("day")
ax.set_ylabel("price")

# Plot adjusted close over time
rcParams['figure.figsize'] = 10, 8 # width 10, height 8

ax = train.plot(x='day', y='daily', style='bx-', grid=True)
ax = cv.plot(x='day', y='daily', style='yx-', grid=True, ax=ax)
ax = test.plot(x='day', y='daily', style='gx-', grid=True, ax=ax)
ax = cv.plot(x='day', y='est_N1', style='rx-', grid=True, ax=ax)
ax = cv.plot(x='day', y='est_N5', style='mx-', grid=True, ax=ax)
ax.legend(['train', 'dev', 'test', 'predictions with N=1', 'predictions with N=5'])
ax.set_xlabel("day")
ax.set_ylabel("Price")
# ax.set_xlim([date(2017, 11, 1), date(2017, 12, 30)])
# ax.set_ylim([127, 137])
ax.set_title('Zoom in to dev set')

est_list = get_preds_lin_reg(df, 'daily', N_opt, 0, num_train+num_cv)
test.loc[:, 'est' + '_N' + str(N_opt)] = est_list
print("RMSE = %0.3f" % math.sqrt(mean_squared_error(est_list, test['daily'])))
print("R2 = %0.3f" % r2_score(test['daily'], est_list))
print("MAPE = %0.3f%%" % get_mape(test['daily'], est_list))
test.head()

# Plot adjusted close over time
rcParams['figure.figsize'] = 10, 8 # width 10, height 8

ax = train.plot(x='day', y='daily', style='b-', grid=True)
ax = cv.plot(x='day', y='daily', style='y-', grid=True, ax=ax)
ax = test.plot(x='day', y='daily', style='g-', grid=True, ax=ax)
ax = test.plot(x='day', y='est_N5', style='r-', grid=True, ax=ax)
ax.legend(['train', 'dev', 'test', 'predictions with N_opt=5'])
ax.set_xlabel("day")
ax.set_ylabel("price")

# Plot adjusted close over time
rcParams['figure.figsize'] = 10, 8 # width 10, height 8

ax = train.plot(x='day', y='daily', style='bx-', grid=True)
ax = cv.plot(x='day', y='daily', style='yx-', grid=True, ax=ax)
ax = test.plot(x='day', y='daily', style='gx-', grid=True, ax=ax)
ax = test.plot(x='day', y='est_N5', style='rx-', grid=True, ax=ax)
ax.legend(['train', 'dev', 'test', 'predictions with N_opt=5'])
ax.set_xlabel("day")
ax.set_ylabel("price")
# ax.set_xlim([date(2018, 6, 1), date(2018, 7, 31)])
# ax.set_ylim([135, 150])
ax.set_title('Zoom in to test set')

# Plot adjusted close over time, only for test set
rcParams['figure.figsize'] = 10, 8 # width 10, height 8
matplotlib.rcParams.update({'font.size': 14})

ax = test.plot(x='day', y='daily', style='gx-', grid=True)
ax = test.plot(x='day', y='est_N5', style='rx-', grid=True, ax=ax)
ax.legend(['test', 'predictions using linear regression'], loc='upper left')
ax.set_xlabel("day")
ax.set_ylabel("Price")
#plt.show()

# from sklearn.preprocessing import StandardScaler
# data = pd.read_csv('data_final.csv')
# data.head()
# # Let’s select some features to explore more :
# df_new = data[['daily']]
# last_30_days = df_new[-30:].values

# scaler = StandardScaler()
# scaler.fit(last_30_days)
# last_30_days_scaled = scaler.transform(last_30_days)

# X_test = []
# X_test.append(last_30_days_scaled)
# X_test = np.array(X_test)

# X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1],1))

# pred_price = regr.predict(X_test)
# pred_price = scaler.inverse_transform(pred_price)

# print('Est Price: ' + str(pred_price))

# Save as 
# csv
test_lin_reg = test
test_lin_reg.to_csv("test_lin_reg.csv")