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


class MomentumEffectAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        self.settings.seed_initial_prices = True
        self.settings.automatic_indicator_warm_up = True
        # Define some parameters.
        self._lookback = 252
        self._liquidity_filter_size = 100
        self._universe_size = 50
        self._portfolio_size = 5
        # Add a universe of US Equities.
        self.universe_settings.resolution = Resolution.DAILY
        self._date_rule = self.date_rules.month_start('SPY')
        self.universe_settings.schedule.on(self._date_rule)
        self._universe = self.add_universe(self._select_assets)
        # Add a warm-up period so the algorithm trades right away
        # instead of waiting for the next month to start.
        self.set_warm_up(timedelta(45))

    def _select_assets(self, fundamentals):
        # Select the most liquid stocks trading above $5/share.
        filtered = [f for f in fundamentals if f.has_fundamental_data and f.price > 5]
        filtered = sorted(filtered, key=lambda f: f.dollar_volume)[-self._liquidity_filter_size:]
        # Select the subset of stocks with highest market caps.
        filtered = sorted(filtered, key=lambda f: f.market_cap, reverse=True)[:self._universe_size]
        return [f.symbol for f in filtered]

    def on_securities_changed(self, changes):
        # As assets enter the universe, add a momentum indicator to each one.
        for security in changes.added_securities:
            security.momentum = self.momp(security, self._lookback)
        # As assets leave the universe, remove their indicator and liquidate it.
        for security in changes.removed_securities:
            self.deregister_indicator(security.momentum)
            self.liquidate(security, 'Removed from universe')

    def on_warmup_finished(self):
        # Add a Scheduled Event to rebalance the portofolio monthly.
        time_rule = self.time_rules.at(8, 0)
        self.schedule.on(self._date_rule, time_rule, self._rebalance)
        # Rebalance today too.
        if self.live_mode:
            self._rebalance()
        else:
            self.schedule.on(self.date_rules.today, time_rule, self._rebalance)

    def _rebalance(self):
        # Get the stocks that have their momentum indicators primed.
        securities = []
        for symbol in self._universe.selected:
            security = self.securities[symbol]
            if security.momentum.is_ready:
                securities.append(security)
        # Select the securities with greatest momentum.
        selected = sorted(securities, key=lambda security: security.momentum)[-self._portfolio_size:]
        # Form an equal-weighted portfolio.
        targets = [PortfolioTarget(security, 1/len(selected)) for security in selected]
        self.set_holdings(targets, True)