# Data format and transformations
import numpy as np
import pandas as pd

# Importing data
from yahoo_finance_api2 import share
from yahoo_finance_api2.exceptions import YahooFinanceError

# Technical analysis
import ta

# Machine Learning
from sklearn.linear_model import LinearRegression
from sklearn import preprocessing, model_selection, svm


HISTORY_LENGTH=100
PREDICTIONS_DAYS=[1,5,10]
PREDICTION_PREFIX='prediction'

def initHistory(share_name):
   """
   This function generates a history from a share name

   Parameters
   ----------
   share_name: str
      Share name in yahoo format
   
   Returns
   -------
   A pandas with last share values and TA
   """
   df = __getLastPrices(share_name, HISTORY_LENGTH)
   if df is not None:
      df = ta.add_all_ta_features(
         df, "open", "high", "low", "close", "volume", fillna=True)
      for day in PREDICTIONS_DAYS:
         column_name = PREDICTION_PREFIX+str(day)
         df[column_name]=np.nan
      for i in range(0,len(df)):
         makePredictions(df, i)
   return df


def readHistory(file_name):
   return pd.read_csv(file_name).head(HISTORY_LENGTH)


def writeHistory(file_name, history):
   # TODO: saved data is not the same as the read one.
   history.to_csv(file_name, index=False)


def updateHistory(history, share_name):
   entry = __getLastPrices(share_name)
   if entry['timestamp'][0] not in history['timestamp']:
      df = history.append(entry, ignore_index=True)
      df = ta.add_all_ta_features(
          df, "open", "high", "low", "close", "volume", fillna=True)
      return makePredictions(df, len(history)-1)
   else:
      return None


def getPredictionError(history, days):
   column_name = PREDICTION_PREFIX + str(days)
   df = pd.DataFrame()
   df['prediction'] = history[column_name]
   df['real'] = history['close'].shift(-days)
   df.dropna(inplace=True)

   return ((df['prediction']-df['real'])/df['real']).abs().median()


def __getLastPrices(share_name, days = 1):
   """
   This function returns last values of a share

   Parameters
   ----------
   share_name: str
      Share name in yahoo format
   
   Returns
   -------
   A pandas with last share values and TA
   """
   my_share = share.Share(share_name)
   symbol_data = None
   try:
      symbol_data = my_share.get_historical(
         share.PERIOD_TYPE_DAY,
         days,
         share.FREQUENCY_TYPE_DAY,
         1)
   except YahooFinanceError as e:
      print(e.message)
   finally:
      if symbol_data is None:
         return None
      else:
         df = pd.DataFrame(symbol_data)
         return df

def __removePredictions(history):
   df = history
   for day in PREDICTIONS_DAYS:
      column_name = PREDICTION_PREFIX+str(day)
      if column_name in df.columns:
         df = df.drop(column_name, 1)
   return df

def __makePrediction(history, element, days=1):
   df = history.drop(['timestamp'],1)
   df = __removePredictions(df)
   df['future']= df['close'].shift(-days)
   df_dropped = df.dropna(inplace=False)
   
   X = np.array(df_dropped.drop('future', 1))
   Y = np.array(df_dropped['future'])
   X = preprocessing.scale(X)

   # Split data into train and test
   X_train, X_test, Y_train, Y_test = model_selection.train_test_split(
      X, Y, test_size=0.2)

   # Perform the regresion on the training data
   clf = LinearRegression()
   clf.fit(X_train, Y_train)
   to_predict = preprocessing.scale(
       np.array(df.drop('future', 1)))[element:element+1]
   prediction = clf.predict(to_predict)
   return prediction

def makePredictions(history, element):
   for day in PREDICTIONS_DAYS:
      column_name = PREDICTION_PREFIX+str(day)
      history.at[element, column_name] = __makePrediction(
          history, element=element, days=day)
   return history

if __name__ == "__main__":
   
   import argparse
   parser = argparse.ArgumentParser(description='Generates a prediction for a share symbol.')
   parser.add_argument('historic_file', help='historic filename')
   parser.add_argument('symbol', help='Share symbol name')
   parser.add_argument('-i', action='store_true', default=False,
                     dest='init',
                     help='Initialise the historic file')

   try:
      args = parser.parse_args()
      if args.init:
         history = initHistory(args.symbol)
      else:
         history = readHistory(args.historic_file)
         history = updateHistory(history, args.symbol)

      writeHistory(args.historic_file, history)
   except:
      pass

   import matplotlib.pyplot as plt
   from matplotlib import style
   style.use('ggplot')
   #df = pd.DataFrame(predictions)
   #df = history['close'].append(df, ignore_index=True)
   #plt.plot(df[:-N])
   #plt.plot(df[-N:])

   #plt.show()