# region imports
from AlgorithmImports import *
# endregion

# https://quantpedia.com/Screener/Details/54
class MomentumandStateofMarketFiltersAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2011, 1, 1)
        self.SetEndDate(2018, 8, 1)
        self.SetCash(100000)
        # add Wilshire 5000 Total Market Index data
        self.wilshire_symbol = self.AddData(Fred, Fred.Wilshire.Price5000, Resolution.Daily).Symbol
        self.W5000Return = self.ROC(self.wilshire_symbol, 252)
        # initialize the RateOfChange indicator of Wilshire 5000 total market index
        history = self.History(self.wilshire_symbol, 500, Resolution.Daily)
        for tuple in history.loc[self.wilshire_symbol].itertuples():
            self.W5000Return.Update(tuple.Index, tuple.value)
        self.Debug("W5000 Rate of Change indicator isReady: "+ str(self.W5000Return.IsReady))    
        self.AddEquity("SPY", Resolution.Daily)
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.At(0, 0), self.rebalance)
        # mark it's the start of each month
        self.month_start = False
        # mark the coarse universe selection has finished
        self.selection = False
        self.momp = {}
        self.lookback = 20*6
        self.long = None
        self.short = None
        self.tlt = self.AddEquity("TLT", Resolution.Daily).Symbol        
    
    def CoarseSelectionFunction(self, coarse):
        coarse = [x for x in coarse if (x.HasFundamentalData and x.AdjustedPrice > 1)]
        for i in coarse:
            if i.Symbol not in self.momp:
                self.momp[i.Symbol] = SymbolData(i.Symbol, self.lookback, self)
            else:
                self.momp[i.Symbol].MOMP.Update(self.Time, i.AdjustedPrice)
        
        if self.month_start:
            self.selection = True
            self.MOMPReady = {symbol: SymbolData for symbol, SymbolData in self.momp.items() if SymbolData.MOMP.IsReady}
            if self.MOMPReady:
                # sort stocks by 6-month momentum
                sortByMOMP = sorted(self.MOMPReady, key = lambda x: self.MOMPReady[x].MOMP.Current.Value, reverse = True)
                self.long = sortByMOMP[:20]
                self.short = sortByMOMP[-20:]
                return self.long+self.short
            else:
                return []
        else: 
            return []   

    def rebalance(self):
        # rebalance every month
        self.month_start = True
            
    def OnData(self, data):
        if self.month_start and self.selection: 
            self.month_start = False
            self.selection = False
            if self.long is None or self.short is None: return
            # if the previous 12 months return on the broad equity market was positive
            if self.W5000Return.Current.Value > 0: 
                stocks_invested = [x.Key for x in self.Portfolio if x.Value.Invested]
                for i in stocks_invested:
                    if i not in self.long+self.short:
                        self.Liquidate(i)

                short = [symbol for symbol in self.short if symbol in data.Bars]
                short_weight = 0.5/len(short)
                # goes short on the prior six-month losers (lowest decile) 
                for short_symbol in short:
                    self.SetHoldings(short_symbol, -short_weight)
                    
                # goes long on the prior six-month winners (highest decile)
                long = [symbol for symbol in self.long if symbol in data.Bars]
                long_weight = 0.5/len(long)
                for long_symbol in long:
                    self.SetHoldings(long_symbol, long_weight)
            else:
                self.Liquidate()
                self.SetHoldings(self.tlt, 1)

class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    
    def __init__(self, symbol, lookback, algorithm):
        self.symbol = symbol
        self.MOMP = MomentumPercent(lookback)

        trade_bars = algorithm.History[TradeBar](symbol, lookback, Resolution.Daily)
        for trade_bar in trade_bars:
            self.MOMP.Update(trade_bar.EndTime, trade_bar.Close)