import numpy as np
class ATRMomentumUniverse(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2019, 1, 7)
        self.SetEndDate(2020, 4, 1)
        self.SetCash(100000)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction) 
        self.atrscan = { }
    
    def CoarseSelectionFunction(self, universe):  
        selected = []
        universe = sorted(universe, key=lambda c: c.DollarVolume, reverse=True)  
        universe = [c for c in universe if c.Price > 10][:100]
        
        for coarse in universe:  
            symbol = coarse.Symbol
            
            if symbol not in self.atrscan:
                history = self.History(symbol, 20, Resolution.Daily)
                self.atrscan[symbol] = SelectionData(symbol, history) 
                
            if self.atrscan[symbol].is_ready():
                atr_quantile = np.quantile([x.Value for x in list(self.atrscan[symbol].atr_container)[0:15]], 0.25)
                if  self.atrscan[symbol].atr_container[0].Value < atr_quantile:
                    selected.append(symbol)
        
        return selected[:10]
        
    def OnSecuritiesChanged(self, changes):
        
        for security in changes.AddedSecurities:
            self.SetHoldings(security.Symbol, 0.1)
            
    def OnData(self, data):
        
        for symbol, selectionData in self.atrscan.items():
            
            if data.Bars.ContainsKey(symbol):
                selectionData.update(data.Bars[symbol])
            
                if self.Portfolio[symbol].Invested and selectionData.atr.Current.Value > 100:
                        self.Liquidate(symbol)
        
class SelectionData():
    def __init__(self, symbol, history):
        self.symbol = symbol
        self.atr = AverageTrueRange(1)
        self.atr.Updated += self.container_update
        self.atr_container = RollingWindow[IndicatorDataPoint](20)
        
        
        for bar in history.itertuples():
            
            tbar = TradeBar(bar.Index[1], self.symbol, bar.open, bar.high, bar.low, bar.close, bar.volume)
            self.atr.Update(tbar)
              
    def is_ready(self):
        return self.atr_container.IsReady
    
    def update(self, bar):
        self.atr.Update(bar)
    
    def container_update(self, sender, updated):
        self.atr_container.Add(updated)