OrderTypeKeys = [
    'Market', 'Limit', 'StopMarket', 'StopLimit', 'MarketOnOpen',
    'MarketOnClose', 'OptionExercise',
]

OrderTypeCodes = dict(zip(range(len(OrderTypeKeys)), OrderTypeKeys))

from orderTypes import OrderTypeCodes
from datetime import timedelta



class MeanReversion(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019,1,1)  # Set Start Date
        self.SetEndDate(2019,3,20)
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)

        # what resolution should the data *added* to the universe be?
        self.UniverseSettings.Resolution = Resolution.Minute

      
        # Number of stock selected for check long and short conditions. 
        # This is used in the Fine Fundamental Filter to select the top 200
        # stocks ordered by Price To Sales and the bottom 200 stock by PS.
        self.long = 200
        self.short = 200
        
        self.longStocks = 0
        self.shortStocks = 0 
        
        self.tradeBarWindow = {}
        
        self.BarPeriod = timedelta(days=1)
        
        # This is the period of our sma indicators
        self.SimpleMovingAveragePeriod = 20
       
        # This is the number of consolidated bars we'll hold in symbol data for reference
        self.RollingWindowSize = 2
        
        # Holds all of our data keyed by each symbol
        self.Data = {}
        
        # List to use to add stocks every time an entry point is triggered.
        self.stocksPurchased = []
        
        self.nextRebalance = self.Time      # Initialize next rebalance time
        self.rebalance_days = 90           # Rebalance every 90 days
        
      
        # The dictionaries BollingBand and SMA20 are used to store the values for each symbol
        # of these two technical indicators. At first these dictionaries are updated with the
        # required window values and then are used in the technical setup for entry points.
        self.BollingerBand = {}
        self.SMA20 = {}
        
        # The three dictionaries below, are used to store order tickets for the market orders
        # entry positions, the stop orders as well as the limit orders for take profits
        # Using the order tickets, we can track orders and get important properties of the 
        # order object such as date, price, quantity, and also use the ticket for cancelling
        # the order for the symbol when the opposite order is filled.
        self.orders = {}
        self.stopOrders = {}
        self.profitOrders = {}
        
        self.stateData = {}
        
        self.AddUniverse(self.MyCoarseFilterFunction, self.FineSelectionFunction)
        
        # The functions technicalState and exitPositions are scheduled functions at specific time 
        # of the day.
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(15, 59), self.technicalState) 
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(10, 0), self.exitPositions) 
        

      # sort the data by daily dollar volume and take the top 'NumberOfSymbols'
    def MyCoarseFilterFunction(self, coarse):
        '''
        This is the first filter of the universe. In this filter we filtered out stocks with 
        price below to 10 and dollar Volume less than 10 millions dollars. Using this filter 
        we can get stocks with considerable volume
        '''
        
        
        if self.Time < self.nextRebalance:
            return Universe.Unchanged
        
        firstUniverse = [x for x in coarse  if x.HasFundamentalData and x.Price > 5 and x.DollarVolume > 1000000] 

        
        return [x.Symbol for x in firstUniverse]
        
    def FineSelectionFunction(self, fine):
        '''
        This function narrow the universe by ordering the first universe based on the Price
        To Sales ratio. Then it creates the longUniverse list with the 100 stocks with 
        the lowest Price to Sales ratio and the shortUniverse list with the 100 stocks with
        the greatest price to sales ratio
        '''
        
        if self.Time < self.nextRebalance:
            return Universe.Unchanged
            
        # Within this line we ordered the universe by Price To Sales ratio. The 
        # variable sortedByPsRatio will be a list where the first items are those
        # items with the lowest Price To Sales ratio(default order is ascending, so
        # lowest values are at first) and the bottom items will have the higest 
        # Price to Sales ratio
        sortedByPsRatio = sorted(fine, key= lambda x: x.ValuationRatios.PSRatio)
        
        # self.Debug(self.Time)
        #self.Debug('length %s' % len(sortedByPsRatio))
        
        # We get the first items of the sortedByPsRatio list by slicing the list 
        # using self.long value. Basically, if self.long is equal to 200, we take
        # the first 200 items of the sortedByPsRatio list, which are the 200 lowest
        # price to sales ratio. Then we call the Symbol method of the item to get
        # the symbol value
        self.longUniverse = [x.Symbol for x in sortedByPsRatio[:self.long]]
        
         # We get the last items for the sortedByPsRatio list by slicing the list
         # wit the last self.short value. So if self.short is 200 the self.shortUniverse 
         # will contain the latest 200 items from the sortedByPsRatio list.
         # Last we call the Symbol method of the item to get the symbol value
        self.shortUniverse = [x.Symbol for x in sortedByPsRatio[-self.short:]]
        
        
        return [x for x in (self.longUniverse + self.shortUniverse)] 
    
            
    def technicalState(self):
        '''
        This function scan for entry points for the stocks selected to go long or 
        go short. Once the technical indicators are ready for each symbol, the 
        function tested out the correspondent condition for go long or go short 
        for the stock selected in the universe, and send a market Order as well
        as a stopOrder and limit Order to stop position or take profit. 
        '''
        
        if self.Time < self.nextRebalance:
            return
        
        self.nextRebalance = self.Time + timedelta(self.rebalance_days)
        
    def exitPositions(self,):
        '''
        This function runs when the portfolio has one ore more position. The function
        will send order to exit positions (long or short) if the current date is 7 days
        greater than the order date. When the order is triggered the two open orders
        for stop loss and take profit are cancelled.
        '''
        
        if not self.Portfolio.Invested:
            return
        
        for symbol, order in self.orders.items():
       
                
            if order.Time.date() + timedelta(days= 7) < self.Time.date() and order.Quantity > 0 and self.Securities[symbol].Invested and symbol in self.stocksPurchased and ((self.Time.date().weekday !=5) or (self.Time.date().weekday !=6)):
                quantity = order.Quantity
                
                self.MarketOrder(symbol,-quantity,self.Time,'LongExitTime')
                self.stocksPurchased.remove(symbol)
                
                self.Debug('Exit after 7 days order for symbol {}'.format(symbol))
                
                # Cancel open orders
                cancelStop = self.stopOrders[symbol].Cancel()
                cancelProfit = self.profitOrders[symbol].Cancel()
                self.Debug("Cancel stop order for symbol {}".format(symbol))
                self.Debug("Cancel profit order for symbol {}".format(symbol))
                
              
                
            if order.Time.date() + timedelta(days=7) < self.Time.date() and order.Quantity < 0 and self.Securities[symbol].Invested and symbol in self.stocksPurchased and ((self.Time.date().weekday !=5) or (self.Time.date().weekday !=6)):
                quantity = order.Quantity
                self.MarketOrder(symbol,-quantity,self.Time,'ShortExitTime')
                self.stocksPurchased.remove(symbol)
                
                openOrders = self.Transactions.GetOpenOrders(symbol)
                
                # Cancel open orders
                #if self.exitByTime !=True:
                    
                cancelStop = self.stopOrders[symbol].Cancel()
                cancelProfit = self.profitOrders[symbol].Cancel()
                self.Debug("Cancel stop order for symbol {}".format(symbol))
                self.Debug("Cancel profit order for symbol {}".format(symbol))
               
                
    def OnSecuritiesChanged(self, changes):
        '''
        Show new securities added to the Universe as well as removed securities
        that dont belong any more to the Universe. At the moment where a security
        is added to the portfolio, we create the technical indicators objects.
        These objects will be uptated on each day in the technicalState function
        '''
        
       
        # Securities that are added to the universe based on the Fundamental Filters
        for security in changes.AddedSecurities:
            symbol = security.Symbol    
            self.Data[symbol] = SymbolData(self, symbol, self.BarPeriod, self.RollingWindowSize)
           
            
          
        
        # Securities that are removed from the universe because dont meet any more
        # fundamental filters.
        for removed in changes.RemovedSecurities:
            
            if removed.Symbol in self.longUniverse:
                self.longUniverse.remove(removed.Symbol)
        #       self.Debug('symbol removed is %s' % removed.Symbol)
               
                
            elif removed.Symbol in self.shortUniverse:
                self.shortUniverse.remove(removed.Symbol)
               
    
               
               
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm.
        In this function we calculate the amount of days for the next rebalance
        date. The variable self.rebalance_days is a variable that is defined
        in the Initailization part of the strategy, and means the number of days
        needed for the next rebalance.
        
        Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
            
        '''
        
        universe = self.longUniverse + self.shortUniverse
            
        for symbol in self.Data.keys():
            symbolData = self.Data[symbol]
            #self.Debug('On date {} symbolData is {}'.format(self.Time.date(),symbolData.IsReady()))
            if symbolData.IsReady() and symbolData.WasJustUpdated(self.Time):
                
                
                
                #self.Debug('barsCount %s' % symbolData.Bars.Count)
                if symbolData.Bars.Count == self.RollingWindowSize:
                
                    yesterdayBar = symbolData.Bars[1]
                    todayBar = symbolData.Bars[0]
                        
                    #
                    self.Debug('yesterdayclose %s' % yesterdayBar.Close)
                    #self.Debug('yesterday Open %s' % yesterdayBar.Open)
                    self.Debug('today close %s' % todayBar.Close)
                    #self.Debug('today Open %s' % todayBar.Open)
                    
                    #self.Debug('current Time %s' % self.Time)
                    #self.Debug('Time of Bars[0] %s' % symbolData.Bars[0].Time)
                    #self.Debug('Time of Bars[1] %s' % symbolData.Bars[1].Time)
                    
                    
                    #self.Debug('EndTime of Bars[0] %s' % symbolData.Bars[0].EndTime)
                    #self.Debug('EndTime of Bars[1] %s' % symbolData.Bars[1].EndTime)
                    
                    #if yesterdayBar and todayBar:
                    yesterdayClose = round(yesterdayBar.Close,2)
                    yesterdayOpen = round(yesterdayBar.Open,2)
                    todayClose = round(todayBar.Close,2)
                    todayOpen = round(todayBar.Open,2)
                
                    
                    upperBB = symbolData.BollingerBand.UpperBand.Current.Value
                    lowerBB = symbolData.BollingerBand.LowerBand.Current.Value
                        
                    sma_20 = round(symbolData.SMA20.Current.Value,2)
                    
                    if symbol in self.shortUniverse:
                        #self.Debug('On {} Symbol {} Yesterday Close {} BBUpper {} Yesterday Open {} Today Close {} Today Open {} SMA20 {}'.format(self.Time.date(),symbol,yesterdayClose,upperBB,yesterdayOpen,todayClose,todayOpen,sma_20))

                        if (yesterdayClose > upperBB) and (yesterdayClose > yesterdayOpen) and (todayClose < upperBB) and (todayClose < todayOpen) and (todayClose > yesterdayOpen) and (todayClose > (sma_20 * 1.1))  and symbol not in self.stocksPurchased:
                            self.Debug('Condition to sell for {}'.format(symbol))
                            self.Debug('On {} Symbol {} Yesterday Close {} BBUpper {} Yesterday Open {} Today Close {} Today Open {} SMA20 {}'.format(self.Time.date(),symbol,yesterdayClose,upperBB,yesterdayOpen,todayClose,todayOpen,sma_20))
                        
                        
                            # Define the tradeSize as 15% of the portfolio. Orders are sended by one stock
                            # so, most of the time the total position is at this level or below.
                            tradeSize = self.Portfolio.Cash * 0.15
                            quantity = int(round(tradeSize / todayClose))
                            stopPrice = round(todayClose * 1.20,2)
                        
                            self.Debug('Holdings Value %s' % self.Portfolio.TotalHoldingsValue)
                    
                            # Place the entry order to go short and inmediately place the stop Order at 5%
                            # of loss and the take profit at the sma_20 level.
                        
                            self.orders[symbol] = self.MarketOrder(symbol, -quantity,self.Time,'ShortTrade')
                        
                            #if self.exitByTime != True:
                            self.stopOrders[symbol] = self.StopMarketOrder(symbol, quantity, stopPrice,'StopShort') 
                            self.profitOrders[symbol] = self.LimitOrder(symbol,quantity,sma_20,'TakeProfitShort')
                                
                            # Every symbol that is purchased, is joined to the stocksPurchased list to avoid
                            # multiple orders for the same symbol on the same moment of time.
                            self.stocksPurchased.append(symbol)
                        
                    if symbol in self.longUniverse:
                        #self.Debug('On {} Symbol {} Yesterday Close {} BBLow {} Yesterday Open {} Today Close {} Today Open {} SMA20 {}'.format(self.Time.date(),symbol,yesterdayClose,lowerBB,yesterdayOpen,todayClose,todayOpen,sma_20))

                        if (yesterdayClose < lowerBB) and (yesterdayClose < yesterdayOpen) and (todayClose > lowerBB) and (todayClose > todayOpen) and (todayClose < yesterdayOpen) and (todayClose < (sma_20 * 0.9)) and symbol not in self.stocksPurchased:
                            self.Debug('Condition to buy for {}'.format(symbol))
                            self.Debug('On {} Symbol {} Yesterday Close {} BBLow {} Yesterday Open {} Today Close {} Today Open {} SMA20 {}'.format(self.Time.date(),symbol,yesterdayClose,lowerBB,yesterdayOpen,todayClose,todayOpen,sma_20))
                        
                            # Define the tradeSize as 15% of the portfolio. Orders are sended by one stock
                            # so, most of the time the total position is at this level or below.
                            tradeSize = self.Portfolio.Cash * 0.15
                            quantity = int(round(tradeSize / todayClose))
                            stopPrice = round(todayClose * 0.80,2)
                        
                            #self.Debug('Holdings Value %s' % self.Portfolio.TotalHoldingsValue)
                    
                            # Place the entry order to go short and inmediately place the stop Order at 5%
                            # of loss and the take profit at the sma_20 level.
                        
                            self.orders[symbol] = self.MarketOrder(symbol, quantity,self.Time,'LongTrade')
                            #if self.exitByTime !=True:
                            
                            self.stopOrders[symbol] = self.StopMarketOrder(symbol, -quantity, stopPrice,'StopLong')
                            self.profitOrders[symbol] = self.LimitOrder(symbol,-quantity,sma_20,'TakeProfitLong')
                    
                            #if self.exitByTime == True:
                            #self.stopOrders[symbol] = self.StopMarketOrder(symbol, -quantity, stopPrice,'StopLong')

                            self.stocksPurchased.append(symbol)
                        
                
 # Override the base class event handler for order events
    def OnOrderEvent(self, orderEvent):
        '''
        Each time an order is submitted or filled, this function is triggered. We use
        this function to cancel opposite function for each symbol. For example, when
        a profit order is filled, we should cancel the stop market order for that symbol.
        
        '''
        
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        #self.Debug("{0}: {1}: {2}: {3}: {4}".format(self.Time, OrderTypeCodes[order.Type], order.Quantity,order.Price,orderEvent))
        
        # OrderTypeCodes[order.Type] == 'StopMarket'  and order.Quantity < 0
        # If statetement to cancel profit limit order when the stop market order for loss is filled
        if order.Status == 3 and self.stopOrders[order.Symbol].OrderId == order.Id  and order.Quantity >0: 
        
            self.stocksPurchased.remove(order.Symbol)
           
            del self.orders[order.Symbol] 
            
            #if self.exitByTime != True:
            self.profitOrders[order.Symbol].Cancel()
            
            if order.Quantity > 0:
                self.Debug('Cancel short profit order for symbol {} because stop is triggered'.format(order.Symbol))
        
        if order.Status == 3 and self.stopOrders[order.Symbol].OrderId == order.Id  and order.Quantity < 0:
            self.stocksPurchased.remove(order.Symbol)
           
            del self.orders[order.Symbol] 
            
            self.profitOrders[order.Symbol].Cancel()
            
            self.Debug('Cancel long profit order for symbol {} because stop is triggered'.format(order.Symbol))
       
        if order.Status == 3 and  order.Id == self.profitOrders[order.Symbol].OrderId and order.Quantity >0: 
            
            self.stocksPurchased.remove(order.Symbol)
            
            self.stopOrders[order.Symbol].Cancel()
            del self.orders[order.Symbol] 
            
            #if order.Quantity > 0:
            self.Debug('At {} Cancel short loss order for symbol {} because profit is triggered'.format(self.Time.date(),order.Symbol))
       
        if order.Status == 3 and  order.Id == self.profitOrders[order.Symbol].OrderId and order.Quantity <0:
            
            self.stocksPurchased.remove(order.Symbol)
            self.stopOrders[order.Symbol].Cancel()
            del self.orders[order.Symbol] 
            
            self.Debug('At {} Cancel long loss order for symbol {} because profit is triggered'.format(self.Time.date(),order.Symbol))
                
     
class SymbolData(object):
    
    def __init__(self, algorithm, symbol, barPeriod, windowSize):
        self.symbol = symbol
        self.algorithm = algorithm
        self.BarPeriod = barPeriod
    
        self.SMA20 = SimpleMovingAverage(symbol,20) 
           
        self.Bars = RollingWindow[TradeBar](windowSize)
            
        self.BollingerBand = algorithm.BB(symbol ,20,2)
            
        consolidator = TradeBarConsolidator(barPeriod)
            
        consolidator.DataConsolidated += self.OnDataConsolidated
            
        algorithm.SubscriptionManager.AddConsolidator(symbol,consolidator)
    
    
    def OnDataConsolidated(self,sender,bar):
        self.SMA20.Update(bar.EndTime, bar.Close)
        self.BollingerBand.Update(bar.EndTime, bar.Close)
        self.Bars.Add(bar)
        
    # Returns true if all the data in this instance is ready (indicators, rolling windows, ect...)
    
    def IsReady(self):
    
        return self.Bars.IsReady and self.SMA20.IsReady and self.BollingerBand.IsReady
    
    # Returns true if the most recent trade bar time matches the current time minus the bar's period, this
    # indicates that update was just called on this instance
    
    def WasJustUpdated(self, current):
        return self.Bars.Count > 0 and self.Bars[0].EndTime == current - self.BarPeriod