#region imports
from AlgorithmImports import *
#endregion
class VixCalendarSpread(QCAlgorithm):
    
    def Initialize(self):

        self.cap = 400000
        self.leverage = 1.
        self.SetStartDate(2021, 12, 8)  # Set Start Date
        self.SetCash(self.cap)
        
        self.vix = self.AddFuture(
            ticker=Futures.Indices.VIX,
            resolution=Resolution.Hour,
            dataNormalizationMode = DataNormalizationMode.Raw,
            dataMappingMode = DataMappingMode.LastTradingDay,
            contractDepthOffset = 0
        )

        self.AddEquity("SPY")
        self.Schedule.On(
            self.DateRules.EveryDay("SPY"),
            self.TimeRules.AfterMarketOpen("SPY", 60),
            self.Rebalance)

        self.orderWeights = {}
        self.isExpired = False
        
    def OnData(self, slice) -> None:
        for changedEvent in slice.SymbolChangedEvents.Values:
            if changedEvent.Symbol == self.vix.Symbol: self.isExpired = True
        
        if not self.isExpired: return

        self.chain = slice.FuturesChains.get(self.vix.Symbol.Canonical)
        
        if not self.chain: return
        
        contracts = [c for c in self.chain.Value]
        if len(contracts) < 1: return
        
        self.isExpired = False
        sortedChain = sorted(contracts, key = lambda x: x.Expiry, reverse=False)
        self.contractSymbols = [x.Symbol for x in sortedChain[:2]]
        front, second = self.contractSymbols
        self.orderWeights[front] = -1
        self.orderWeights[second] = 1
        self.Debug(f"Short {sortedChain[0].Symbol}, long {sortedChain[1].Symbol}")
    
    def Rebalance(self):
        if len(self.orderWeights) < 0: return
        for symbol, holding in self.Portfolio.items():
            if symbol in self.orderWeights:
                self.orderWeights[symbol] -= holding.Quantity
            self.orderWeights[symbol] = -holding.Quantity

        for symbol, quantity in self.orderWeights.items():
            if quantity!=0:
                self.MarketOrder(symbol, quantity)