SEED = 2019
# TYPE = 'classification'
TYPE = 'regression'
# FEATURE = 'value'
FEATURE = 'change'

from math import sqrt
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
from datetime import timedelta
from sklearn.preprocessing import MinMaxScaler, LabelEncoder
from sklearn.metrics import mean_squared_error, accuracy_score, classification_report, confusion_matrix
from keras.models import Sequential
from keras.layers import Dense, LSTM, Dropout, BatchNormalization
from keras.callbacks import ModelCheckpoint, EarlyStopping, TensorBoard
np.random.seed(SEED)
from tensorflow import set_random_seed
set_random_seed(SEED)


TEST_START_TIME = pd.datetime(2014, 1, 1)
TEST_END_TIME = pd.datetime(2015, 12, 31)
TRAIN_START_TIME = pd.datetime(2009, 1, 1)
TRAIN_END_TIME = pd.datetime(2013, 12, 31)

# convert series to supervised learning
def series_to_supervised(data, n_in=1, n_out=1, dropnan=True):
	n_vars = 1 if type(data) is list else data.shape[1]
	df = pd.DataFrame(data)
	cols, names = list(), list()
	# input sequence (t-n, ... t-1)
	for i in range(n_in, 0, -1):
		cols.append(df.shift(i))
		names += [('var%d(t-%d)' % (j+1, i)) for j in range(n_vars)]
	# forecast sequence (t, t+1, ... t+n)
	for i in range(0, n_out):
		cols.append(df.shift(-i))
		if i == 0:
			names += [('var%d(t)' % (j+1)) for j in range(n_vars)]
		else:
			names += [('var%d(t+%d)' % (j+1, i)) for j in range(n_vars)]
	# put it all together
	agg = pd.concat(cols, axis=1)
	agg.columns = names
	# drop rows with NaN values
	if dropnan:
		agg.dropna(inplace=True)
	return agg

def transform_same_date_range(df1, df2):
	date_start = max(df1['Date'].min(), df2['Date'].min())
	date_end = min(df1['Date'].max(), df2['Date'].max())
	df1 = df1[(df1['Date'] >= date_start) & (df1['Date'] <= date_end)]
	df2 = df2[(df2['Date'] >= date_start) & (df2['Date'] <= date_end)]
	print('{} to {}'.format(date_start, date_end))
	return df1, df2

def get_rmse(y_true, y_pred):
	return mean_squared_error(y_true, y_pred) ** 0.5

def plot_predict_result(model_name, y_predict):
	DJI_plot = combine[combine.index >= TEST_START_TIME][['DJI_value']]
	if FEATURE == 'change':
		DJI_plot['change_predict'] = y_predict.cumsum()
		DJI_plot['DJI_predict'] = DJI_plot['change_predict'] + combine[combine.index < TEST_START_TIME]['DJI_value'].iloc[-1]
		title = 'Dow Index Prediction Trend By Change ({})'.format(model_name)
	else:
		DJI_plot['DJI_predict'] = y_predict
		title = 'Dow Index Prediction Trend By Value ({})'.format(model_name)

	plt.figure(figsize = (12,8))
	plt.plot(combine[['DJI_value']], label = 'Real')
	plt.plot(DJI_plot[['DJI_predict']], label = 'Predict')
	plt.legend()
	plt.title(title, size = 18)
	plt.savefig('images/' + title.replace(' ', '_').replace('(', '_').replace(')', '_'), dpi=200, bbox_inches="tight")


def get_result(test_y, test_pred, model_name):
	if TYPE == 'regression':
		if FEATURE == 'change':
			print('TEST RMSE: {}'.format(get_rmse(test_y.cumsum()+ combine[combine.index < TEST_START_TIME]['DJI_value'].iloc[-1], test_pred.cumsum() + combine[combine.index < TEST_START_TIME]['DJI_value'].iloc[-1])))
		else:
			print('TEST RMSE: {}'.format(get_rmse(test_y, test_pred)))
	else:
		print('Test Accuracy: {}'.format(accuracy_score(test_y, test_pred)))
		print(classification_report(test_y, test_pred))
		print(confusion_matrix(test_y, test_pred))

		pd.DataFrame(confusion_matrix(test_y, test_pred)).to_csv('results/cm_{}_{}_{}.csv'.format(TYPE, FEATURE, model_name))
		pd.DataFrame(classification_report(test_y, test_pred, output_dict=True)).transpose().to_csv('results/cr_{}_{}_{}.csv'.format(TYPE, FEATURE, model_name))

# load DJI dataset
DJI = pd.read_csv('data/DJI_yahoo.csv', parse_dates=['Date'], infer_datetime_format='%d/%m/%Y')
DJI.describe()

DJI['DJI_value'] = DJI['Adj Close']
DJI = DJI[['Date', 'DJI_value']]
combine = DJI

