# ref
# Alexi Muci, A simple VIX Strategy, https://www.quantconnect.com/forum/discussion/2657/a-simple-vix-strategy
# Tony Cooper, Easy Volatility Investing, https://www.ssrn.com/abstract=2255327

from QuantConnect import *
from QuantConnect.Algorithm import *

import pandas as pd
import numpy as np

from my_custom_data import CboeVix, CboeVxV


class VIXStrategyByRatio(QCAlgorithm):
    
    def Initialize(self):
        # SVXY inception date 10/3/2011
        # VXX inception date 1/29/2009
        
        self.SetStartDate(2011, 10, 15)
        #self.SetStartDate(2019, 1, 15)
        self.SetEndDate(datetime.now().date() - timedelta(1))
        self.SetCash(10000)
            
        #self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin) 
        # ^^ with a low alpha model that trades daily, you get eaten alive by brokerage fees if you close positions everyday with starting capital of 10k.
        self.SetBrokerageModel(BrokerageName.AlphaStreams)

        # add the #2 ETFs (short and long VIX futures)
        # swap xiv with svxy, since xiv is wiped out.
        self.XIV = self.AddEquity("SVXY", Resolution.Daily).Symbol
        self.VXX = self.AddEquity("VXX", Resolution.Daily).Symbol
        self.SPY = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.SHY = self.AddEquity("SHY", Resolution.Daily).Symbol
        
        self.SetBenchmark("SPY")
        
        # Define symbol and "type" of custom data: used for signal ratio
        self.VIX = self.AddData(CboeVix, "VIX").Symbol
        self.VXV = self.AddData(CboeVxV, "VXV").Symbol
        
        self.window_len = 252
        hist = self.History([self.VIX], self.window_len, Resolution.Daily)
        
        self.window_vix = RollingWindow[float](self.window_len)
        for close in hist.loc[self.VIX]['close']:
            self.window_vix.Add(close)
        
        self.window_vix_date = RollingWindow[str](self.window_len)
        for item in hist.index:
            self.window_vix_date.Add(item[-1].strftime('%Y-%m-%d'))
        
        hist = self.History([self.VXV], self.window_len, Resolution.Daily)
        
        self.window_vxv = RollingWindow[float](self.window_len)
        for close in hist.loc[self.VXV]['close']:
            self.window_vxv.Add(close)
            
        self.window_vxv_date = RollingWindow[str](self.window_len)
        for item in hist.index:
            self.window_vxv_date.Add(item[-1].strftime('%Y-%m-%d'))
            
        # Define the Schedules
        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.AfterMarketOpen(self.XIV, 0),
            Action(self.Balance)
        )
        #self.Schedule.On(
        #    self.DateRules.EveryDay(),
        #    self.TimeRules.BeforeMarketClose(self.XIV, 5),
        #    Action(self.Close)
        #)
        self.SetWarmUp(timedelta(self.window_len))
        
    
    # TODO: verify data correctness.
    # TODO: is there a better way to collect data?    
    def OnData(self, data):
        
        if data.ContainsKey(self.VIX):
            self.window_vix.Add(data[self.VIX].Price)
            self.window_vix_date.Add(data.Time.strftime('%Y-%m-%d'))
            
        if data.ContainsKey(self.VXV):
            self.window_vxv.Add(data[self.VXV].Price)
            self.window_vxv_date.Add(data.Time.strftime('%Y-%m-%d'))
        
    def Close(self):
        for x in [self.SPY,self.SHY,self.XIV,self.VXX]:
            self.SetHoldings(x, 0.0)
            
    def Balance(self):

        if not self.window_vxv.IsReady: return
        if not self.window_vix.IsReady: return
        if not self.window_vxv_date.IsReady: return
        if not self.window_vix_date.IsReady: return
        
        # TODO: compute needs to be moved to ??? alpha compute method.
        # then alpha can be used to aid in construct portfolio.
        
        df = pd.DataFrame()
        # flip the data, so last item is the most current.
        vix_date_array = [i for i in self.window_vix_date][::-1]
        vxv_date_array = [i for i in self.window_vxv_date][::-1]
        vix_price = np.array([i for i in self.window_vix])[::-1]
        vxv_price = np.array([i for i in self.window_vxv])[::-1]
        
        df['vix_date'] = vix_date_array
        df['vxv_date'] = vxv_date_array
        df['VIX'] = vix_price
        df['VXV'] = vxv_price
        
        # avoid look ahead bias.
        df['VIX'] = df['VIX'].shift(1)
        df['VXV'] = df['VXV'].shift(1)
        
        df["cob"] = df["VIX"]/df["VXV"]
        
        df["cob"] = df["cob"].rolling(10).median() # "Strategy 3, Vratio10 by Tony Cooper"
        cob = df["cob"].iloc[-1]
        
        XIV_qnty = self.Portfolio[self.XIV].Quantity
        VXX_qnty = self.Portfolio[self.VXX].Quantity
        
        self.Log("{},{},{},{},{}".format(
            df['vix_date'].iloc[-1],
            df['vxv_date'].iloc[-1],
            df['VIX'].iloc[-1],
            df['VXV'].iloc[-1],
            cob))
        
        # cob = vix/vxv
        # cob < 1 , vix < vxv: contago
        # cob > 1 , vix > vxv: backwardation (1 mnth more expensive than 3 mnth future)
        # https://en.wikipedia.org/wiki/Contango
        # https://en.wikipedia.org/wiki/Normal_backwardation
        # 
        #np.nanpercentile(df['VIX/VXV'],[30,40,50,60,70,80,90])
        # >>> array([0.86649373, 0.88454818, 0.9025271 , 0.92344436, 0.94629521, 0.97491226, 1.01362785])
        
        # short vol (buy XIV)
        if cob < 0.95:
            situation = 'short_volatility'
            self.Log("short VOL")
            #self.Notify.Email("XXXXX@gmail.com", "IB Algo Execution", "long XIV"); self.Log("short VOL")
            Insight(self.VXX, timedelta(days=1), InsightType.Price, InsightDirection.Down, confidence=1.0)
            Insight(self.XIV, timedelta(days=1), InsightType.Price, InsightDirection.Up, confidence=1.0)
        
        # long vol (buy VXX)
        elif cob > 1.05:
            situation = 'long_volatility'
            self.Log("long VOL")
            #self.Notify.Email("XXXX@gmail.com", "IB Algo Execution", "long VXX"); self.Log("long VOL")
            Insight(self.XIV, timedelta(days=1), InsightType.Price, InsightDirection.Down, confidence=1.0)
            Insight(self.VXX, timedelta(days=1), InsightType.Price, InsightDirection.Up, confidence=1.0)
            
        # flat ? - not really sure what to do - chicken mode
        else:
            situation = 'flat'
            self.Log("Flat")
            #self.Notify.Email("xxxxxxxxxx@gmail.com", "IB Algo Execution", "Flat position"); self.Log("Flat")
            Insight(self.XIV, timedelta(days=1), InsightType.Price, InsightDirection.Flat, confidence=0.3)
            Insight(self.VXX, timedelta(days=1), InsightType.Price, InsightDirection.Flat, confidence=0.3)
            
        iama_braveheart = {
            'short_volatility':{ self.SPY: 0.0, self.SHY: 0.0, self.XIV: 1.0, self.VXX: 0.0, },
            'flat':            { self.SPY: 0.0, self.SHY: 0.0, self.XIV: 0.0, self.VXX: 0.0, },
            'long_volatility': { self.SPY: 0.0, self.SHY: 0.0, self.XIV: 0.0, self.VXX: 1.0, },
        }
        ihave_babies_and_30_year_home_mortgage = { # blend in some SHY and SPY to reduce drawdowns.
            'short_volatility':{ self.SPY: 0.5, self.SHY: 0.0, self.XIV: 0.5, self.VXX: 0.0, },
            'flat':            { self.SPY: 0.5, self.SHY: 0.5, self.XIV: 0.0, self.VXX: 0.0, },
            'long_volatility': { self.SPY: 0.3, self.SHY: 0.7, self.XIV: 0.0, self.VXX: 0.0, },
        }
        
        mydict = ihave_babies_and_30_year_home_mortgage # use braveheart param for wow factor.
        for k,v in mydict[situation].items():
            self.SetHoldings(k,v)

