from AlgorithmImports import *
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()

class TensorFlowNeuralNetworkAlgorithm(QCAlgorithm):

    def Initialize(self):
        #test 2018 to 2023
        self.SetStartDate(2020, 2, 15)  # Set Start Date
        self.SetEndDate(2020, 4, 15) # Set End Date
        self.SetCash(100000)  # Set Strategy Cash

        spy = self.AddEquity("SPY", Resolution.Daily, fillForward=False) # Add Equity
        self.symbols = [spy.Symbol] # potential trading symbols pool (in this algorithm there is only 1). 
        #!!!Optimize for returns (7 is norm)
        self.lookback = 7# number of previous days for training
        self.lookbackTwo = 2 #!Second set lockback periods

        # Feed in 100 days of daily data before the start date
        self.SetWarmUp(timedelta(days=14), Resolution.Daily)
        # Feed in data for 100 trading days before the start date
        self.SetWarmUp(14, Resolution.Daily)

        #Benchmark
        self.SetBenchmark("SPY")
        self.lastBenchmarkValue = None
        self.BenchmarkPerformance = self.Portfolio.TotalPortfolioValue



        #Checks if spy is traded on the day, if so proccesses NetTrain every 60mins during trading hours, Trade function initiated at end of train.
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.Every(timedelta(minutes=60)), Action(self.NetTrain)) #Lower number is better for more up to date trades(60 is good)
        #self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 10), Action(self.NetTrain))#for testing speed purposes, change resolution to daily




        
    def add_layer(self, inputs, in_size, out_size, activation_function=None):
        # add one more layer and return the output of this layer
        # this is one NN with only one hidden layer
        Weights = tf.Variable(tf.random_normal([in_size, out_size]))
        biases = tf.Variable(tf.zeros([1, out_size]) + 0.1)
        Wx_plus_b = tf.matmul(inputs, Weights) + biases
        if activation_function is None:
            outputs = Wx_plus_b
        else:
            outputs = activation_function(Wx_plus_b)
        return outputs
    

    def NetTrain(self):
        #return if warmup period in initialize is not complete
        if self.IsWarmingUp: return
        
        # Daily historical data is used to train the machine learning model
        history = self.History(self.symbols, self.lookback + 1, Resolution.Daily)
        
        # model: use prices_x to fit prices_y; key: symbol; value: according price
        self.prices_x, self.prices_y = {}, {}
        
        # key: symbol; values: prices for sell or buy 
        self.sell_prices, self.buy_prices = {}, {}
        
        for symbol in self.symbols:
            if not history.empty:
                # Daily historical data is used to train the machine learning model 
                # use open prices to predict the next days'
                self.prices_x[symbol] = list(history.loc[symbol.Value]['open'][:-1])
                self.prices_y[symbol] = list(history.loc[symbol.Value]['open'][1:])
                
        
        for symbol in self.symbols:
            if symbol in self.prices_x:
                # create numpy array
                x_data = np.array(self.prices_x[symbol]).astype(np.float32).reshape((-1,1))
                y_data = np.array(self.prices_y[symbol]).astype(np.float32).reshape((-1,1))
                
                # define placeholder for inputs to network
                xs = tf.placeholder(tf.float32, [None, 1])
                ys = tf.placeholder(tf.float32, [None, 1])
                
                # add hidden layer
                l1 = self.add_layer(xs, 1, 10, activation_function=tf.nn.relu)
                # add output layer
                prediction = self.add_layer(l1, 10, 1, activation_function=None)
                
                # the error between prediciton and real data
                loss = tf.reduce_mean(tf.reduce_sum(tf.square(ys - prediction),
                                     reduction_indices=[1]))
                # use gradient descent and square error
                train_step = tf.train.GradientDescentOptimizer(0.1).minimize(loss)
                
                # the following is precedure for tensorflow
                sess = tf.Session()
                
                init = tf.global_variables_initializer()
                sess.run(init)
                
                for i in range(200):
                    # training
                    sess.run(train_step, feed_dict={xs: x_data, ys: y_data})
                
            
            # predict today's price
            y_pred_final = sess.run(prediction, feed_dict = {xs: y_data})[0][-1]
            
            # get sell prices and buy prices as trading signals
            self.sell_prices[symbol] = y_pred_final - np.std(y_data)
            self.buy_prices[symbol] = y_pred_final + np.std(y_data)


            
            # Call trade---Look into better implementation Boneyard and attempts below
            #if self.Time.day % 5 == 0:
            #    self.Trade()
            #else:
                #self.Trade()

            #if symbol in :
            #    self.Trade()
            #self.Trade()

            #if data.ContainsKey("SPY"):
            #    self.Trade()

            for symbol in self.symbols:
                if not history.empty:
                    self.Trade()

            
        

    def Trade(self):
        ''' 
        Enter or exit positions based on relationship of the open price of the current bar and the prices defined by the machine learning model.
        Liquidate if the open price is below the sell price and buy if the open price is above the buy price 
        ''' 

        #Experiment with just selling position, shorting position, or leverageed short based on brokerage limits.
        for holding in self.Portfolio.Values:
            if holding.Symbol in self.CurrentSlice and self.CurrentSlice[holding.Symbol] is not None:
                    if self.CurrentSlice[holding.Symbol].Open < self.sell_prices[holding.Symbol]:
                        #self.Liquidate(holding.Symbol)
                        self.SetHoldings("SPY", -1)
                    if self.CurrentSlice[holding.Symbol].Open > self.buy_prices[holding.Symbol]:
                        self.SetHoldings(holding.Symbol, 1 / len(self.symbols))

        #Benchmark
        benchmark = self.Securities["SPY"].Close
        if self.lastBenchmarkValue is not  None:
           self.BenchmarkPerformance = self.BenchmarkPerformance * (benchmark/self.lastBenchmarkValue)
        self.lastBenchmarkValue = benchmark
        self.Plot("Strategy vs Benchmark", "Portfolio Value", self.Portfolio.TotalPortfolioValue)
        self.Plot("Strategy vs Benchmark", "Benchmark", self.BenchmarkPerformance)