# # merge FX and DJI dataset
# FX_df = pd.read_csv('data/DTWEXB.csv', parse_dates=['DATE'], na_values='.', infer_datetime_format='%Y-%m-%d').ffill().rename(columns = {'DTWEXB':'FX_Weighted_value', 'DATE': 'Date'})
# DJI, FX_df = transform_same_date_range(DJI, FX_df)
# combine = DJI.merge(FX_df, how = 'left')
#
# # load consumer price index data
# CPI = pd.read_csv('data/CPI.csv')
# CPI.columns = ['Year'] + list(range(1, 13))
# CPI = pd.melt(CPI, id_vars=['Year'], value_vars=list(range(1, 13)), var_name='Month', value_name='CPI_value').sort_values(['Year', 'Month'])
# CPI['Date'] = pd.to_datetime(CPI['Year'].astype(str) + '/' + CPI['Month'].astype(str), format = '%Y/%m')
# CPI = CPI.dropna()
# CPI = pd.DataFrame(data = {'Date': pd.date_range(CPI['Date'].min(), CPI['Date'].max() + timedelta(30))}).merge(CPI, how = 'left').ffill()
#
# # merge CPI into the combined dataset
# combine, CPI = transform_same_date_range(combine, CPI[['Date', 'CPI_value']])
# combine = combine.merge(CPI, how = 'left')
# combine.isnull().sum()
#
# # Load interest rate data
# IR = pd.read_csv('data/IR.csv', parse_dates=['DATE'], infer_datetime_format='%Y-%m-%d')
# IR.columns = ['Date', 'IR_value']
# IR = pd.DataFrame(data = {'Date': pd.date_range(IR['Date'].min(), IR['Date'].max() + timedelta(29))}).merge(IR, how = 'left').ffill()
#
# # merge CPI into the combined dataset
# combine, IR = transform_same_date_range(combine, IR)
# combine = combine.merge(IR, how = 'left')
# combine.isnull().sum()
#
# # load S&P 500 index
# SP500 = pd.read_csv('data/SP500.csv', parse_dates=['Date'], infer_datetime_format='%Y-%m-%d')
# SP500['SP500_value'] = SP500['Adj Close']
# SP500 = SP500[['Date', 'SP500_value']]
#
# # merge SP500 into the combined dataset
# combine, SP500 = transform_same_date_range(combine, SP500)
# combine = combine.merge(SP500, how = 'left')
# combine.isnull().sum()
#
# News = pd.read_csv("data/combined_stock_data.csv", parse_dates=[1])[['Date', 'Label', 'Subjectivity', 'Objectivity', 'Positive', 'Negative', 'Neutral']].ffill()
# News.columns = ['Date'] + ['News_' + col for col in ['Label', 'Subjectivity', 'Objectivity', 'Positive', 'Negative', 'Neutral']]
# combine, News = transform_same_date_range(combine, News)
# combine = combine.merge(News, how = 'left')
# combine.isnull().sum()

combine = combine.set_index(['Date'])
# Convert to stationary
for col in combine.columns:
	if col.endswith('value'):
		combine[col.replace('value', 'change')] = combine[col] - combine[col].shift(1)
combine.describe()

combine = combine[(combine.index >= TRAIN_START_TIME) & (combine.index <= TEST_END_TIME)]
combine.head()
combine.tail()

combine['DJI_direction'] = (combine['DJI_change'] > 0).astype(int)
values = combine[[col for col in combine.columns if col.endswith(FEATURE) or col.startswith('News')]]

LAG_DAYS = 10
attributes = values.columns
for day in range(1, LAG_DAYS + 1):
	for attribute in attributes:
		values['{}(t-{})'.format(attribute, day)] = values[attribute].shift(day)

values = values.iloc[LAG_DAYS:]
if TYPE == 'regression':
	if FEATURE == 'change':
		y = combine['DJI_change'].iloc[LAG_DAYS:]
	else:
		y = combine['DJI_value'].iloc[LAG_DAYS:]
else:
	y = combine['DJI_direction'].iloc[LAG_DAYS:]
flatten = lambda l: [item for sublist in l for item in sublist]
X = values[flatten([['{}(t-{})'.format(attribute, i) for i in range(1, 11)] for attribute in attributes])]
y.shape
X.shape

train_X, test_X = X.loc[:TEST_START_TIME], X.loc[TEST_START_TIME:]
train_y, test_y = y.loc[:TEST_START_TIME], y.loc[TEST_START_TIME:]


# Scale
scaler = MinMaxScaler(feature_range=(0, 1))
train_X[train_X.columns] = scaler.fit_transform(train_X)
test_X[train_X.columns] = scaler.transform(test_X)

train_df = train_X.copy()
test_df = test_X.copy()
if TYPE == 'regression':
	if FEATURE == 'change':
		train_df['DJI_change'] = train_y
		test_df['DJI_change'] = test_y
	else:
		train_df['DJI_value'] = train_y
		test_df['DJI_value'] = test_y
else:
	train_df['DJI_direction'] = train_y
	test_df['DJI_direction'] = test_y

train_df.to_csv('../data/dow_train_{}_{}.csv'.format(TYPE, FEATURE))
test_df.to_csv('../data/dow_test_{}_{}.csv'.format(TYPE, FEATURE))
