#region imports
from AlgorithmImports import *

from collections import deque
import statsmodels.api as sm
#endregion

class TermStructureOfVixAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2012, 1, 1)
        self.set_end_date(2024, 11, 1)
        self.set_cash(10_000_000)

        # Add a security initializer so we can trade ES contracts immediately.
        self.set_security_initializer(
            BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices))
        )

        self._vix_index = self.add_index("VIX")
        self._vx_future = self.add_future(Futures.Indices.VIX, leverage=1)
        self._vx_future.set_filter(14, 60)
        self._vx_future.prices = pd.Series()
        self._vx_future.trading_days_until_expiry = pd.Series()
        self._vx_future.invested_contract = None

        self._es_future = self.add_future(
            Futures.Indices.SP_500_E_MINI,
            data_normalization_mode=DataNormalizationMode.BACKWARDS_RATIO,
            data_mapping_mode=DataMappingMode.OPEN_INTEREST,
            contract_depth_offset=0
        )
        self._es_future.prices = pd.Series()

        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),
            lambda: setattr(self, '_get_liquid_price', True)
        )

        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._lookback = 252 # Not defined in the paper.
        self._enable_hedge = False #True

        self.set_warm_up(timedelta(self._lookback*1.5))

    def on_data(self, data):
        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: # Is the chain always sorted by expiry?
                    # 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 < 10:
                        continue
                    # Wait for the selected contract to have a narrow spread.
                    if (contract.ask_price - contract.bid_price > 0.1 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

        # "The spot VIX values and the mini-S&P 500 futures prices used in the study are averages of the open and closing quotes during the minute in which the above conditions for VIX futures contracts hold." (p.6)
        vix_price = (self._vix_index.open + self._vix_index.close) / 2
        es_price = (self._es_future.open + self._es_future.close) / 2 # "the front rollover adjusted mini-S&P 500 futures prices" (p.5)

        self.plot('ES', 'Price', es_price)
        self.plot('VIX', 'Future Price', vx_price)
        self.plot('VIX', 'Spot Price', vix_price)

        basis = (vx_price - vix_price)
        in_contango = basis > 0
        self.plot('In Contango?', 'Value', int(in_contango))

        daily_roll = basis / trading_days_until_expiry
        self.plot('Daily Roll', 'Value', daily_roll)

        # "the size of the mini-S&P futures hedge is based on out of sample hedge ratio 
        # estimates. The hedge ratios are constructed from regressions of VIX futures price changes on a
        # constant and on contemporaneous percentage changes of the front mini-S&P 500 futures contract
        # both alone and multiplied by the number of business days that the VIX futures contract is from
        # settlement" (pp. 11-12)
        t = self.time
        self._es_future.prices.loc[t] = es_price
        self._vx_future.prices.loc[t] = vx_price
        self._vx_future.trading_days_until_expiry.loc[t] = trading_days_until_expiry

        self._es_future.prices = self._es_future.prices.iloc[-self._lookback:]
        self._vx_future.prices = self._vx_future.prices.iloc[-self._lookback:]
        self._vx_future.trading_days_until_expiry = self._vx_future.trading_days_until_expiry.iloc[-self._lookback:]

        # Wait until there are sufficient samples to fit the regression model.
        if len(self._vx_future.prices) < self._lookback or self.is_warming_up:
            return
        vx_changes = self._vx_future.prices.pct_change()[1:]
        es_returns = self._es_future.prices.pct_change()[1:]
        X = pd.DataFrame(
            {
                'es_returns': es_returns,
                'es_returns*tts': es_returns * self._vx_future.trading_days_until_expiry[1:]
            },
            index=es_returns.index
        )
        X = sm.add_constant(X)
        y = vx_changes
        model = sm.OLS(y, X).fit()
        beta_1 = model.params[1]
        beta_2 = model.params[2]

        # "The β1 coefficient should be significantly negative in light of the tendency of VIX futures prices to move inversely to equity returns." (p. 12)
        self.plot('Regression B1', 'Value', beta_1)
        # "The β2 coefficient should be significantly positive if the reaction of VIX futures prices to equity returns is more subdued the further contracts are from settlement." (p. 12)
        self.plot('Regression B2', 'Value', beta_2)

        # Calculate the hedge ratio. Equation 4 on page 13.
        hedge_ratio = (
            beta_1 * 1_000
            + beta_2 * trading_days_until_expiry * 1_000
        ) / (0.01 * es_price * 50)
        # "The average hedge ratio is close to one mini-S&P futures contract per VIX futures contract" (p. 13) with a range from about 0.5 to 2.
        self.plot('Hedge Ratio', 'Value', hedge_ratio)

        # 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

            # Add ES hedge.
            if self._enable_hedge:
                es_contract_symbols = self.future_chain_provider.get_future_contract_list(self._es_future.symbol, self.time)
                es_contract_symbol = sorted(
                    [s for s in es_contract_symbols if s.id.date >= self._vx_future.invested_contract.expiry], 
                    key=lambda symbol: symbol.id.date
                )[0]
                self.add_future_contract(es_contract_symbol)
                self.set_holdings(es_contract_symbol, weight * -np.sign(hedge_ratio))
            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))))