class SmoothBlueGuanaco(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2020, 1, 1)  # Set Start Date
        self.SetEndDate(2021,1,1)
        self.SetCash(100000)  # Set Strategy Cash
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.data = {}
        self.closingPrices = {}
        self.symbols =  ["SPY", "IGV", "XLE", "VNQ", "XLF", "XME", "XHB", "TBT", "KWEB", "OIH","XBI"]
        for symbol in self.symbols:
            self.AddEquity(symbol, Resolution.Daily)
            self.data[symbol] = self.MOM(symbol, 30, Resolution.Daily)
            self.closingPrices[symbol] = self.Identity(symbol, Resolution.Daily, Field.Close)
            self.Log("{} symbol:{} closingPrice:{}".format(self.Time, symbol, self.closingPrices[symbol]))
      
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Friday), self.TimeRules.AfterMarketOpen("SPY",180), self.Rebalance) # Rebalances at 12:30 EST
        self.Schedule.On(self.DateRules.On(2022, 5, 31), self.TimeRules.AfterMarketOpen("SPY", 150), self.Rebalance)

        self.Settings.FreePortfolioValuePercentage=0
        self.Leverage=1.4
        self.assetsToHoldLong=2 # number of assets to hold long in portfolio 
        self.assetsToHoldShort=2 # number of assets to hold short in portfolio 
        

    def OnData(self, data: Slice):
        if self.IsWarmingUp: return
            
        #self.Debug("{} Entered OnData".format(self.Time))
        for symbol in self.symbols:
            if data.Bars.ContainsKey(symbol) and data.Bars[symbol] is not None:
                self.closingPrices[symbol] = data.Bars[symbol].Close
                
        for symbol in self.symbols:
            self.Plot('closingPrices',  symbol, self.closingPrices[symbol])

    def Rebalance(self):
        
        topLong = pd.Series(self.data).sort_values(ascending = False)[:self.assetsToHoldLong]
        topShort = pd.Series(self.data).sort_values(ascending = True)[:self.assetsToHoldShort]
        #self.Debug("{} Entering Rebalance with topLong:{}".format(self.Time, topLong ))
        #self.Debug("{} Entering Rebalance with topShort:{}".format(self.Time, topShort ))
        
        for kvp in self.Portfolio:
            security_key = kvp.Key
            security_hold = kvp.Value
            # liquidate the security which is no longer in the top3 momentum list
            if security_hold.Invested and self.Portfolio[security_key].Quantity > 1 :
                # For longs, liquidate position if momentum drops below 1 
                #self.Debug("Entered long")
                if (security_hold.Symbol.Value not in topLong.index or (security_hold.Invested and self.data[security_hold.Symbol.Value].Current.Value < 2)):
                    #self.Debug("{} Elim Long pos {} @{} Qty:{} MOM:{}".format(self.Time, security_hold, self.closingPrices.get(security_hold.Symbol, "None"),self.Portfolio[security_key].Quantity, self.data[security_hold.Symbol.Value].Current.Value))
                    self.Liquidate(security_hold.Symbol)
            elif security_hold.Invested and self.Portfolio[security_key].Quantity < 1 :
                #self.Debug("Entered Short")
                if security_hold.Symbol.Value not in topShort.index or (security_hold.Invested and (self.data[security_hold.Symbol.Value].Current.Value > -2)):
                # For longs, liquidate position if momentum drops below 1 
                    #self.Debug("{} Elim Short pos {} @{} Qty:{} MOM:{}".format(self.Time, security_hold, self.closingPrices.get(security_hold.Symbol, "None"), self.Portfolio[security_key].Quantity, self.data[security_hold.Symbol.Value].Current.Value))
                    self.Liquidate(security_hold.Symbol)
            
        # Buy top long stocks 
        added_symbolsLong = []        
        for symbol in topLong.index:
            self.Debug("{} TOP LONG Symbol:{} Mom:{} Close:{}".format(self.Time, symbol, self.data[symbol].Current.Value, self.closingPrices[symbol]))
            if not self.Portfolio[symbol].Invested and self.data[symbol].Current.Value > 1:
                added_symbolsLong.append(symbol)
                
         # Sell top short stocks 
        added_symbolsShort = []        
        for symbol in topShort.index:
            self.Debug("{} TOP SHORT Symbol:{} Mom:{} Close:{}".format(self.Time, symbol, self.data[symbol].Current.Value, self.closingPrices[symbol]))
            if not self.Portfolio[symbol].Invested and self.data[symbol].Current.Value < -1:
                added_symbolsShort.append(symbol)
                
        added_symbols = added_symbolsLong+added_symbolsShort
        count = 0
        for added in added_symbolsLong:

            lastPrice = self.closingPrices[added]
            
            if not self.closingPrices[added] ==0: 
                sharesToBuy= (self.Portfolio.MarginRemaining*self.Leverage) / (len(added_symbols)-count) // lastPrice
            self.Debug("{} Shares to buy Symbol:{} Close:{} RemMargin:{:.2f} shares:{}"
                    .format(self.Time, added, lastPrice , self.Portfolio.MarginRemaining, sharesToBuy))
            
            self.MarketOrder(added, sharesToBuy)
            count =+1
                
        for added in added_symbolsShort:
            lastPrice = self.closingPrices[added]
            
            sharesToBuy= - (self.Portfolio.MarginRemaining*self.Leverage) / (len(added_symbols)-count) // lastPrice
            self.Debug("{} Shares to sell Symbol:{} Close:{} RemMargin:{:.2f} shares:{}"
                    .format(self.Time, added, lastPrice , self.Portfolio.MarginRemaining, sharesToBuy))
            
            self.MarketOrder(added, sharesToBuy)