#region imports
from AlgorithmImports import *
#endregion


class TermStructureOfVixAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(10_000_000)
        # Add the VIX Index and Futures.
        self._vix_index = self.add_index("VIX")
        self._vx_future = self.add_future(Futures.Indices.VIX)
        self._vx_future.set_filter(14, 60)
        # Create a Scheduled Event to update the portfolio.
        self._get_liquid_price = False
        self.schedule.on(
            self.date_rules.every_day(self._vx_future.symbol),
            self.time_rules.before_market_close(self._vx_future.symbol, 15),
            self._enable_trading
        )
        # Set some parameters.
        self._daily_roll_entry_threshold = 0.05 # The paper uses +/-0.1
        self._daily_roll_exit_threshold = 0.025 # The paper uses +/-0.05
        self._min_days_to_expiry = 10
        self._spread_threshold = 0.1
    
    def _enable_trading(self):
        self._get_liquid_price = True

    def on_data(self, data):
        # For the last 15 minutes of the trading day, select the VIX Futures contract that has:
        #  - >=10 trading days until expiry
        #  - <=$0.1 bid-ask spread
        vx_price = None
        selected_contract = None
        trading_days_until_expiry = None
        if self._get_liquid_price:
            chain = data.future_chains.get(self._vx_future.symbol)
            # Get the VIX Future price.
            if chain:
                for contract in chain:
                    # Select the contract with at least 10 trading days until expiry.
                    trading_days_until_expiry = self._trading_days_until_expiry(contract)
                    if trading_days_until_expiry < self._min_days_to_expiry:
                        continue
                    # Wait for the selected contract to have a narrow spread.
                    if (contract.ask_price - contract.bid_price > self._spread_threshold and
                        # If the spread isn't narrow before market close, use the prices right before market close.
                        self._vx_future.exchange.hours.is_open(self.time + timedelta(minutes=1), extended_market_hours=False)):
                        break
                    vx_price = (contract.ask_price + contract.bid_price) / 2
                    selected_contract = contract
        if not vx_price:
            return
        self._get_liquid_price = False

        # Calculate the basis and daily roll.
        vix_price = (self._vix_index.open + self._vix_index.close) / 2
        basis = (vx_price - vix_price)
        in_contango = basis > 0
        daily_roll = basis / trading_days_until_expiry
        self.plot('VIX', 'Future Price', vx_price)
        self.plot('VIX', 'Spot Price', vix_price)
        self.plot('In Contango?', 'Value', int(in_contango))
        self.plot('Daily Roll', 'Value', daily_roll)

        # Look for trade entries.
        # "Short VIX futures positions are entered when the VIX futures basis is in
        # contango and the daily roll exceeds .10 VIX futures points ($100 per day) and long VIX futures
        # positions are entered when the VIX futures basis is in backwardation and the daily roll is less
        # than -.10 VIX futures points" (p. 14)
        if not self.portfolio.invested:
            if in_contango and daily_roll > self._daily_roll_entry_threshold:
                weight = -0.5 # Enter short VIX Future trade.
            elif not in_contango and daily_roll < -self._daily_roll_entry_threshold:
                weight = 0.5 # Enter long VIX Future trade.
            else:
                return
            self.set_holdings(selected_contract.symbol, weight, tag=f"Entry: Expires on {selected_contract.expiry}")
            self._vx_future.invested_contract = selected_contract
            return
        
        # Look for trade exits.
        # Exit long VIX Future trade when daily_roll < self._daily_roll_exit_threshold or when the contract expires next day.
        # Exit short VIX Future trade when daily_roll > -self._daily_roll_exit_threshold or when the contract expires next day.
        trading_days_until_expiry = self._trading_days_until_expiry(self._vx_future.invested_contract)
        holding = self.portfolio[self._vx_future.invested_contract.symbol]
        tag = ""
        if trading_days_until_expiry <= 1:
            tag = f"Exit: Expires in {trading_days_until_expiry} trading day(s)"
        elif holding.is_long and daily_roll < self._daily_roll_exit_threshold:
            tag = f"Exit: daily roll ({daily_roll}) < threshold"
        elif holding.is_short and daily_roll > -self._daily_roll_exit_threshold:
            tag = f"Exit: daily roll ({daily_roll}) > -threshold"
        if tag:
            self.liquidate(tag=tag)
            self._vx_future.invested_contract = None
            
    def _trading_days_until_expiry(self, contract):
        return len(list(self.trading_calendar.get_trading_days(self.time, contract.expiry - timedelta(1))))