import numpy as np
from QuantConnect.Python import PythonQuandl
from QuantConnect.Data.Custom import *
import pandas as pd
from datetime import timedelta
from decimal import Decimal

class TermStructureOfVixAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2018, 1, 1)   # Set Start Date
        self.SetEndDate(2018, 8, 1)     # Set End Date
        self.SetCash(100000)          # Set Strategy Cash
        # Add Quandl VIX price (daily)
        self.AddData(QuandlVix, "CBOE/VIX", Resolution.Daily, TimeZones.Chicago)           
        # Add VIX futures contract data 
        self.AddFuture(Futures.Indices.VIX).SetFilter(timedelta(0), timedelta(days=182))
        # Add E-mini S&P500 futures contract data 
        self.AddFuture(Futures.Indices.SP500EMini).SetFilter(self.UniverseFunc)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        self.entryLong = None
        self.latest_VX = None
        self.latest_ES = None
        self.vix = Identity("VIX")
        self.vx_identity = Identity("VX")
        self.es_identity = Identity("ES")
        self.PlotIndicator("VIX", self.vix)
        self.PlotIndicator("VIX", self.vx_identity)
        self.PlotIndicator("ES", self.es_identity)

    def UniverseFunc(self, universe):
        return universe.FrontMonth()
        
    def OnData(self, data):
        if data.ContainsKey("CBOE/VIX"):
            self.vix.Update(self.Time, self.Securities["CBOE/VIX"].Price)
        for chain in data.FutureChains:
            if chain.Key.Value == "VX": 
                front_VX = sorted(chain.Value, key = lambda x: x.Expiry)[0]
                self.latest_VX = front_VX.LastPrice
            if chain.Key.Value == "ES": 
                front_ES = sorted(chain.Value, key = lambda x: x.Expiry)[-1]
                self.latest_ES = front_ES.LastPrice
                if not self.Portfolio.Invested and self.entryLong:
                    self.MarketOrder(front_ES.Symbol, 1)
                    self.entryDate = self.Time
                if self.Portfolio.Invested and self.Time > self.entryDate + timedelta(days=15):
                    self.Liquidate()
    
    def OnEndOfDay(self):
        if self.latest_VX and self.latest_ES:
            self.vx_identity.Update(self.Time, self.latest_VX)
            self.es_identity.Update(self.Time, self.latest_ES)
            self.entryLong = self.vix < self.vx_identity


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