#region imports
from AlgorithmImports import *
#endregion

class ShortTermReversalWithFutures(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)                                 
        self.SetEndDate(2020, 9, 1)                               
        self.SetCash(10000000) 

        self.tickers = [Futures.Currencies.CHF, 
                        Futures.Currencies.GBP, 
                        Futures.Currencies.CAD, 
                        Futures.Currencies.EUR,
                        Futures.Indices.NASDAQ100EMini, 
                        Futures.Indices.Russell2000EMini, 
                        Futures.Indices.SP500EMini, 
                        Futures.Indices.Dow30EMini]
        self.length = len(self.tickers)

        self.symbol_data = {}
        
        for ticker in self.tickers:
            future = self.AddFuture(ticker,
                    resolution = Resolution.Daily,
                    extendedMarketHours = True,
                    dataNormalizationMode = DataNormalizationMode.BackwardsRatio,
                    dataMappingMode = DataMappingMode.OpenInterest,
                    contractDepthOffset = 0
                )
            future.SetLeverage(1)
            self.symbol_data[future.Symbol] = SymbolData(self, future)

    def OnData(self, slice):
        for symbol, symbol_data in self.symbol_data.items():
            # Update SymbolData
            symbol_data.Update(slice)

            # Rollover
            if slice.SymbolChangedEvents.ContainsKey(symbol):
                changed_event = slice.SymbolChangedEvents[symbol]
                old_symbol = changed_event.OldSymbol
                new_symbol = changed_event.NewSymbol
                tag = f"Rollover - Symbol changed at {self.Time}: {old_symbol} -> {new_symbol}"
                quantity = self.Portfolio[old_symbol].Quantity

                # Rolling over: to liquidate any position of the old mapped contract and switch to the newly mapped contract
                self.Liquidate(old_symbol, tag = tag)
                self.MarketOrder(new_symbol, quantity // self.Securities[new_symbol].SymbolProperties.ContractMultiplier, tag = tag)
        
        # Check if weekly consolidated bars are at their updatest
        if not all([symbol_data.IsReady 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.VolumeReturn)[:int(self.length*0.5)] +
            sorted(self.symbol_data.values(), key=lambda x: x.OpenInterestReturn)[-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.CalculateOrderQuantity(short_symbol, -0.3)
        multiplier = self.Securities[short_symbol].SymbolProperties.ContractMultiplier
        self.MarketOrder(short_symbol, qty // multiplier)

        qty = self.CalculateOrderQuantity(long_symbol, 0.3)
        multiplier = self.Securities[long_symbol].SymbolProperties.ContractMultiplier
        self.MarketOrder(long_symbol, qty // multiplier)


class SymbolData:
    def __init__(self, algorithm, future):
        self._future = future
        self.Symbol = future.Symbol
        self._is_volume_ready = False
        self._is_oi_ready = False
        self._is_return_ready = False

        # 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.OnVolumeRocUpdated
        self._oi_roc.Updated += self.OnOiRocUpdated
        self._return.Updated += self.OnReturnUpdated

        # 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.DataConsolidated += self.OnTradeBarConsolidated
        self.oi_consolidator.DataConsolidated += 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 IsReady(self):
        return self._volume_roc.IsReady and self._oi_roc.IsReady \
            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 VolumeReturn(self):
        return self._volume_roc.Current.Value

    @property
    def OpenInterestReturn(self):
        return self._oi_roc.Current.Value

    @property
    def Return(self):
        return self._return.Current.Value
    
    def Update(self, slice):
        if slice.Bars.ContainsKey(self.Symbol):
            self.consolidator.Update(slice.Bars[self.Symbol])
            
            oi = OpenInterest(slice.Time, self.Symbol, self._future.OpenInterest)
            self.oi_consolidator.Update(oi)

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

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

    def OnReturnUpdated(self, sender, updated):
        self._is_return_ready = True
    
    def OnTradeBarConsolidated(self, sender, bar):
        self._volume_roc.Update(bar.EndTime, bar.Volume)
        self._return.Update(bar.EndTime, bar.Close)
    
    # Define a consolidation period method
    def consolidation_period(self, dt):
        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)