import math
import pandas_datareader as web
import numpy as np
from numpy  import array
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from keras.models import Sequential
from keras.layers import Dense, LSTM, Dropout
import matplotlib.pyplot as plt
from pandas_datareader import data as pdr
import csv
from google.colab import files
import io
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.optimizers import Adam
import datetime as dt

uploaded = files.upload()
df = pd.read_csv(io.BytesIO(uploaded['AIndexEodPrices.csv']),index_col=0)
del df['S_INFO_WINDCODE'] # delete the column we don't need

# Save the first column as trade date
# Save the next five columns as data

training_set = df.iloc[:len(df)//4*3 , 1:6].values  # Use first 75% data for training
test_set = df.iloc[len(df)//4*3:, 1:6].values # Use last 25% data for testing

sc  = MinMaxScaler(feature_range=(0, 1))
training_set_scaled = sc.fit_transform(training_set)
testing_set_scaled  = sc.transform(test_set) 
n_future=1
n_past =30

def data_split(data, n_future, n_past):
    X = []
    y = []
    for i in range(n_past, len(data) - n_future +1):
      X.append(data[i - n_past:i, 0:5])
      y.append(data[i + n_future - 1:i + n_future,3])
    return np.array(X), np.array(y)

X_train, y_train = data_split(training_set_scaled, n_future,n_past)
X_train          = X_train.reshape(X_train.shape[0], X_train.shape[1], 5)

X_test, y_test   = data_split(testing_set_scaled, n_future,n_past)
X_test           = X_test.reshape(X_test.shape[0], X_test.shape[1], 5)

model1 = keras.models.Sequential([
     layers.Input((n_past, 5)),
     layers.Reshape((n_past, 5, 1)),
     layers.Conv2D(filters=64,
                           kernel_size=3,
                           strides=1,
                           padding="same",
                           activation="relu"),
    layers.MaxPooling2D(pool_size=2, strides=1, padding="same"),
    layers.Dropout(0.3),
    layers.Reshape((n_past, -1)),
    layers.LSTM(128, return_sequences=True),
    layers.LSTM(64, return_sequences=False),
    layers.Dense(32, activation="relu"),
    layers.Dense(y_train.shape[1])
])

model1.compile(optimizer='adam', loss='mean_squared_error')
history1 = model1.fit(X_train, y_train, batch_size=1, epochs=10)

predicted_stock_price = model1.predict(X_test)                      
scaled_prediction = predicted_stock_price * (sc.data_max_[2] - sc.data_min_[2]) + sc.data_min_[2]
scaled_true = y_test * (sc.data_max_[2] - sc.data_min_[2]) + sc.data_min_[2]

train = df[:len(df)//4*3+30]
valid = df[len(df)//4*3+30:]

valid['Predictions'] = scaled_prediction
#Visualize the data
plt.figure(figsize=(160,80))
plt.title('Model')
plt.plot(train['S_DQ_CLOSE'])
plt.plot(valid[['S_DQ_CLOSE', 'Predictions']])
plt.legend(['Train', 'Val', 'Predictions'], loc='lower right')
plt.show()

def returns(days,df,model):
  balance=[0]
  signal = pd.DataFrame()
  data=df.iloc[:,1:6].values
  scaler  = MinMaxScaler(feature_range=(0, 1))
  scaled_data=scaler.fit_transform(data)
  prediction_hist=[]
  profit = []
  datelist_train = list(df['TRADE_DT'][-days - 1:])
  datelist_train = [dt.datetime.strptime(str(date), '%Y%m%d').date() for date in datelist_train]

  # return strategy
  for i in range(days):
    x=[]
    x.append(scaled_data[i:i+30, 0:5])
    x=np.array(x)
    x=x.reshape(x.shape[0],x.shape[1],5)
    prediction=model.predict(x, verbose = False)
    prediction = prediction * (scaler.data_max_[2] - scaler.data_min_[2]) + scaler.data_min_[2]
    prediction_hist.append(prediction)
    if prediction>1.015*data[i+29][3]:
      profit.append(data[i+30][3]-data[i+29][3])
      balance.append((profit[i]+balance[i]))
      tmp = pd.DataFrame(["Long"])
      tmp.index = [datelist_train[i + 1]]
      signal = signal.append(tmp)
    elif prediction < 0.985*data[i+29][3]:
      profit.append(data[i+29][3]-data[i+30][3])
      balance.append(profit[i] + balance[i])
      tmp = pd.DataFrame(["Short"])
      tmp.index = [datelist_train[i + 1]]
      signal = signal.append(tmp)
    else:
      profit.append(0)
      balance.append(balance[i])
    
    
    # long strategy
  stock_price = df.iloc[29:,4:5].values
  long_return = [0]
  for i in range(days):
    long_return.append(long_return[i] + stock_price[i + 1] - stock_price[i])
  #print(profit)
  #print(prediction_hist)
  #print(stock_price)
  #plt.plot(np.array(prediction_hist).reshape(days,1),c='blue')
  #plt.plot(stock_price[1:,],c='red')
  signal.columns = ["Action"]
  balance = pd.DataFrame(balance)
  balance.index = datelist_train
  plt.plot(balance,c='blue')
  long_return = pd.DataFrame(long_return)
  long_return.index = datelist_train
  plt.plot(long_return, c = 'red')
  return signal, balance[days:]

returns(365,df[-395:],model1)