import numpy as np
class ATRMomentumUniverse(QCAlgorithm):
    
    def Initialize(self):
        self.SetStartDate(2020, 4, 7)
        self.SetEndDate(2020, 5, 6)
        self.SetCash(10000)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction) 
        
        self.assets = { }                                                       # List of traded assets
        
    def CoarseSelectionFunction(self, universe):  
        
        selected = []
        
        universe = sorted(universe, key=lambda c: c.DollarVolume, reverse=True) # Sorts the whole 8k universe after volume 
        universe = [c for c in universe if c.Price > 10][:1000]                  # Take the first 100 assets if they trade above 10$
        
        for coarse in universe:                                                 # For every asset dict in the universe
            symbol = coarse.Symbol                                              # Take the symbol name out of the dict
            
            if symbol not in self.assets:                                       # If it is not in the traded asset list yet
                history = self.History(symbol, 20, Resolution.Daily)            # Pull the history of the last 20 days
                
                self.assets[symbol] = SelectionData(symbol, history)            # Insert history into the data class and run it for the symbol
        
        selected = list(self.assets.keys())
        
        return selected[:1]
        
        
        
    #def OnSecuritiesChanged(self, changes):
        
        #for security in changes.AddedSecurities:
            #self.SetHoldings(security.Symbol, 0.1)
            #self.StopMarketOrder(security.Symbol, -self.Portfolio[security.Symbol].Quantity, security.Price*(0.995))
            
            
            
    def OnData(self, data):
        
        for symbol, selectionData in self.assets.items():
            
            if data.Bars.ContainsKey(symbol):
                selectionData.update(data.Bars[symbol])
                
                print('Ready')
                
                if self.Portfolio[symbol].Invested and selectionData.atr.Current.Value > 0.6:
                    self.Liquidate(symbol)
                    
                if self.assets[symbol].is_ready():                              # If the symbol is tradable
                    atr_quantile = np.quantile([x.Value for x in list(self.assets[symbol].atr_container)[0:15]], 0.25)  # Calculate the quantile value of the last 15 days of the ATR
                    
                    print('Ready2')
                    
                    self.Log((atr_quantile, self.assets[symbol].macd_container_slow[0], self.assets[symbol].macd_container_fast[0]))
                    
                    # Trading Logic Long
                    if self.assets[symbol].atr_container[0].Value <= atr_quantile and self.assets[symbol].macd_container_slow[0].Value > self.assets[symbol].macd_container_fast[0].Value and self.assets[symbol].macd_container_slow[0].Value > 0:
                        #selected_longs.append(symbol)
                        self.StopMarketOrder(security.Symbol, 1, data.Bars[symbol].High[1])
                        
                    # Trading Logic Short
                    if self.assets[symbol].atr_container[0].Value <= atr_quantile and self.assets[symbol].macd_container_slow[0].Value < self.assets[symbol].macd_container_fast[0].Value and self.assets[symbol].macd_container_slow[0].Value < 0:
                        #selected_shorts.append(symbol)
                        self.StopMarketOrder(security.Symbol, -1, data.Bars[symbol].Low[1])
                        
                        
class SelectionData():
    def __init__(self, symbol, history):
        self.symbol = symbol
        
        self.macd = MovingAverageConvergenceDivergence(12, 26, 9)
        self.macd.Updated += self.macd_container_update
        self.macd_container_slow = RollingWindow[IndicatorDataPoint](3)
        self.macd_container_fast = RollingWindow[IndicatorDataPoint](3)
        
        self.atr = AverageTrueRange(1)
        self.atr.Updated += self.atr_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)
            #self.macd.Update(bar.close)
              
    def is_ready(self):
        return self.atr_container.IsReady
    
    def update(self, bar):
        self.atr.Update(bar)
        self.macd.Update(bar.EndTime, bar.Close)
    
    def macd_container_update(self, sender, updated):
        self.macd_container_slow.Add(self.macd.Slow.Current)
        self.macd_container_fast.Add(self.macd.Fast.Current)
        
    def atr_container_update(self, sender, updated):
        self.atr_container.Add(updated)