from datetime import timedelta

class BasicTemplateOptionsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 01, 01)
        self.SetEndDate(2016, 06, 30)
        self.SetCash(40000)
        equity = self.AddEquity("IBM", Resolution.Minute)
        option = self.AddOption("IBM", Resolution.Minute)
        self.symbol = option.Symbol

        # set our strike/expiry filter for this option chain
        option.SetFilter(-3, +3, timedelta(0), timedelta(60))
        # use the underlying equity as the benchmark
        self.SetBenchmark(equity.Symbol)
        
    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 (slice.OptionChains.Count < 1) or (len(contract_list) < 1): 
            return   
 
        if not self.Portfolio.HoldStock: 
            self.Buy("IBM",100)     # buy 100 the underlying stock
            self.SellCall(optionchain) # sell the call option
        # if the option contract expires, choose a new call option to sell       
        if slice.Delistings.Count > 0 and self.Portfolio.HoldStock:
            for x in slice.Delistings:
                if x.Key == self.contract_symbol:
                    self.SellCall(optionchain)

    def SellCall(self,optionchain):

        for i in optionchain:
            if i.Key != self.symbol: continue
            chain = i.Value
            call = [x for x in chain if x.Right == 0] # filter the call options contracts
            # sorted the contracts according to their expiration dates and choose the ATM options
            contracts = sorted(sorted(call, key = lambda x: abs(chain.Underlying.Price - x.Strike)), 
                                            key = lambda x: x.Expiry, reverse=True)
    
            if len(contracts) == 0: continue
            contract = contracts[0]
            self.Log("Price: {0}; ATM Strike: {1}".format(self.Securities["IBM"].Price, contract.Strike))
            self.contract_symbol = contract.Symbol
            self.Sell(self.contract_symbol, 1) # short the call options