class LongStraddleAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 1, 1)
        self.SetEndDate(2019, 1, 1)
        self.SetCash(100000)
        self.equity = self.AddEquity("SPY", Resolution.Minute)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.underlyingsymbol = self.equity.Symbol
        # use the underlying equity GOOG as the benchmark
        self.SetBenchmark(self.underlyingsymbol)
        self.Schedule.On(self.DateRules.EveryDay(),self.TimeRules.At(15,30),self.close_pos)
        self.call_cost = 0
        self.put_cost = 0

    def OnData(self,slice):

        if not self.Portfolio.Invested:
            contracts = self.OptionChainProvider.GetOptionContractList(self.underlyingsymbol, self.Time.date())
            if self.Time.hour == 9 and self.Time.minute == 32:
                self.TradeOptions(contracts)
                self.call_cost = self.Securities[self.call.Value].Price
                self.put_cost = self.Securities[self.put.Value].Price
                
        if self.Portfolio.Invested:        
        
            #if price of underlying goes beyond the range of (- 2%) it will liquidate
            if self.Securities[self.call.Value].Price < 0.98 * self.call_cost:
                self.Liquidate(self.call.Value)
                self.Debug("call liquidated")
            elif self.Securities[self.put.Value].Price < 0.98 * self.put_cost:
                self.Liquidate(self.put.Value)
                self.Debug("put liquidated")

    def TradeOptions(self,contracts):
        # run CoarseSelection method and get a list of contracts expire within 0 to 8 days from now on
        # and the strike price between rank -1 to rank 1
        filtered_contracts = self.CoarseSelection(self.underlyingsymbol, contracts, -1, 1, 0, 8)
        if filtered_contracts is None: return
        expiry = sorted(filtered_contracts,key = lambda x: x.ID.Date, reverse=True)[0].ID.Date
        # filter the call options from the contracts expire on that date
        call = [i for i in filtered_contracts if i.ID.Date == expiry and i.ID.OptionRight == 0]
        # sorted the contracts according to their ATM strike prices
        call_contracts = sorted(call , key = lambda x: abs(x.ID.StrikePrice-self.Securities[self.underlyingsymbol].Price))
        if len(call_contracts) <= 0: return
        self.call = call_contracts[0]
        
        
        for i in filtered_contracts:
            if i.ID.Date == expiry and i.ID.OptionRight == 1 and i.ID.StrikePrice ==call_contracts[0].ID.StrikePrice:
                self.put = i

        self.AddOptionContract(self.call, Resolution.Minute)
        self.AddOptionContract(self.put, Resolution.Minute)

        self.Sell(self.call.Value ,1)
        self.Sell(self.put.Value ,1)

    def CoarseSelection(self, underlyingsymbol, symbol_list, min_strike_rank, max_strike_rank, min_expiry, max_expiry):


        # fitler the contracts based on the expiry range
        contract_list = [i for i in symbol_list if min_expiry < (i.ID.Date.date() - self.Time.date()).days < max_expiry]
        # find the strike price of ATM option
        if len(contract_list) <= 0: return
        atm_strike = sorted(contract_list,
                            key = lambda x: abs(x.ID.StrikePrice - self.Securities[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: 
            min_strike = strike_list[atm_strike_rank + min_strike_rank]
            max_strike = strike_list[atm_strike_rank + max_strike_rank]
        except:
            min_strike = strike_list[0]
            max_strike = strike_list[-1]
        # filter the contracts based on the range of the strike price rank
        filtered_contracts = [i for i in contract_list if i.ID.StrikePrice >= min_strike and i.ID.StrikePrice <= max_strike]

        return filtered_contracts
        
    def close_pos(self):
        if self.Portfolio.Invested:
            self.Liquidate()