# https://quantpedia.com/Screener/Details/14
class MomentumEffectAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2014, 1, 1)  # Set Start Date
        self.SetEndDate(2018, 6, 1)    # Set Start Date       
        self.SetCash(100000)           # Set Strategy Cash

        self.UniverseSettings.Resolution = Resolution.Daily

        self.mom = {}           # Dict of Momentum indicator keyed by Symbol
        self.lookback = 21*12   # Momentum indicator lookback period
        self.num_coarse = 100   # Number of symbols selected at Coarse Selection
        self.num_fine = 50      # Number of symbols selected at Fine Selection
        self.num_long = 5       # Number of symbols with open positions

        self.month = -1
        self.symbols = []

        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)


    def CoarseSelectionFunction(self, coarse):
        '''Drop securities which have no fundamental data or have too low prices.
        Select those with highest by dollar volume'''
        
        if self.month == self.Time.month:
            return self.symbols

        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 5],
            key=lambda x: x.DollarVolume, reverse=True)
            
        return [x.Symbol for x in selected[:self.num_coarse]]


    def FineSelectionFunction(self, fine):
        '''Select security with highest market cap'''
        
        if self.month == self.Time.month:
            return self.symbols

        for f in fine:
            f.MarketCap = (f.ValuationRatios.PERatio *
                           f.EarningReports.BasicEPS.TwelveMonths *
                           f.EarningReports.BasicAverageShares.ThreeMonths)
                        
        selected = sorted([f for f in fine if f.MarketCap > 0],
            key=lambda x: f.MarketCap, reverse=True)

        self.symbols = [x.Symbol for x in selected[:self.num_fine]]
        
        return self.symbols


    def OnData(self, data):

        # Update the indicator
        for symbol, mom in self.mom.items():
            mom.Update(self.Time, self.Securities[symbol].Close)

        if self.month == self.Time.month:
            return

        self.month = self.Time.month

        # Selects the securities with highest momentum
        sorted_mom = sorted(self.mom, key=lambda x: self.mom[x].Current.Value, reverse=True)
        selected = sorted_mom[:self.num_long]

        # Liquidate securities that are not in the list
        for symbol in self.symbols:
            if symbol not in selected:
                self.Liquidate(symbol)

        # Buy selected securities
        for symbol in selected:
            self.SetHoldings(symbol, 1/self.num_long)


    def OnSecuritiesChanged(self, changes):

        # Warm up the indicator with history price for newly added securities
        addedSymbols = [x.Symbol for x in changes.AddedSecurities]
        history = self.History(addedSymbols, 1 + self.lookback, Resolution.Daily)
        history = history.close.unstack(level=0)

        for symbol in addedSymbols:
            if symbol not in self.mom:
                self.mom[symbol] = Momentum(self.lookback)
                ticker = str(symbol)
                if ticker in history:
                    for time, value in history[ticker].items():
                        item = IndicatorDataPoint(symbol, time, value)
                        self.mom[symbol].Update(item)

        # Clean up data for removed securities and Liquidate
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if self.mom.pop(symbol, None) is not None:
                self.Liquidate(symbol)