import pandas as pd

class UglyYellowGreenGalago(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2006, 5, 2)  # Set Start Date
        self.SetEndDate(2006, 6, 1)
        self.SetCash(100000)  # Set Strategy Cash
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetUniverseSelection(QC500UniverseSelectionModel())
        self.SetSecurityInitializer(self.CustomSecurityInitializer)
        self.symbol_data_by_symbol = {}
    
        
    def CustomSecurityInitializer(self, security):
        security.SetFeeModel(ConstantFeeModel(0))
        security.SetSlippageModel(ConstantSlippageModel(0))
        security.SetFillModel(NoBidAskSpreadModel())       
        
    def OnData(self, data):
        # Gather historical prices
        closes = pd.DataFrame()
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if not (data.ContainsKey(symbol) and data[symbol] is not None and symbol_data.IsReady):
                continue
            closes[symbol] = symbol_data.closes
        if closes.empty:
            return
        
        weight = 1/len(closes.columns)
        for symbol in closes.columns:
            self.SetHoldings(symbol, weight)
        self.Quit()
        
        
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            self.symbol_data_by_symbol[security.Symbol] = SymbolData(self, security.Symbol)
        
        for security in changes.RemovedSecurities:
            self.Liquidate(security.Symbol)
            symbol_data = self.symbol_data_by_symbol.pop(security.Symbol, None)
            if symbol_data:
                symbol_data.dispose()


class SymbolData:
    def __init__(self, algorithm, symbol, lookback=6):
        self.algorithm = algorithm
        self.symbol = symbol
        self.lookback = lookback
        
        # Setup consolidator to track trailing prices
        self.consolidator = TradeBarConsolidator(1)
        self.consolidator.DataConsolidated += self.consolidation_handler
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)
        
        self.closes = pd.Series()
        history = algorithm.History(symbol, lookback + 1, Resolution.Daily)
        if history.empty or 'close' not in history.columns:
            return
        self.closes = history.loc[symbol].close
        
    def consolidation_handler(self, sender, bar):
        if bar.EndTime in self.closes.index:
            return
        self.closes = self.closes.append(pd.Series([bar.Close], index=[bar.EndTime])).iloc[-(self.lookback + 1):]
        
    @property
    def IsReady(self):
        return len(self.closes) == self.lookback + 1
            
        
    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)
        
        
class NoBidAskSpreadModel(FillModel):

    def MarketFill(self, asset, order):
        fill = super().MarketFill(asset, order)
        fill.FillPrice = asset.Price # Set fill price to last traded price
        return fill