#region imports
from AlgorithmImports import *

from collections import deque
#endregion
# https://quantpedia.com/Screener/Details/155


class MomentumReversalCombinedWithVolatility(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
        # Define some parameters.
        self._num_portfolios = 6
        self._lookback = 20*6
        # Add a universe of US Equities.
        self._date_rule = self.date_rules.month_start('SPY')
        self.universe_settings.schedule.on(self._date_rule)
        self.universe_settings.resolution = Resolution.DAILY
        self._universe = self.add_universe(self._select_assets)
        # We'll churn 1/6 of the portfolio during each rebalance.
        # Let's add a deque to help us keep track.
        self._portfolios = deque(maxlen=self._num_portfolios)
        # Add a warm-up period so the algorithm trades on deployment
        # instead of waiting for the next month.
        self.set_warm_up(timedelta(self._lookback))

    def _select_assets(self, fundamentals):
        # Select stocks trading above $5.
        filtered = [
            f for f in fundamentals 
            if f.has_fundamental_data and f.price > 5 and not np.isnan(f.market_cap)
        ]
        # Take the 50% of stocks with that have higher market caps.
        filtered = sorted(filtered, key=lambda f: f.market_cap)[-int(len(filtered)/2):]
        return [f.symbol for f in filtered]
        
    def on_securities_changed(self, changes):
        # As assets enter the universe, warm up their trailing price data.
        for security in changes.added_securities:            
            security.session.size = self._lookback
            for bar in self.history[TradeBar](security, self._lookback):
                security.session.update(bar)
    
    def on_warmup_finished(self):
        # Add a Scheduled Event to rebalance the portfolio each month.
        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 securities in the universe that have sufficient
        # trailing data.
        securities = []
        for symbol in self._universe.selected:
            security = self.securities[symbol]
            if security.session.is_ready and security.price:
                securities.append(security)
        if len(securities) < 100:
            return
        # Select the most volatile subset of stocks.
        volatile_stocks = sorted(securities, key=self._volaility)[-int(0.2*len(securities)):]
        # Sort stocks by their trailing returns and then split into long/short sides.
        sorted_by_return = sorted(volatile_stocks, key=self._return)
        longs = sorted_by_return[-int(0.2*len(sorted_by_return)):]
        shorts = sorted_by_return[:int(0.2*len(sorted_by_return))]        
        # If the portfolio is full, liquidate the oldest 1/6th of holdings.
        if len(self._portfolios) == self._portfolios.maxlen:
            for security, quantity in list(self._portfolios)[0].items():
                if security.is_tradable and security.price:
                    self.market_on_open_order(security, -quantity)
        # For this new 1/6th of the portfoilo, form a long-short portfolio.
        quantity_by_security = (
            self._get_target_quantities(longs, 1) 
            | self._get_target_quantities(shorts, -1)
        )
        self._portfolios.append(quantity_by_security)
        for security, quantity in quantity_by_security.items():
            self.market_on_open_order(security, quantity)
    
    def _get_target_quantities(self, securities, trade_direction):
        # Calculate the trade quantity of each asset for this rebalance.
        quantity_by_security = {}
        for security in securities:
            quantity = trade_direction * int(
                self.portfolio.total_portfolio_value 
                / 2 
                / len(securities) 
                / security.price 
                / self._num_portfolios
            )
            if quantity:
                quantity_by_security[security] = quantity
        return quantity_by_security

    def _get_history(self, security):
        # Get the trailing bars from the security.session.
        return [bar.close for bar in list(security.session)[::-1]][:-1]

    def _volaility(self, security):
        history = self._get_history(security)
        # one week (5 trading days) prior to the beginning of each month is skipped 
        prices = np.array(history)[:-5]
        returns = (prices[1:]-prices[:-1])/prices[:-1]
        # calculate the annualized realized volatility
        return np.std(returns)*np.sqrt(250/len(returns))

    def _return(self, security):
        history = self._get_history(security)
        # one week (5 trading days) prior to the beginning of each month is skipped 
        prices = np.array(history)[:-5]
        # calculate the annualized realized return
        return (prices[-1]-prices[0])/prices[0]