#region imports
from AlgorithmImports import *
#endregion


class ShortTimeReversal(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2016, 1, 1)
        self.set_end_date(2021, 1, 1)
        self.set_cash(1000000)
        
        self.universe_settings.resolution = Resolution.DAILY
        self.add_universe(self._select_coarse)
        
        self._dollar_volume_selection_size = 100
        self._roc_selection_size = int(0.1 * self._dollar_volume_selection_size)
        
        self._lookback = 22
        self._roc_by_symbol = {}
        self._week = 0

    def _select_coarse(self, coarse):
        # We should keep a dictionary for all securities that have been selected
        for cf in coarse:
            symbol = cf.symbol
            if symbol in self._roc_by_symbol:
                self._roc_by_symbol[symbol].update(cf.end_time, cf.adjusted_price)

        # Refresh universe each week
        week_number = self.time.date().isocalendar()[1]
        if week_number == self._week:
            return Universe.UNCHANGED
        self._week = week_number

        # sort and select by dollar volume
        sorted_by_dollar_volume = sorted(coarse, key=lambda x: x.dollar_volume, reverse=True)
        selected = {cf.symbol: cf for cf in sorted_by_dollar_volume[:self._dollar_volume_selection_size]} 
        
        # New selections need a history request to warm up the indicator
        symbols = [k for k in selected.keys()
            if k not in self._roc_by_symbol or not self._roc_by_symbol[k].is_ready]

        if symbols:
            history = self.history(symbols, self._lookback+1, Resolution.DAILY)
            if history.empty:
                self.log(f'No history for {", ".join([x.value for x in symbols])}')
            history = history.close.unstack(0)

            for symbol in symbols:
                symbol_id = symbol.id.to_string()
                if symbol_id not in history:
                    continue

                # Create and warm-up the RateOfChange indicator
                roc = RateOfChange(self._lookback)
                for time, price in history[symbol_id].dropna().items():
                    roc.update(time, price)
                
                if roc.is_ready:
                    self._roc_by_symbol[symbol] = roc
        
        # Sort the symbols by their ROC values
        selected_rate_of_change = {}
        for symbol in selected.keys():
            if symbol in self._roc_by_symbol:
                selected_rate_of_change[symbol] = self._roc_by_symbol[symbol]
        sorted_by_rate_of_change = sorted(selected_rate_of_change.items(), key=lambda kv: kv[1], reverse=True)
        
        # Define the top and the bottom to buy and sell
        self._roc_top = [x[0] for x in sorted_by_rate_of_change[:self._roc_selection_size]]
        self._roc_bottom = [x[0] for x in sorted_by_rate_of_change[-self._roc_selection_size:]]
        
        return self._roc_top + self._roc_bottom

    def on_data(self, data):
        # Rebalance
        for symbol in self._roc_top:
            self.set_holdings(symbol, -0.5/len(self._roc_top))
        for symbol in self._roc_bottom:
            self.set_holdings(symbol, 0.5/len(self._roc_bottom))
        
        # Clear the list of securities we have placed orders for
        # to avoid new trades before the next universe selection
        self._roc_top.clear() 
        self._roc_bottom.clear()

    def on_securities_changed(self, changes):
        for security in changes.removed_securities:
            self.liquidate(security.symbol, 'Removed from Universe')