from datetime import timedelta

class BasicTemplateOptionsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 01, 01)
        self.SetEndDate(2016, 03, 01)
        self.SetCash(100000)
        equity = self.AddEquity("IBM", Resolution.Minute)
        option = self.AddOption("IBM", Resolution.Minute)
        self.symbol = option.Symbol
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        # set our strike/expiry filter for this option chain
        option.SetFilter(-3, +3, timedelta(0), timedelta(30))
        # use the underlying equity as the benchmark
        self.SetBenchmark(equity.Symbol)
        self.call = "IBM" # Initialize the call contract
        
    def OnData(self,slice):
        if not self.Portfolio[self.call].Invested and self.MarketOpen():
            self.TradeOptions(slice) # sell the call option
        
        # if the option contract expires, print out the price and position information      
        if slice.Delistings.Count > 0:
            if [x.Key == self.call for x in slice.Delistings]:
                self.Log("stock IBM quantity: {0}".format(self.Portfolio["IBM"].Quantity))
                self.Log("{0} quantity: {1}".format(self.call.Value, self.Portfolio[self.call].Quantity))
                self.Log("The stock price at Expiry S(T): {}".format(self.Securities["IBM"].Price))

    def TradeOptions(self,slice):
        if slice.OptionChains.Count == 0: return   
        for i in slice.OptionChains:
            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: x.Expiry, reverse=True), 
                                            key = lambda x: abs(chain.Underlying.Price - x.Strike))
    
            if len(contracts) == 0: return    
            contract = contracts[0] 
            self.call = contract.Symbol
            self.Sell(self.call, 1) # short the call options
            if self.Portfolio["IBM"].Quantity == 0:
                self.Buy("IBM",100)     # buy 100 the underlying stock
                self.Log("The stock price at time 0 S(0): {}".format(self.Securities["IBM"].Price))  
                
    def MarketOpen(self):
        return self.Time.hour != 0 and self.Time.minute == 1