from datetime import timedelta

class CoveredCallOptionsAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2011, 01, 01)
        self.SetEndDate(2017, 02, 01)
        self.SetCash(100000)
        self.AddEquity("SPY", Resolution.Daily)
        equity = self.AddEquity("QQQ", Resolution.Minute)
        self.underlyingsymbol = equity.Symbol
        # use the underlying equity as the benchmark
        self.SetBenchmark(equity.Symbol)
        self.call = "QQQ" # Initialize the call contract
        self.rebalance = False;
                
        # schedule an event to fire at the beginning of the month, the symbol is optional
        # if specified, it will fire the first trading day for that symbol of the month,
        # if not specified it will fire on the first day of the month
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY"), Action(self.Rebalance))
        
        
    def OnData(self,slice):
        if not self.rebalance:
            return
        
        self.rebalance = False
            
        self.Log("rebalance=true")
        contracts = self.OptionChainProvider.GetOptionContractList(self.underlyingsymbol, self.Time.date())
        if len(contracts) == 0: return
        filtered_contracts = self.InitialFilter(self.underlyingsymbol, contracts, -3, 3, 0, 30)
        call = [x for x in filtered_contracts if x.ID.OptionRight == 0] 
        # sorted the contracts according to their expiration dates and choose the ATM options
        contracts = sorted(sorted(call, key = lambda x: abs(self.Securities["QQQ"].Price- x.ID.StrikePrice)), 
                                        key = lambda x: x.ID.Date, reverse=True)
        self.call = contracts[0]
        #numberOfShares = self.Portfolio.Cash / self.Securities["QQQ"].Price
        #numberOfOptions = ( 100 % numberOfShares)
        #numberOfShares = numberOfOptions * 100
        self.AddOptionContract(self.call, Resolution.Minute)
        self.Sell(self.call, 6) # short the call options
        if self.Portfolio["QQQ"].Quantity == 0:
          #  self.Log("numberOfShares: " + str(numberOfShares) +" : "+ str(numberOfOptions))
            self.Buy("QQQ",600)     # buy 100 the underlying stock
           # self.Log("The stock price at time 0 S(0): {0}".format(self.Securities["QQQ"].Price))
            
        #if not self.Portfolio[self.call].Invested and self.Time.hour != 0 and self.Time.minute == 1:
       #     self.Log("OnData(self,slice)")

        # 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 QQQ quantity: {0}".format(self.Portfolio["QQQ"].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["QQQ"].Price))

    def Rebalance(self):
        self.rebalance = True;
    
    def InitialFilter(self, underlyingsymbol, symbol_list, 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 '''
            
        if len(symbol_list) == 0 : return
        # 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
        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]
           
        filtered_contracts = [i for i in contract_list if i.ID.StrikePrice >= min_strike and i.ID.StrikePrice <= max_strike]

        return filtered_contracts