# region imports
from AlgorithmImports import *
# endregion
# https://quantpedia.com/Screener/Details/54


class MomentumAndStateOfMarketFiltersAlgorithm(QCAlgorithm):

    _blacklisted_assets = [
        "DYN RRWGXYWP5HET",   # /datasets/issue/16852
        "NIHD SML8IVC2ZZHH",  # /datasets/issue/16854
        "HERO TD6B1QW5DBHH",  # /datasets/issue/16856
    ]  

    def initialize(self):
        self.set_start_date(2011, 1, 1)
        self.set_end_date(2018, 8, 1)
        self.set_cash(100000)
        self.settings.minimum_order_margin_portfolio_percentage = 0
        # add Wilshire 5000 Total Market Index data
        self.wilshire_symbol = self.add_data(Fred, Fred.Wilshire.PRICE_5000, Resolution.DAILY).symbol
        self.w5000_return = 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 t, value in history.loc[self.wilshire_symbol]['value'].items():
            self.w5000_return.update(t, value)
        self.debug("W5000 Rate of Change indicator isReady: " + str(self.w5000_return.is_ready))    
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._coarse_selection_function)
        self.schedule.on(
            self.date_rules.month_start(Symbol.create("SPY", SecurityType.EQUITY, Market.USA)), 
            self.time_rules.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.add_equity("TLT", Resolution.DAILY).symbol        
    
    def _coarse_selection_function(self, coarse):
        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.adjusted_price)
        
        if self._month_start:
            coarse_by_symbol = {c.symbol: c for c in coarse}
            self._selection = True
            momp_ready = {
                symbol: symbol_data 
                for symbol, symbol_data in self._momp.items() 
                if (symbol_data.momp.is_ready and
                    symbol not in self._blacklisted_assets and
                    symbol in coarse_by_symbol and
                    coarse_by_symbol[symbol].adjusted_price > 5 and
                    coarse_by_symbol[symbol].has_fundamental_data)
            }
            if momp_ready:
                # sort stocks by 6-month momentum
                sort_by_momp = sorted(momp_ready, key=lambda x: momp_ready[x].momp.current.value, reverse=True)
                self._long = sort_by_momp[:20]
                self._short = sort_by_momp[-20:]
                return self._long + self._short
        return []   

    def _rebalance(self):
        # rebalance every month
        self._month_start = True
            
    def on_data(self, data):
        if self._month_start and self._selection: 
            self._month_start = False
            self._selection = False
            if not (self._long and self._short): 
                return
            
            targets = []

            # if the previous 12 months return on the broad equity market was positive
            if self.w5000_return.current.value > 0: 
                longs = [symbol for symbol in self._long if symbol in data.bars]
                long_weight = 0.25/len(longs)
                for symbol in longs:
                    targets.append(PortfolioTarget(symbol, long_weight))

                shorts = [symbol for symbol in self._short if symbol in data.bars]
                short_weight = -0.25/len(shorts)
                for symbol in shorts:
                    targets.append(PortfolioTarget(symbol, short_weight))
            else:
                targets.append(PortfolioTarget(self._tlt, 1))
            
            self.set_holdings(targets, True)


class SymbolData:
    
    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.end_time, trade_bar.close)