class RetrospectiveBlueWhale(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)
        self.AddUniverse(self.SelectCoarse, self.SelectFine)
        self.UniverseSettings.Resolution = Resolution.Minute
        self.leverageNumber = 3
        self.UniverseSettings.Leverage = self.leverageNumber
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.03))
        
        self.mompLookback = self.GetParameter("momp-lookback")
        self.mompEntry = self.GetParameter("momp-entry")
        self.mompLookback = int(self.mompLookback)
        self.mompEntry = int(self.mompEntry)
        
        self.symbolDataBySymbol = {}
        
        self.coarseNumber = 100
        self.fineNumber = 10
        self.lastMonth = -1
        
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen(self.spy, 30),        
                 self.SetTrades)

    def OnData(self, data):
        if self.IsWarmingUp:
            return
                
    def SetTrades(self):
        for symbol, symbolData in self.symbolDataBySymbol.items():
            rsi = symbolData.Rsi.Current.Value
            momp = symbolData.Momp.Current.Value
            ema = symbolData.Ema.Current.Value
            
            #self.Debug(str(symbol) + " " + str(momp))
        
            if momp < (-self.mompEntry) and not self.Portfolio[symbol].IsShort and ema < 9:
                if (-self.mompEntry) >= momp >= (-self.mompEntry)-10:
                    self.SetHoldings([PortfolioTarget(symbol, -(1/self.fineNumber)*(self.leverageNumber-1))])
                elif (-self.mompEntry)-10 > momp >= (-self.mompEntry)-20:
                    self.SetHoldings([PortfolioTarget(symbol, -(1/self.fineNumber)*self.leverageNumber)])
                elif momp < (-self.mompEntry)-20:
                    ''' ARMAGEDON '''
                    self.Liquidate()
                    self.SetHoldings([PortfolioTarget(symbol, -1.0*self.leverageNumber)])
                    break
            elif momp > self.mompEntry and not self.Portfolio[symbol].IsLong and ema > 9:
                if self.mompEntry <= momp <= self.mompEntry+10:
                    self.SetHoldings([PortfolioTarget(symbol, (1/self.fineNumber)*(self.leverageNumber-1))])
                elif self.mompEntry+10 < momp <= self.mompEntry+20:
                    self.SetHoldings([PortfolioTarget(symbol, (1/self.fineNumber)*self.leverageNumber)])
                elif momp > self.mompEntry+20:
                    ''' ARMAGEDON '''
                    self.Liquidate()
                    self.SetHoldings([PortfolioTarget(symbol, 1.0*self.leverageNumber)])
                    break
        
    def OnSecuritiesChanged(self, changes):
        added = [x for x in changes.AddedSecurities]
        removed = [x for x in changes.RemovedSecurities]
        
        for x in added:
            symbol = x.Symbol
            if symbol == self.spy:
                continue
            
            rsi = self.RSI(symbol, 65, MovingAverageType.Simple, Resolution.Minute)
            sma = self.SMA(symbol, 65, Resolution.Minute, Field.Volume)
            obv = self.OBV(symbol, Resolution.Minute)
            
            momp = self.MOMP(symbol, self.mompLookback, Resolution.Minute, Field.Volume)
            ema = self.EMA(symbol, 14, Resolution.Daily, Field.Close)
            #priceRoc = self.ROC(symbol, 14, Resolution.Daily, Field.Close)
            history = self.History(symbol, 300, Resolution.Minute)
            historyDaily = self.History(symbol, 14, Resolution.Daily)
            
            for tuple in history.loc[symbol].itertuples():
                rsi.Update(tuple.Index, tuple.close)
                sma.Update(tuple.Index, tuple.volume)
                tradebar = TradeBar(tuple.Index, symbol, tuple.open, tuple.high, tuple.low, tuple.close, tuple.volume)
                obv.Update(tradebar)
                momp.Update(tuple.Index, tuple.volume)
                
            for tuple in historyDaily.loc[symbol].itertuples():
                ema.Update(tuple.Index, tuple.close)
                
            symbolData = SymbolData(symbol, rsi, sma, momp, ema)
            self.symbolDataBySymbol[symbol] = symbolData
        
    def SelectCoarse(self, coarse):
        if self.lastMonth == self.Time.month:
            return Universe.Unchanged
        self.lastMonth = self.Time.month
        
        filteredCoarse = sorted([x for x in coarse if x.Price > 50 and x.DollarVolume > 1000], key = lambda x: x.DollarVolume, reverse = True)
        return [x.Symbol for x in filteredCoarse][:self.coarseNumber]
    
    def SelectFine(self, fine):
        filteredFine = sorted([x for x in fine], key = lambda x: x.MarketCap, reverse = True)
        return [x.Symbol for x in filteredFine][:self.fineNumber]
        
class SymbolData:
    def __init__(self, symbol, rsi, sma, momp, ema):
        self.Symbol = symbol
        self.Rsi = rsi
        self.Sma = sma
        self.Momp = momp
        self.Ema = ema