class SymbolData(object):
    def __init__(self, symbol):
        self.symbol = symbol
        self.sma_fast = SimpleMovingAverage(10)
        self.sma_slow = SimpleMovingAverage(30)
        self.is_ready = False
        self.has_fundamentals = False
        self.pct_diff = None
        
    def update_warmup(self, time, price):
        self.is_ready = all([
            self.sma_fast.Update(time, price),
            self.sma_slow.Update(time, price)
        ])

    def update(self, stock):
        self.has_fundamentals = stock.HasFundamentalData
        
        self.update_warmup(stock.EndTime, stock.Price)

        if self.is_ready:
            sma_fast = float(self.sma_fast.Current.Value)
            sma_slow = float(self.sma_slow.Current.Value)
            self.pct_diff = (sma_fast - sma_slow) / sma_slow
    

class TimeoutProblemAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2016, 12, 29)
        self.SetEndDate(2016, 12, 30)
        self.SetCash(100000)
        
        self.WarmUp = 30
        self.MaxCandidates = 100

        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.warmed_up = False
        self.symbol_data = {}

        self.AddEquity("SPY", Resolution.Daily)
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", -10), self.before_trading_start)
        
    def CoarseSelectionFunction(self, coarse):
        if not self.warmed_up:
            for stock in coarse:
                self.AddEquity(stock.Symbol.Value)
            history = self.History([stock.Symbol for stock in coarse], self.WarmUp, Resolution.Daily)
            history_tickers = history.index.get_level_values(0)
            for stock in coarse:
                if stock.Symbol.Value in history_tickers:
                    self.symbol_data[stock.Symbol] = SymbolData(stock.Symbol)
                    data = history.loc[stock.Symbol.Value]
                    for time in data.index:
                        close_price = data.loc[time]['close']
                        self.symbol_data[stock.Symbol].update_warmup(time, close_price)
                self.RemoveSecurity(stock.Symbol)
                self.Securities.Remove(stock.Symbol)
            self.Debug("Warmed up {} stocks".format(len(self.symbol_data)))
            self.warmed_up = True

        coarse_symbols = [stock.Symbol for stock in coarse]
        for symbol in list(self.symbol_data.keys()):
            if symbol not in coarse_symbols:
                del self.symbol_data[symbol]
        
        for stock in coarse:
            if stock.Symbol not in self.symbol_data:
                self.symbol_data[stock.Symbol] = SymbolData(stock.Symbol)
            self.symbol_data[stock.Symbol].update(stock)

        return [data.symbol for data in self.symbol_data.values() if \
                       data.is_ready and \
                       data.has_fundamentals]

    def FineSelectionFunction(self, fine):
        self.Debug('Fine collection is {} long'.format(len(fine)))
        security_filter = [self.symbol_data[x.Symbol] for x in fine if \
            x.SecurityReference.IsPrimaryShare and \
            x.SecurityReference.SecurityType == 'ST00000001' and \
            x.SecurityReference.ExchangeId != 'OTC' and \
            not x.CompanyReference.IsLimitedPartnership and \
            not x.SecurityReference.IsDepositaryReceipt]

        stock_data = sorted([data for data in security_filter], key=lambda x: x.pct_diff)
        todays_universe = [data.symbol for data in stock_data[:self.MaxCandidates]]
        self.Debug('Today\'s universe: {}'.format(', '.join([sym.Value for sym in todays_universe])))
        return todays_universe

    def before_trading_start(self):
        if not self.warmed_up:
            return
        
        self.Debug('before_trading_start: {}'.format(self.Securities.Count))