from QuantConnect.Indicators import SimpleMovingAverage

class ParticleVerticalCompensator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019,1,2)
        self.SetEndDate(2019,1,5)
        self.SetCash(100000)
        
        self.UniverseSettings.Resolution = Resolution.Daily

        self.AddUniverse(self.MyCoarseFilterFunction)
        
        self.indicators_by_symbol = {}
        self.sma_period = 10
        
    def MyCoarseFilterFunction(self, coarse):
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        filtered = [ x.Symbol for x in sortedByDollarVolume
                      if x.Price > 5 and x.DollarVolume > 400000000 ]
        return filtered[:3]
    
    def OnSecuritiesChanged(self, changes):
        for added in changes.AddedSecurities:
            self.indicators_by_symbol[added.Symbol] = SimpleMovingAverage(added.Symbol, self.sma_period)
            
            # Warm up indicator
            history = self.History(added.Symbol, self.sma_period, Resolution.Daily).loc[added.Symbol]
            for idx, row in history.iterrows():
                self.indicators_by_symbol[added.Symbol].Update(idx, row['close'])
            
            # Register indicator for new data
            self.RegisterIndicator(added.Symbol, self.indicators_by_symbol[added.Symbol], Resolution.Daily)
            
            self.Log(f"{added.Symbol}  {self.indicators_by_symbol[added.Symbol].Current.Value}")
            
        for removed in changes.RemovedSecurities:
            self.indicators_by_symbol.pop(removed.Symbol, None)
        
    def OnData(self, data):
        for symbol in self.indicators_by_symbol:
            indicator = self.indicators_by_symbol[symbol]
            self.Log(f"{symbol}: {indicator.Current.Value}")