
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.preprocessing import RobustScaler
import ta
from datetime import timedelta
import sys
import myenv

# Function to prepare and compute technical indicator and prepare to neuron network by scenario
# Input: df -- pandas dataframe
#        scenario -- scenario for generate tech indicators. All is for all indicators from ta library. dpIndicators for indicators from me dp
def prepareIndicator(df, scenario):
    if scenario == 'all':
        df = ta.add_all_ta_features(df, open="Open", high="High", low="Low", close="Close", volume="Volume", fillna=True)
    if scenario == 'dpIndicators':
        #Oscilator rsi from ta library
        rsiDf = ta.momentum.RSIIndicator(df.Close, 14, True)
        df['rsi'] = rsiDf.rsi()
        # Oscilator stochactic RSI
        stochDf = ta.momentum.StochRSIIndicator(df.Close, 14, 3, 3, True)
        df['stochRsi'] = stochDf.stochrsi()
        df['stochRsi_D'] = stochDf.stochrsi_d()
        df['stochRsi_K'] = stochDf.stochrsi_k()
        # Oscilator stochacstic classic
        stochNormal = ta.momentum.StochasticOscillator(df.High, df.Low, df.Close, 14, 3, True)
        df['stochOsc'] = stochNormal.stoch()
        df['stochSignal'] = stochNormal.stoch_signal()
        #trend indicator Moving Average Convergence Divergence
        macd = ta.trend.MACD(df.Close, 26, 12, 9, True)
        df['macd'] = macd.macd()
        #trend indicator Simple Moving Acerage
        sma = ta.trend.SMAIndicator(df.Close, 10, True)
        df['sma'] = sma.sma_indicator()
        #volume indicator Money Flow Index (MFI)
        mfi = ta.volume.MFIIndicator(df.High, df.Low, df.Close, df.Volume, 14, True)
        df['mfi'] = mfi.money_flow_index()


    # Delete all input data except Close value
    df.drop(['Open', 'High', 'Low', 'Adj Close', 'Volume'], axis=1, inplace=True)
    return df


df = pd.read_csv(myenv.SOURCE_DATASET)

df['Date'] = pd.to_datetime(df.Date)
df.set_index('Date', inplace=True)
# change scenario for generate prepare data. Example: all, dpIndicators
df = prepareIndicator(df, myenv.TECH_SCENARIO)

#scale data to <-1,1> for neural network
scaler = RobustScaler()
df = pd.DataFrame(scaler.fit_transform(df), columns=df.columns, index=df.index)

df.to_pickle(myenv.PREPARED_DATASET)  

# script to load from pickle format
#df1 = pd.read_pickle(myenv.PREPARED_DATASET) 