import heapq

class PensiveBlackSalamander(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 28)
        self.SetCash(100000) 
        self.AddUniverse(self.MyCoarseFilterFunction)
        self.selected = []
        self.UniverseSettings.Resolution = Resolution.Daily
        self.dataBySymbol = {}
        self.AddEquity("SPY")
        self.Schedule.On(self.DateRules.MonthStart("SPY"), \
                 self.TimeRules.AfterMarketOpen("SPY"), \
                 self.liquidate)
    def liquidate(self):
        self.Liquidate()
        

    def MyCoarseFilterFunction(self, coarse):
         
         
         sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
         filtered = [ x for x in sortedByDollarVolume 
                      if x.Price > 10 and x.DollarVolume > 10000000 ]
        # probably want some different logic for filtered or none at all
#         filtered = filtered[:20]
         filtered = filtered
         self.selected = []
         self.heap = []
         for x in filtered:
             if x.Symbol not in self.dataBySymbol:
                 self.dataBySymbol[x.Symbol] = SymbolData(self,x.Symbol)
             self.dataBySymbol[x.Symbol].roc.Update(self.Time,x.Price)
             self.dataBySymbol[x.Symbol].rocd.Update(self.Time,x.Price)
             if self.dataBySymbol[x.Symbol].IsReady:
                 #https://docs.python.org/3/library/heapq.html
                heapq.heappush(self.heap, (self.dataBySymbol[x.Symbol].rocd.Current.Value, x.Symbol))

        # choose the 5 largest in the heap
         for h in heapq.nlargest(5,self.heap):
             self.selected.append(h[1])
         
         self.Debug(str(self.selected))
         return self.selected
         
    def OnData(self, data):
        for symbol in self.selected:
            if not self.Portfolio[symbol].Invested:
                self.SetHoldings(symbol,1/len(self.selected))
    
    def OnSecuritiesChanged(self,changes):
        for x in changes.RemovedSecurities:
            self.dataBySymbol.pop(x.Symbol,None)
            self.Liquidate(x.Symbol)
class SymbolData:
    def __init__(self,algo,symbol):
        self.roc = RateOfChange(252)
        self.rocd = IndicatorExtensions.Of(Delay(21), self.roc)
        self.symbol = symbol
        history = algo.History(symbol,252,Resolution.Daily)
        if not  history.empty:
            for time, row in history.loc[symbol].iterrows():
                self.roc.Update(time,row['close'])
                self.rocd.Update(time,row['close'])
    @property
    def IsReady(self):
        return self.roc.IsReady
        return self.rocd.IsReady