from datetime import timedelta

class longButterflyExample(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 4, 1)
        self.SetEndDate(2017, 5, 30)
        self.SetCash(100000)
        equity = self.AddEquity("SPY", Resolution.Minute)
        option = self.AddOption("SPY", Resolution.Minute)
        self.symbol = option.Symbol
        # set our strike/expiry filter for this option chain
        option.SetFilter(-15, 15, timedelta(30), timedelta(60))
        
    def OnData(self,slice):
        optionchain = slice.OptionChains
        for i in slice.OptionChains:
            if i.Key != self.symbol: continue
            chains = i.Value
            contract_list = [x for x in chains]
            # if there is no contracts in this optionchain, pass the instance
            if (slice.OptionChains.Count == 0) or (len(contract_list) == 0): return   
             # if there is no securities in portfolio, trade the options 
            if not self.Portfolio.Invested: 
                self.TradeOptions(optionchain)

    def TradeOptions(self,optionchain):
    
        for i in optionchain:
            if i.Key != self.symbol: continue
            chain = i.Value
            # sorted the optionchain by expiration date and choose the furthest date
            expiry = sorted(chain,key = lambda x: x.Expiry, reverse=True)[0].Expiry
            # filter the call options from the contracts expires on that date
            call = [i for i in chain if i.Expiry == expiry and i.Right == 0]
            # sorted the contracts according to their strike prices 
            call_contracts = sorted(call,key = lambda x: x.Strike)    
            if len(call_contracts) == 0: continue
            # choose the deep OTM call option
            self.call = call_contracts[-1]
            # 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 chain if i.Expiry == expiry and i.Right == 1], key = lambda x: x.Strike)
            # choose the deep OTM put option
            self.put = put_contracts[0]
            
            # Sell OTM Call and OTM Put
            self.Sell(self.call.Symbol ,1)
            self.Sell(self.put.Symbol ,1)
            
            # Locate ATM contracts
            # Put
            put = [i for i in chain if i.Right == OptionRight.Put]
            putContracts =  sorted(sorted(put, key=lambda x: x.Expiry, reverse=False), 
                                            key=lambda x: abs(x.UnderlyingLastPrice - x.Strike))
            self.atmPut = putContracts[0]
            
            # Call
            call = [i for i in chain if i.Right == OptionRight.Call]
            callContracts =  sorted(sorted(call, key=lambda x: x.Expiry, reverse=False), 
                                            key=lambda x: abs(x.UnderlyingLastPrice - x.Strike))
            self.atmCall = callContracts[0]
            
            # Buy ATM Call contract and ATM Put contract
            self.Buy(self.atmCall.Symbol ,1)
            self.Buy(self.atmPut.Symbol ,1)
            
    
    def OnOrderEvent(self, orderEvent):
        self.Log(str(orderEvent))