import numpy as np
from QuantConnect.Python import PythonQuandl
from QuantConnect.Data.Custom import *
# from expiry import calendar
import datetime
import pandas as pd
from datetime import timedelta
from collections import deque
import statsmodels.api as sm
from decimal import Decimal

class TermStructureOfVixAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2011, 1, 1)   # Set Start Date
        self.SetEndDate(2011, 9, 1)     # Set End Date
        self.SetCash(10000000)          # Set Strategy Cash
        


        self.vix = self.AddData(QuandlVix, "CBOE/VIX", Resolution.Daily).Symbol              # Add Quandl VIX price (daily)
        self.vx1 = self.AddData(QuandlFutures, "CHRIS/CBOE_VX1", Resolution.Daily).Symbol    # Add Quandl VIX front month futures data (daily)
        self.es1 = self.AddData(QuandlFutures, "CHRIS/CME_ES1", Resolution.Daily).Symbol     # Add Quandl E-mini S&P500 front month futures data (daily)

        self.Settings.FreePortfolioValuePercentage = 0.5
        
        self.canonical_symbol_by_index = {}
        for index in [Futures.Indices.VIX, Futures.Indices.SP500EMini]:
            future = self.AddFuture(index)
            future.SetFilter(timedelta(0), timedelta(days=180))
            self.canonical_symbol_by_index[index] = future.Symbol
            
        self.front_VX = None
        self.front_ES = None
        # request the history to warm-up the price and time-to-maturity 
        hist = self.History([self.vx1, self.es1], timedelta(days=450), Resolution.Daily)
        settle = hist['settle'].unstack(level=0)
        
        # the rolling window to save the front month VX future price        
        self.price_VX = deque(maxlen=252)   
        # the rolling window to save the front month ES future price        
        self.price_ES = deque(maxlen=252)   
        # the rolling window to save the time-to-maturity of the contract
        self.days_to_maturity = deque(maxlen=252) 
        
        expiry_date = self.get_expiry_calendar()
        df = pd.concat([settle, expiry_date], axis=1, join='inner')

        for index, row in df.iterrows():
            self.price_VX.append(row[str(self.vx1)])
            self.price_ES.append(row[str(self.es1)])
            self.days_to_maturity.append((row['expiry']-index).days)
        
        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))

    def OnData(self, data):
        has_data = False
        if self.front_VX and self.front_ES:
            has_data = data.ContainsKey(self.front_ES.Symbol) and \
                        data[self.front_ES.Symbol] is not None and \
                        data.ContainsKey(self.front_VX.Symbol) and \
                        data[self.front_VX.Symbol] is not None
        
        # select the nearest VIX and E-mini S&P500 futures with at least 10 trading days to maturity 
        # if the front contract expires, roll forward to the next nearest contract
        for chain in data.FutureChains:
            if chain.Key == self.canonical_symbol_by_index[Futures.Indices.VIX]:
                if self.front_VX is None or ((self.front_VX.Expiry-self.Time).days <= 1):
                    contracts = list(filter(lambda x: x.Expiry >= self.Time + timedelta(days = 10), chain.Value))
                    self.front_VX = sorted(contracts, key = lambda x: x.Expiry)[0]
            if chain.Key == self.canonical_symbol_by_index[Futures.Indices.SP500EMini]:
                if self.front_ES is None or ((self.front_ES.Expiry-self.Time).days <= 1):
                    contracts = list(filter(lambda x: x.Expiry >= self.Time + timedelta(days = 10), chain.Value))
                    self.front_ES = sorted(contracts, key = lambda x: x.Expiry)[0]
        
        if data.ContainsKey(self.vx1) and data.ContainsKey(self.es1):
            # update the rolling window price and time-to-maturity series every day
            if self.front_VX and self.front_ES:
                self.price_VX.append(float(self.Securities[self.vx1].Price))
                self.price_ES.append(float(self.Securities[self.es1].Price))
                self.days_to_maturity.append((self.front_VX.Expiry-self.Time).days)
            
                # calculate the daily roll
                daily_roll = (self.Securities[self.vx1].Price - self.Securities[self.vix].Price)/(self.front_VX.Expiry-self.Time).days
                self.Plot("Trade", "VIX", self.Securities[self.vix].Price)
                self.Plot("Trade", "VIX Futures", self.Securities[self.vx1].Price)
                self.Plot("Trade", "Daily Roll", daily_roll)
                
                if not has_data:
                    return
                
                if not self.Portfolio[self.front_VX.Symbol].Invested:
                    # Short if the contract is in contango with adaily roll greater than 0.10 
                    if daily_roll > 0.1:
                        hedge_ratio = self.CalculateHedgeRatio()
                        self.Plot("Trade", "Sell", self.Securities[self.vx1].Price)
                        self.SetHoldings(self.front_VX.Symbol, -0.5)
                        self.SetHoldings(self.front_ES.Symbol, -0.5*hedge_ratio)
                        self.Log(f"Short at {self.Time}; has_data {has_data}")
                    # Long if the contract is in backwardation with adaily roll less than -0.10
                    elif daily_roll < -0.1:
                        hedge_ratio = self.CalculateHedgeRatio()
                        self.Plot("Trade", "Buy", self.Securities[self.vx1].Price)
                        self.SetHoldings(self.front_VX.Symbol, 0.5)
                        self.SetHoldings(self.front_ES.Symbol, 0.5*hedge_ratio)
                        self.Log(f"Long at {self.Time}; has_data {has_data}")
                
                # exit if the daily roll being less than 0.05 if holding short positions                 
                if self.Portfolio[self.front_VX.Symbol].IsShort and daily_roll < 0.05:
                    self.Liquidate()
                    self.front_VX = None
                    self.front_ES = None
                    return
                
                # exit if the daily roll being greater than -0.05 if holding long positions                    
                if self.Portfolio[self.front_VX.Symbol].IsLong and daily_roll > -0.05:
                    self.Liquidate()
                    self.front_VX = None
                    self.front_ES = None
                    return

                
        if self.front_VX and self.front_ES and has_data:
            # if these exit conditions are not triggered, trades are exited two days before it expires
            if self.Portfolio[self.front_VX.Symbol].Invested and self.Portfolio[self.front_ES.Symbol].Invested: 
                if (self.front_VX.Expiry-self.Time).days <=2 or (self.front_ES.Expiry-self.Time).days <=2:
                    self.Log(f"Liquidate for expiry at {self.Time}; has_data {has_data}")
                    self.Liquidate()
                    self.front_VX = None
                    self.front_ES = None
                    return
                
                    
    def CalculateHedgeRatio(self):
        
        price_VX = np.array(self.price_VX)
        price_ES = np.array(self.price_ES)
        delta_VX = np.diff(price_VX)
        res_ES = np.diff(price_ES)/price_ES[:-1]*100
        tts = np.array(self.days_to_maturity)[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 = abs((1000*beta_1 + beta_2*((self.front_VX.Expiry-self.Time).days)*1000)/(0.01*50*float(self.Securities[self.es1].Price)))
        
        self.Plot("Trade", "Hedge Ratio", hedge_ratio)
        
        return hedge_ratio

    def get_expiry_calendar(self):
        # import the futures expiry calendar
        url = "https://www.dropbox.com/s/5k4rbuzfsfn3w0h/expiry.csv?dl=1"
        data = self.Download(url).split('\r\n')
        expiry = [x.split(',')[1] for x in data][1:]
        date = [x.split(',')[0] for x in data][1:]
        df_date = pd.DataFrame(expiry, index = date, columns = ['expiry'])
        df_date.index = pd.to_datetime(df_date.index)
        df_date['expiry'] = pd.to_datetime(df_date['expiry'])
        idx = pd.date_range('01-01-2011', '04-19-2019')
        # populate the date index and backward fill the dataframe    
        return df_date.reindex(idx, method='bfill')

class QuandlVix(PythonQuandl):
    def __init__(self):
        self.ValueColumnName = "vix Close"

class QuandlFutures(PythonQuandl):
    def __init__(self):
        self.ValueColumnName = "settle"