#region imports
from AlgorithmImports import *

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


# Use Minute resolution data.
# To get closing VIX quotes, get first quote during last 15 minutes of trading day where the bid-ask spread is no greater than 0.1 VIX future point or $100, as one VIX futures point is worth $1K. In the absense of such a quote, use the final bid/ask quotes of the day.
# Use front rollover adjusted mini-S&P 500 futures prices.
# The spot VIX values and the mini-S&P 500 futures prices are the averages of the open and closing price during the minute in which the above conditions for VIX futures contract hold


# To calculate daily roll
#  the spread b/w the price of the front VIX futures contract that has 10+ business days until settlement and the VIX, scaled by the number of business days until settlement.


# ** Use number of business days to settlement, not number of days.

# Enter positions in the last 15 minutes of the trading day at the first instance that the bid-ask spread is no greater than 0.1, and in the absense of such opportunities, at the final bid-ask spread recorded at the close.
#  - compare mid-point of VIX futures to average of the open and close price of the VIX during the same minute.

# ** Do regression to get the Beta coef => Then use coeff to calculate hedge ratio
#    ** Hedge should average around 1 contract for every 1 VIX contract.

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.market_order(es_contract_symbol, (1 if weight > 0 else -1) * -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))))