from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class TheSqueezeUniverseSelection(FundamentalUniverseSelectionModel):
    
    def __init__(self, period = 20):
        super().__init__(False, None)
        
        self.period = period
        self.indicators = {}
    
    def SelectCoarse(self, algorithm, coarse):  
        selected = []
        universe = sorted(coarse, 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.indicators:
                # 1. Call history to get an array of 200 days of history data
                history = algorithm.History(symbol, 200, Resolution.Daily)
                
                #2. Adjust SelectionData to pass in the history result
                self.indicators[symbol] = SelectionData(history, self.period) 

            indicator = self.indicators[symbol]
            
            indicator.update(algorithm.Time, coarse.AdjustedPrice)
            
            if indicator.is_ready() and \
                indicator.BollingerUpper < indicator.KeltnerUpper and \
                indicator.BollingerLower > indicator.KeltnerLower:
                
                selected.append(symbol)
        
        return selected[:10]
        
class SelectionData():
    #3. Update the constructor to accept a history array
    def __init__(self, history, period):
        self.bollinger = BollingerBands(period, 2, MovingAverageType.Simple)
        self.keltner = KeltnerChannels(period, 1.5, MovingAverageType.Simple)
        #4. Loop over the history data and update the indicatorsc
        for bar in history.itertuples():
            tradeBar = TradeBar(bar.Index[1], bar.Index[0], bar.open, bar.high, bar.low, bar.close, bar.volume, timedelta(1))
            self.bollinger.Update(bar.Index[1], bar.close)
            self.keltner.Update(tradeBar)
    
    @property
    def BollingerUpper(self):
        return float(self.bollinger.UpperBand.Current.Value)
        
    @property
    def BollingerLower(self):
        return float(self.bollinger.LowerBand.Current.Value)
        
    @property
    def KeltnerUpper(self):
        return float(self.keltner.UpperBand.Current.Value)
        
    @property
    def KeltnerLower(self):
        return float(self.keltner.LowerBand.Current.Value)
    
    
    def is_ready(self):
        return self.bollinger.IsReady and self.keltner.IsReady
    
    def update(self, time, value):
        return self.bollinger.Update(time, value)
            

from Universe3 import TheSqueezeUniverseSelection

class UglySkyBlueRhinoceros(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 2, 15)  # Set Start Date
        self.SetEndDate(2021, 4, 1)
        self.SetCash(100000)  # Set Strategy Cash

        self.SetUniverseSelection(TheSqueezeUniverseSelection())