from QuantConnect.Python import PythonQuandl # quandl data not CLOSE
from QuantConnect.Python import PythonData # custom data
from QuantConnect.Data import SubscriptionDataSource

from datetime import datetime, timedelta
import decimal

class CboeVix(PythonData):
    '''CBOE Vix Download Custom Data Class'''
    def GetSource(self, config, date, isLiveMode):
        url_vix = "http://www.cboe.com/publish/scheduledtask/mktdata/datahouse/vixcurrent.csv"
        return SubscriptionDataSource(url_vix, 
                                      SubscriptionTransportMedium.RemoteFile)
    def Reader(self, config, line, date, isLiveMode):
        if not (line.strip() and line[0].isdigit()): return None
        # New CboeVix object
        index = CboeVix();
        index.Symbol = config.Symbol
        try:
            # Example File Format:
            # Date          VIX Open    VIX High VIX Low    VIX Close
            # 01/02/2004    17.96    18.68     17.54        18.22
            #print line
            data = line.split(',')
            date = data[0].split('/')
            index.Time = datetime(int(date[2]), int(date[0]), int(date[1]))
            index.Value = decimal.Decimal(data[4])
            index["Open"] = float(data[1])
            index["High"] = float(data[2])
            index["Low"] = float(data[3])
            index["Close"] = float(data[4])
        except ValueError:
            # Do nothing
            return None
#       except KeyError, e:
#          print 'I got a KeyError - reason "%s"' % str(e)
        return index


# NB: CboeVxV class ==  CboeVix class, except for the URL
class CboeVxV(PythonData):
    '''CBOE VXV Download Custom Data Class'''
    
    def GetSource(self, config, date, isLiveMode):
        url_vxv = "http://www.cboe.com/publish/scheduledtask/mktdata/datahouse/vix3mdailyprices.csv"
        return SubscriptionDataSource(url_vxv, 
                                      SubscriptionTransportMedium.RemoteFile)
    def Reader(self, config, line, date, isLiveMode):
        if not (line.strip() and line[0].isdigit()): return None
        index = CboeVxV();
        index.Symbol = config.Symbol
        try:
        # Example File Format:
        #                 OPEN    HIGH    LOW        CLOSE
        # 12/04/2007    24.8    25.01    24.15    24.65
            data = line.split(',')
            date = data[0].split('/')
            index.Time = datetime(int(date[2]), int(date[0]), int(date[1]))
            index.Value = decimal.Decimal(data[4])
            index["Open"] = float(data[1])
            index["High"] = float(data[2])
            index["Low"] = float(data[3])
            index["Close"] = float(data[4])
        except ValueError:
                # Do nothing
                return None
        return index

# for using VIX futures settle in calc. ratios like VIX/VIX1
class QuandlFuture(PythonQuandl):
    '''Custom quandl data type for setting customized value column name. 
       Value column is used for the primary trading calculations and charting.'''
    def __init__(self):
        # Define ValueColumnName: cannot be None, Empty or non-existant column name
        # If ValueColumnName is "Close", do not use PythonQuandl, use Quandl:
        # self.AddData[QuandlFuture](self.VIX1, Resolution.Daily)
        self.ValueColumnName = "Settle"