#region imports
from AlgorithmImports import *

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


class CombiningMomentumEffectWithVolume(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._date_rule = self.date_rules.month_start('SPY')
        self.universe_settings.schedule.on(self._date_rule)
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._select_assets)

        # Define some members to help with the ROC calculation.
        self._lookback = 252
        self._roc_by_symbol = {}

        self.set_warmup(timedelta(45))
        self._portfolios = deque(maxlen=3)


    def _select_assets(self, fundamentals):
        # Drop stocks which have no fundamental data.
        fundamentals = [f for f in fundamentals if f.has_fundamental_data]
        # Update the ROC of each asset and calculate its turnover.
        candidates = []
        turnover_by_symbol = {}
        for f in fundamentals:
            if f.symbol not in self._roc_by_symbol:
                roc = self._roc_by_symbol.setdefault(f.symbol, RateOfChange(self._lookback))
                for bar in self.history[TradeBar](f.symbol, self._lookback, Resolution.DAILY):
                    roc.update(bar.end_time, bar.close)
            else:
                roc = self._roc_by_symbol[f.symbol]
                roc.update(self.time, f.adjusted_price)
            if (not roc.is_ready or 
                not f.earning_reports.basic_average_shares.three_months or
                not f.volume):
                continue
            candidates.append(f.symbol)
            turnover_by_symbol[f.symbol] = f.earning_reports.basic_average_shares.three_months / f.volume
        # Split candidates into top and bottom baskets based on their ROC.
        sorted_by_roc = sorted(candidates, key=lambda symbol: self._roc_by_symbol[symbol])
        roc_size = int(0.2*len(sorted_by_roc))
        top_roc = sorted_by_roc[-roc_size:]
        bottom_roc = sorted_by_roc[:roc_size]
        # Select the assets with the greatest turnover for the long & short baskets.
        assets_per_side = int(0.01*roc_size)
        self._longs = sorted(top_roc, key=lambda symbol: turnover_by_symbol[symbol])[-assets_per_side:]
        self._shorts = sorted(bottom_roc, key=lambda symbol: turnover_by_symbol[symbol])[-assets_per_side:]
        return self._longs + self._shorts
            
    def on_warmup_finished(self):
        # Add a Scheduled event to rebalance the portfolio yearly.
        time_rule = self.time_rules.at(8, 0)
        self.schedule.on(self._date_rule, time_rule, self._rebalance)
        # Rebalance the portfolio today too.
        if self.live_mode:
            self._rebalance()
        else:
            self.schedule.on(self.date_rules.today, time_rule, self._rebalance)

    def _rebalance(self):
        if self.is_warming_up:
            return
            
        # 1/3 of the portfolio is rebalanced every month
        if len(self._portfolios) == self._portfolios.maxlen:
            for position in list(self._portfolios)[0]:
                position.exit()
                
        self._portfolios.append(
            self._open_trades(self._shorts, -1) + 
            self._open_trades(self._longs, 1)
        )
    
    def _open_trades(self, symbols, direction):
        new_positions = []
        for symbol in symbols:
            security = self.securities[symbol]
            quantity = direction * int(
                0.5*self.portfolio.total_portfolio_value 
                / security.price 
                / len(symbols)
                / self._portfolios.maxlen
            )
            if quantity:
                new_positions.append(MonthlyPosition(self, security, quantity))
        return new_positions
        

class MonthlyPosition:

    def __init__(self, algorithm, equity, quantity):
        self._algorithm = algorithm
        self._equity = equity
        self._quantity = quantity
        # Enter the position.
        algorithm.market_order(equity, quantity)

    def exit(self):
        if self._equity.is_tradable:
            self._algorithm.market_order(self._equity, -self._quantity)