from datetime import timedelta

class VIXCallProtection(QCAlgorithm):

    def Initialize(self):
        # set start/end date for backtest
        self.SetStartDate(2019, 10, 1)
        self.SetEndDate(2020, 10, 1)
        # set starting balance for backtest
        self.SetCash(100000)
        
        # add asset
        self.equity = self.AddEquity("SPY", Resolution.Minute)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
     
        # add VIX data to get price/close data
        self.data_vix = self.AddData( CBOE, "VIX" )
        self.data_vix.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        # VIX index for option chain
        self.index_vix = self.AddIndex('VIX').Symbol
        #self.index_vix = self.AddEquity('SPY').Symbol
        
        # initialize the option contract with empty string
        self.contract = str()
        self.contractsAdded = set()
        
        self.port_value_options = 0
        
        self.ordered = False
        
        # parameters ------------------------------------------------------------
        self.DaysBeforeExp = 2 # number of days before expiry to exit
        self.DTE = 30 # target days till expiration
        self.OTM = 0.1 # target percentage OTM
        self.lookbackIV = 150 # lookback length of IV indicator
        self.percentage = 0.95 # percentage of portfolio for underlying asset
        self.options_alloc = 0.01 # percentage of options
        # ------------------------------------------------------------------------
    
        # schedule Plotting function 30 minutes after every market open
        self.Schedule.On(self.DateRules.EveryDay(self.equity.Symbol), \
                        self.TimeRules.AfterMarketOpen(self.equity.Symbol, 30), \
                        self.Plotting)

        # warmup for IV indicator of data
        self.SetWarmUp(timedelta(self.lookbackIV)) 
        
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data
        '''
        if(self.IsWarmingUp):
            return
        
        # buy underlying asset
        if not self.Portfolio[self.equity.Symbol].Invested:
            self.SetHoldings(self.equity.Symbol, self.percentage)
        
        #  buy the option
        if self.Securities[self.data_vix.Symbol].Close > 10:
            if self.ordered == False:
                self.FindOption(data)
                self.BuyOption(data)
            
        if self.contract:
            # calculate value of options
            self.port_value_options = self.Portfolio[self.contract.Symbol].Quantity * self.Securities[self.contract.Symbol].Close
            
            # close put before it expires
            #if (self.contract.ID.Date - self.Time) <= timedelta(self.DaysBeforeExp):
            #    self.Liquidate(self.contract.Symbol)
            #    self.RemoveSecurity(self.contract.Symbol)
                
            #    self.Log("Closed: too close to expiration")
            #    self.contract = str()
                
    def BuyOption(self, data):
        if self.contract == str():
            return
        
        # if not invested and option data added successfully, buy option
        if not self.Portfolio[self.contract.Symbol].Invested: #and data.ContainsKey(self.contract.Symbol):
            port_value = self.Portfolio[self.equity.Symbol].Quantity * self.Securities[self.equity.Symbol].Price
            contract_value = self.Securities[self.contract.Symbol].AskPrice
            if contract_value > 0:
                #num_contract = round (port_value * self.options_alloc / contract_value)
                num_contract =100
                if self.contract.BidSize >= 100:
                    self.MarketOrder(self.contract.Symbol, num_contract)
                    self.ordered = True                

        if self.ordered == False:
            self.contract = str()
            
        
    
    def FindOption(self, data):
        # get option data if self.contract empty
        if self.contract == str():
            
            # self.underlyingPrice = self.Securities[self.data_vix.Symbol].Price
            self.underlyingPrice = 20
            min_strike = self.underlyingPrice - 10
            max_strike = self.underlyingPrice + 200
            min_expiry = self.Time + timedelta(days=20)
            max_expiry = self.Time + timedelta(days=60)
            sorted_function = lambda x: x.BidSize
            self.contract = self.GetOption(data, self.index_vix, OptionRight.Call, min_strike, max_strike, min_expiry, max_expiry, sorted_function, True)
            if self.contract == None:
                self.contract = str()
    
            
    def GetOption(self,slice,underlying,optionright,MinStrike,\
                  MaxStrike,MinExpiry,MaxExpiry,sortedfunction,reverse):
                      
        ## Get list of Options Contracts for a specific time
        chainProvider = self.OptionChainProvider 
        contracts = chainProvider.GetOptionContractList(underlying, self.Time)  
        
        if len(contracts) == 0: return
       
        filtered_options = [x for x in contracts if x.ID.OptionRight == optionright and\
                                                     x.ID.StrikePrice >= MinStrike and\
                                                     x.ID.StrikePrice <= MaxStrike and\
                                                     x.ID.Date > MinExpiry and\
                                                     x.ID.Date < MaxExpiry]
           
        if len(filtered_options) == 0: 
            return
           
        added_contracts = []
           
        for x in filtered_options:
            option = self.AddOptionContract(x, Resolution.Minute)
            optionsymbol = self.Securities[option.Symbol]
            added_contracts.append(optionsymbol)
            #added_contracts.append(option)
            
 
        sorted_options=sorted(added_contracts,key=sortedfunction,reverse=reverse)
        selected = sorted_options[0]
        if selected.BidSize == 0:
            selected = None
           
        for x in sorted_options:
            if x != selected:
                self.RemoveSecurity(x.Symbol)
                
        return selected            
   
    

    def Plotting(self):

        # plot underlying's price
        self.Plot("Data Chart Portfolio", self.equity.Symbol, self.Securities[self.equity.Symbol].Close)
        # plot strike of put option
        #self.Plot("Data Chart Options", self.contract, self.port_value_options)
        
        
        #option_invested = [x.Key for x in self.Portfolio if x.Value.Invested and x.Value.Type==SecurityType.Option]
        #if option_invested:
        #    self.Plot("Data Chart", "strike", option_invested[0].ID.StrikePrice)

    def OnOrderEvent(self, orderEvent):
        # log order events
        self.Log(str(orderEvent))