class LongStrangleAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 7, 1)
        # self.SetEndDate(2017, 5, 30)
        self.SetCash(10000)
        equity = self.AddEquity("SPY", Resolution.Minute)
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.underlyingsymbol = equity.Symbol
        self.SetBenchmark(equity.Symbol)
        
        tlt = self.AddEquity("TLT", Resolution.Hour)
        self.tlt_sma_long = self.SMA("TLT", 200, Resolution.Hour)
        self.tlt_sma_short = self.SMA("TLT", 50, Resolution.Hour)
        
        self.Schedule.On(
            self.DateRules.Every(DayOfWeek.Thursday), 
            self.TimeRules.BeforeMarketClose("SPY", 60), 
            self.ClosePositions
        ) 
        
        self.Schedule.On(
            self.DateRules.Every(DayOfWeek.Monday), 
            self.TimeRules.AfterMarketOpen("SPY", 120), 
            self.OpenPosition
        ) 
        
        # self.SetWarmUp(200)

    def OpenPosition(self):

        ''' OptionChainProvider gets the option chain provider,
            used to get the list of option contracts for an underlying symbol.
            Then you can manually filter the contract list returned by GetOptionContractList.
            The manual filtering will be limited to the information
            included in the Symbol (strike, expiration, type, style) and/or prices from a History call '''

        if not self.Portfolio.Invested:
            filtered_contracts = self.InitialFilter(-10, 10, 4, 7)
            if len(filtered_contracts) == 0: return
                
            if self.tlt_sma_long < self.tlt_sma_short:
                # sorted the optionchain by expiration date and choose the furthest date
                expiry = sorted(filtered_contracts,key = lambda x: x.ID.Date, reverse=False)[0].ID.Date
                # filter the call options from the contracts expires on that date
                call = [i for i in filtered_contracts if i.ID.Date == expiry and i.ID.OptionRight == OptionRight.Call]
                # sorted the contracts according to their strike prices 
                call_contracts = sorted(call,key = lambda x: x.ID.Date)    
                # choose the deep OTM call option
                self.call = call_contracts[-1]
                        
                self.AddOptionContract(self.call, Resolution.Minute)

                self.MarketOrder(self.call, 1)
            else:
                # sorted the optionchain by expiration date and choose the furthest date
                expiry = sorted(filtered_contracts,key = lambda x: x.ID.Date, reverse=False)[0].ID.Date
                # select the put options which have the same expiration date with the call option 
                # sort the put options by strike price
                put_contracts = sorted([i for i in filtered_contracts if i.ID.Date == expiry and i.ID.OptionRight == OptionRight.Put], key = lambda x: x.ID.Date)
                # choose the deep OTM put option
                self.put = put_contracts[0]
                        
                self.AddOptionContract(self.put, Resolution.Minute)
                
                self.MarketOrder(self.put, 1)         


    def InitialFilter(self, min_strike_rank, max_strike_rank, min_expiry, max_expiry):
        
        ''' This method is an initial filter of option contracts 
            according to the range of strike price and the expiration date '''
            
        contracts = self.OptionChainProvider.GetOptionContractList(self.underlyingsymbol, self.Time.date())
        if len(contracts) == 0 : return []
        # fitler the contracts based on the expiry range
        contract_list = [i for i in contracts if min_expiry < (i.ID.Date.date() - self.Time.date()).days < max_expiry]
        if len(contract_list) == 0: return []
        
        self.Log(contract_list)
        
        # find the strike price of ATM option
        atm_strike = sorted(contract_list,
                            key = lambda x: abs(x.ID.StrikePrice - self.Securities[self.underlyingsymbol].Price))[0].ID.StrikePrice
        strike_list = sorted(set([i.ID.StrikePrice for i in contract_list]))
        # find the index of ATM strike in the sorted strike list
        atm_strike_rank = strike_list.index(atm_strike)
        try: 
            strikes = strike_list[(atm_strike_rank + min_strike_rank):(atm_strike_rank + max_strike_rank)]
        except:
            strikes = strike_list
        filtered_contracts = [i for i in contract_list if i.ID.StrikePrice in strikes]
        
        self.Log(filtered_contracts)

        return filtered_contracts 

    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))
        
    def ClosePositions(self):
        if self.Portfolio.Invested:
            self.Liquidate()