from QuantConnect.Data.Custom.CBOE import *
from datetime import timedelta

class VixAlgo(QCAlgorithm):

    target_pct_change_time = datetime.min
    
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2020, 5, 15) 
        self.SetCash(10000)
        
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol
        self.cboeVix = self.AddData(CBOE, "VIX").Symbol
        self.vixPrevious = None
        self.vixLatest = None
        self.pct_change = None
        self.SetWarmup(5, Resolution.Daily)
        self.stop_price = None
        
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen(self.spy, 10), Action(self.Buy))
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.AfterMarketOpen(self.spy, 60), Action(self.Sell))
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.Every(timedelta(minutes=1)), Action(self.StopLoss))
        
#        ''' for charting benchmark over equity curve'''
#        self.benchmarkTicker = 'SPY'
#        self.SetBenchmark(self.benchmarkTicker)
#        self.initBenchmarkPrice = None
#        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.BeforeMarketClose(self.spy, 1), Action(self.PlotBenchmark))

    def OnData(self, data):
        if data.ContainsKey(self.cboeVix):
            self.vixPrevious = self.vixLatest
            self.vixLatest = data.Get(CBOE, self.cboeVix).Close
            
        if self.IsWarmingUp:
            return 
        
        if not data.Bars.ContainsKey("SPY"):
            return
        
        if self.vixPrevious != 0:
            self.pct_change = (self.vixLatest - self.vixPrevious) / self.vixPrevious
            #self.Log("VIX percent change is {}".format(self.pct_change))
            
#        self.UpdateBenchmarkValue()
            
    def Buy(self):
        # keep buying for 3 days after the pct change target occurs
        if (self.Time - self.target_pct_change_time).days < 4:
            print('Test check')
            self.SetHoldings("SPY", 1)    
            self.stop_price = self.Securities["SPY"].Price
            return
        
        if self.pct_change is not None and self.pct_change > 0.05:
            self.SetHoldings("SPY", 1)
            self.stop_price = self.Securities["SPY"].Price
            self.target_pct_change_time = self.Time
            self.Log("VIX percent change of {} is greater than 5%. Trading today".format(self.pct_change))
        
    def Sell(self):
        self.SetHoldings("SPY", 0)
        
    def StopLoss(self):
        if self.stop_price is not None:
            if self.Securities["SPY"].Price < self.stop_price * 0.97:
                self.SetHoldings("SPY", 0)
    
#    def UpdateBenchmarkValue(self):
#        ''' Simulate buy and hold the Benchmark '''
#        if self.initBenchmarkPrice is None:
#            self.initBenchmarkCash = self.Portfolio.Cash
#            self.initBenchmarkPrice = self.Benchmark.Evaluate(self.Time)
#            self.benchmarkValue = self.initBenchmarkCash
#        else:
#            currentBenchmarkPrice = self.Benchmark.Evaluate(self.Time)
#            self.benchmarkValue = (currentBenchmarkPrice / self.initBenchmarkPrice) * self.initBenchmarkCash
            
#    def PlotBenchmark(self):
#        self.Plot('Strategy Equity', self.benchmarkTicker, self.benchmarkValue)