import numpy as np
import tensorflow as tf

class TensorFlowNeuralNetwork(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1)  # Set Start Date
        self.SetEndDate(2018, 1, 1) # Set End Date
        
        self.SetCash(100000)  # Set Strategy Cash
        spy = self.AddEquity("SPY", Resolution.Minute)
        
        self.symbols = [spy.Symbol]
        self.lookback = 30
        
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.AfterMarketOpen("SPY", 28), Action(self.NetTrain))
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.AfterMarketOpen("SPY", 30), Action(self.Trade))

    def OnData(self, data):
        self.data = data
        
    def add_layer(self, inputs, in_size, out_size, activation_function=None):
        # add one more layer and return the output of this 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):
        history = self.History(self.symbols, self.lookback + 1, Resolution.Daily)
        
        self.prices_x, self.prices_y = {}, {}
        self.sell_prices, self.buy_prices = {}, {}
        
        for symbol in self.symbols:
            if not history.empty:
                self.prices_x[symbol.Value] = list(history.loc[symbol.Value]['open'][:-1])
                self.prices_y[symbol.Value] = list(history.loc[symbol.Value]['open'][1:])
        
        for symbol in self.symbols:
            if symbol.Value in self.prices_x:
                # create data
                x_data = np.array(self.prices_x[symbol.Value]).astype(np.float32).reshape((-1,1))
                y_data = np.array(self.prices_y[symbol.Value]).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]))
                train_step = tf.train.GradientDescentOptimizer(0.1).minimize(loss)
                
                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]
            # self.Debug(f'pred price: {y_pred_final}')
            
            self.sell_prices[symbol.Value] = y_pred_final - np.std(y_data)
            self.buy_prices[symbol.Value] = y_pred_final + np.std(y_data)
        
    def Trade(self):
        # Trending strategy
        for i in self.Portfolio.Values:
            # liquidate
            if self.data[i.Symbol.Value].Open < self.sell_prices[i.Symbol.Value] and i.Invested:
                self.Liquidate(i.Symbol)
            
            # buy
            if self.data[i.Symbol.Value].Open > self.buy_prices[i.Symbol.Value] and not i.Invested:
                self.SetHoldings(i.Symbol, 1 / len(self.symbols))