class MuscularTanHorse(QCAlgorithm):

    def Initialize(self):
        self.SetEndDate(2021, 3, 3)  # Set Start Date
        self.SetStartDate(2021,3,1)
        equities = ['MSFT']
        for ticker in equities:
            self.AddEquity(ticker, Resolution.Daily)
            option = self.AddOption(ticker, Resolution.Daily)
            option.SetFilter(self.UniverseFunc)
        
        self.SetCash(100000)  # Set Strategy Cash

    def getDTE(self, expiry, currentDate):
        expiryTS = pd.to_datetime(expiry)
        currentDateTS = pd.to_datetime(currentDate)
        diff = (expiryTS - currentDateTS) / np.timedelta64(1,'D')
        return diff

    def addInfoToDF(self, rawDF, priceHistory): 
        #trying to just add the DTE for now
        testSeries = rawDF.apply(lambda row: self.getDTE(row['expiry'], self.Time), axis=1) #this works, returns series
        rawDF['DTE'] = rawDF.apply(lambda row: self.getDTE(row['expiry'], self.Time), axis=1) #this does not work, KeyError on 'DTE'
        newDF = rawDF.assign(d=testSeries.values) #does not work, KeyError on 'd'
  
        

    def UniverseFunc(self, universe):
        return universe.IncludeWeeklys().Expiration(TimeSpan.FromDays(0),
                                               TimeSpan.FromDays(20)).Strikes(-3,3)
    def OnData(self, data):
        self.optionData = data.OptionChains
      
        for stockSymbol in data.OptionChains:
            optionchain = stockSymbol.Value #extract this stock's option chain
            underlyingTicker = optionchain.Underlying.Symbol.ToString() #extract the stock's ticker
            underlyingPriceHist = self.History(self.Symbol(underlyingTicker),365) #get recent price history for further calculation
            
            rawDF = pd.DataFrame([[x.Right,float(x.Strike),x.Expiry,float(x.BidPrice),float(x.AskPrice)] for x in optionchain],index=[x.Symbol.Value for x in optionchain], columns=['type', 'strike', 'expiry', 'ask price', 'bid price']) #put options data into DataFrame
            
            dfWithData = self.addInfoToDF(rawDF, underlyingPriceHist) #pass df to helper function to add data
        
        return