#region imports
from AlgorithmImports import *
#endregion


class ShortTermReversalWithFutures(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2019, 1, 1)                                 
        self.set_end_date(2020, 9, 1)                               
        self.set_cash(10000000) 

        self.set_security_initializer(BrokerageModelSecurityInitializer(
            self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        tickers = [Futures.Currencies.CHF, 
                   Futures.Currencies.GBP, 
                   Futures.Currencies.CAD, 
                   Futures.Currencies.EUR,
                   Futures.Indices.NASDAQ_100_E_MINI, 
                   Futures.Indices.RUSSELL_2000_E_MINI, 
                   Futures.Indices.SP_500_E_MINI, 
                   Futures.Indices.DOW_30_E_MINI]
        self._length = len(tickers)

        self._symbol_data = {}
        
        for ticker in tickers:
            future = self.add_future(ticker,
                resolution=Resolution.DAILY,
                extended_market_hours=True,
                data_normalization_mode=DataNormalizationMode.BACKWARDS_RATIO,
                data_mapping_mode=DataMappingMode.OPEN_INTEREST,
                contract_depth_offset=0)
            future.set_leverage(1)
            self._symbol_data[future.symbol] = SymbolData(self, future)

    def on_data(self, data):
        for symbol, symbol_data in self._symbol_data.items():
            # Update SymbolData
            symbol_data.update(data)

            # Rollover
            if data.symbol_changed_events.contains_key(symbol):
                changed_event = data.symbol_changed_events[symbol]
                old_symbol = changed_event.old_symbol
                new_symbol = changed_event.new_symbol
                tag = f"Rollover - Symbol changed at {self.time}: {old_symbol} -> {new_symbol}"
                quantity = self.portfolio[old_symbol].quantity // self.securities[new_symbol].symbol_properties.contract_multiplier

                # Rolling over: to liquidate any position of the old mapped contract and switch to the newly mapped contract
                self.liquidate(old_symbol, tag=tag)
                if quantity: 
                    self.market_order(new_symbol, quantity, tag=tag)
        
        # Check if weekly consolidated bars are at their updatest
        if not all([symbol_data.is_ready for symbol_data in self._symbol_data.values()]):
            return
        
        # Flag to avoid undesired rebalance
        for symbol_data in self._symbol_data.values():
            symbol_data.is_volume_ready = False
            symbol_data.is_oi_ready = False
            symbol_data.is_return_ready = False
        
        # Select stocks with most weekly extreme return out of lowest volume change and highest OI change
        trade_group = set(
            sorted(self._symbol_data.values(), key=lambda x: x.volume_return)[:int(self._length*0.5)] 
            + sorted(self._symbol_data.values(), key=lambda x: x.open_interest_return)[-int(self._length*0.5):]
        )
        sorted_by_returns = sorted(trade_group, key=lambda x: x.return_)
        short_symbol = sorted_by_returns[-1].mapped
        long_symbol = sorted_by_returns[0].mapped
        
        for symbol in self.portfolio.keys():
            if self.portfolio[symbol].invested and symbol not in [short_symbol, long_symbol]:
                self.liquidate(symbol)
        
        # Adjust for contract mulitplier for order size
        qty = self.calculate_order_quantity(short_symbol, -0.3) // self.securities[short_symbol].symbol_properties.contract_multiplier
        if qty:
            self.market_order(short_symbol, qty)

        qty = self.calculate_order_quantity(long_symbol, 0.3) // self.securities[long_symbol].symbol_properties.contract_multiplier
        if qty:
            self.market_order(long_symbol, qty)


class SymbolData:

    def __init__(self, algorithm, future):
        self.is_volume_ready = False
        self.is_oi_ready = False
        self.is_return_ready = False
        self._future = future
        self._symbol = future.symbol

        # create ROC(1) indicator to get the volume and open interest return, and handler to update state
        self._volume_roc = RateOfChange(1)
        self._oi_roc = RateOfChange(1)
        self._return = RateOfChange(1)
        self._volume_roc.updated += self._on_volume_roc_updated
        self._oi_roc.updated += self._on_oi_roc_updated
        self._return.updated += self._on_return_updated

        # Create the consolidator with the consolidation period method, and handler to update ROC indicators
        self._consolidator = TradeBarConsolidator(self._consolidation_period)
        self._oi_consolidator = OpenInterestConsolidator(self._consolidation_period)
        self._consolidator.data_consolidated += self._on_trade_bar_consolidated
        self._oi_consolidator.data_consolidated += lambda sender, oi: self._oi_roc.update(oi.time, oi.value)

        # warm up
        history = algorithm.history[TradeBar](future.symbol, 14, Resolution.DAILY)
        oi_history = algorithm.history[OpenInterest](future.symbol, 14, Resolution.DAILY)
        for bar, oi in zip(history, oi_history):
            self._consolidator.update(bar)
            self._oi_consolidator.update(oi)
    
    @property
    def is_ready(self):
        return (
            self._volume_roc.is_ready and self._oi_roc.is_ready and 
            self._is_volume_ready and self._is_oi_ready and self._is_return_ready
        )

    @property
    def mapped(self):
        return self._future.mapped

    @property
    def volume_return(self):
        return self._volume_roc.current.value

    @property
    def open_interest_return(self):
        return self._oi_roc.current.value

    @property
    def return_(self):
        return self._return.current.value
    
    def update(self, data):
        if data.bars.contains_key(self._symbol):
            self._consolidator.update(data.bars[self._symbol])
            
            oi = OpenInterest(data.time, self._symbol, self._future.open_interest)
            self._oi_consolidator.update(oi)

    def _on_volume_roc_updated(self, sender, updated):
        self._is_volume_ready = True

    def _on_oi_roc_updated(self, sender, updated):
        self._is_oi_ready = True

    def _on_return_updated(self, sender, updated):
        self._is_return_ready = True
    
    def _on_trade_bar_consolidated(self, sender, bar):
        self._volume_roc.update(bar.end_time, bar.volume)
        self._return.update(bar.end_time, bar.close)
    
    def _consolidation_period(self, dt):
        # Define a consolidation period method
        period = timedelta(7)

        dt = dt.replace(hour=0, minute=0, second=0, microsecond=0)
        weekday = dt.weekday()
        if weekday > 2:
            delta = weekday - 2
        elif weekday < 2:
            delta = weekday + 5
        else:
            delta = 0
        start = dt - timedelta(delta)

        return CalendarInfo(start, period)