#region imports
from AlgorithmImports import *
#endregion
from collections import deque
import statsmodels.api as sm

class TermStructureOfVixAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2011, 1, 1)   # Set Start Date
        #self.SetEndDate(2017, 1, 1)     # Set End Date
        self.SetCash(1000000)          # Set Strategy Cash

        vix = self.AddIndex("VIX", Resolution.Daily)
        self.vix = vix.Symbol 
        self.vix_multiplier = vix.SymbolProperties.ContractMultiplier
        self.vx1 = self.AddFuture(Futures.Indices.VIX,
                resolution = Resolution.Daily,
                extendedMarketHours = True,
                dataNormalizationMode = DataNormalizationMode.BackwardsRatio,
                dataMappingMode = DataMappingMode.OpenInterest,
                contractDepthOffset = 0
            )
        self.vx1_multiplier = self.vx1.SymbolProperties.ContractMultiplier
        self.es1 = self.AddFuture(Futures.Indices.SP500EMini,
                resolution = Resolution.Daily,
                extendedMarketHours = True,
                dataNormalizationMode = DataNormalizationMode.BackwardsRatio,
                dataMappingMode = DataMappingMode.OpenInterest,
                contractDepthOffset = 0
            )
        self.es1_multiplier = self.es1.SymbolProperties.ContractMultiplier
        
        # the rolling window to save the front month VX future price
        self.price_VX = RollingWindow[float](252)
        # the rolling window to save the front month ES future price
        self.price_ES = RollingWindow[float](252)
        # the rolling window to save the time-to-maturity of the contract
        self.days_to_maturity = RollingWindow[float](252)
        
        stockPlot = Chart("Trade")
        stockPlot.AddSeries(Series("VIX", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("VIX Futures", SeriesType.Line, 0))
        stockPlot.AddSeries(Series("Buy", SeriesType.Scatter, 0))
        stockPlot.AddSeries(Series("Sell", SeriesType.Scatter, 0))
        stockPlot.AddSeries(Series("Daily Roll", SeriesType.Scatter, 0))
        stockPlot.AddSeries(Series("Hedge Ratio", SeriesType.Scatter, 0))

        self.SetWarmUp(253, Resolution.Daily)

    def OnData(self, data):
        if data.Bars.ContainsKey(self.vx1.Symbol) and data.Bars.ContainsKey(self.es1.Symbol):
            # update the rolling window price and time-to-maturity series every day
            vx1_price = data.Bars[self.vx1.Symbol].Close
            self.price_VX.Add(float(data.Bars[self.vx1.Symbol].Close))
            self.price_ES.Add(float(data.Bars[self.es1.Symbol].Close))
            self.days_to_maturity.Add((self.vx1.Mapped.ID.Date-self.Time).days)
            
        if self.IsWarmingUp or not self.price_VX.IsReady or not self.price_ES.IsReady or not self.days_to_maturity.IsReady\
            or (self.vx1.Mapped.ID.Date - self.Time).days == 0:
            return

        if data.Bars.ContainsKey(self.vx1.Symbol) and data.Bars.ContainsKey(self.es1.Symbol):
            # calculate the daily roll
            daily_roll = (vx1_price*self.vx1_multiplier - self.Securities[self.vix].Price*self.vix_multiplier)\
                /(self.vx1.Mapped.ID.Date - self.Time).days
            self.Plot("Trade", "VIX", vx1_price)
            self.Plot("Trade", "VIX Futures", vx1_price)
            self.Plot("Trade", "Daily Roll", daily_roll)
            
            if not self.Portfolio[self.vx1.Mapped].Invested:
                # Short if the contract is in contango with adaily roll greater than 0.10 
                if daily_roll > 0.1:
                    hedge_ratio = self.CalculateHedgeRatio(data.Bars[self.es1.Symbol].Close)
                    self.Plot("Trade", "Sell", vx1_price)
                    qty = self.CalculateOrderQuantity(self.vx1.Mapped, -1)
                    self.MarketOrder(self.vx1.Mapped, qty // self.vx1.SymbolProperties.ContractMultiplier)
                    qty = self.CalculateOrderQuantity(self.es1.Mapped, -1*hedge_ratio)
                    self.MarketOrder(self.es1.Mapped, qty // self.es1.SymbolProperties.ContractMultiplier)
                # Long if the contract is in backwardation with adaily roll less than -0.10
                elif daily_roll < -0.1:
                    hedge_ratio = self.CalculateHedgeRatio(data.Bars[self.es1.Symbol].Close)
                    self.Plot("Trade", "Buy", vx1_price)
                    qty = self.CalculateOrderQuantity(self.vx1.Mapped, 1)
                    self.MarketOrder(self.vx1.Mapped, qty // self.vx1.SymbolProperties.ContractMultiplier)
                    qty = self.CalculateOrderQuantity(self.es1.Mapped, 1*hedge_ratio)
                    self.MarketOrder(self.es1.Mapped, qty // self.es1.SymbolProperties.ContractMultiplier)
            
            # exit if the daily roll being less than 0.05 if holding short positions                 
            if self.Portfolio[self.vx1.Mapped].IsShort and daily_roll < 0.05:
                self.Liquidate()
                return
            
            # exit if the daily roll being greater than -0.05 if holding long positions                    
            if self.Portfolio[self.vx1.Mapped].IsLong and daily_roll > -0.05:
                self.Liquidate()
                return
                
        if self.vx1.Mapped and self.es1.Mapped:
            # if these exit conditions are not triggered, trades are exited two days before it expires
            if self.Portfolio[self.vx1.Mapped].Invested and self.Portfolio[self.es1.Mapped].Invested: 
                if (self.vx1.Mapped.ID.Date-self.Time).days <= 2 or (self.es1.Mapped.ID.Date-self.Time).days <= 2:
                    self.Liquidate()
                    
    def CalculateHedgeRatio(self, es1_price):
        price_VX = np.array(list(self.price_VX))[::-1]/self.vx1_multiplier
        price_ES = np.array(list(self.price_ES))[::-1]/self.es1_multiplier
        delta_VX = np.diff(price_VX)
        res_ES = np.diff(price_ES)/price_ES[:-1]*100
        tts = np.array(list(self.days_to_maturity))[::-1][1:]
        df = pd.DataFrame({"delta_VX":delta_VX, "SPRET":res_ES, "product":res_ES*tts}).dropna()
        # remove rows with zero value
        df = df[(df != 0).all(1)]
        y = df['delta_VX'].astype(float)
        X = df[['SPRET', "product"]].astype(float)
        X = sm.add_constant(X)
        model = sm.OLS(y, X).fit()
        beta_1 = model.params[1]
        beta_2 = model.params[2]
        
        hedge_ratio = (beta_1 + beta_2*((self.vx1.Mapped.ID.Date-self.Time).days))/float(es1_price)
        
        self.Plot("Trade", "Hedge Ratio", hedge_ratio)
        
        return hedge_ratio