from datetime import timedelta

class SmoothBlueLion(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 9, 1)  # Set Start Date
        self.SetEndDate(2021, 9, 1)
        self.SetCash(100000)  # Set Strategy Cash
        # self.AddEquity("SPY", Resolution.Minute)
        self.AddUniverse(self.MyCoarseFilterFunction)
        self.UniverseSettings.Resolution = Resolution.Minute
        self.DTE = 30
        self.ExpDelta = 20 # num days to look before/after desired DTE (i.e if only monthly exp, might need to go 50 days out)
        self.contractsAdded = set()
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Raw))
        self.contractsWithoutPrices = 0
        self.contractsWithPrices = 0

    
    def OnData(self, data):
        if not (data.Time.hour == 10 and data.Time.minute == 0):
            return
        
        
        for security in self.ActiveSecurities.Values:
            if security.Symbol.SecurityType == SecurityType.Equity:
                contracts = self.OptionsFilter(security.Symbol)
                
                if contracts == str(): continue
            
                for contract in contracts:
                    if contract.ID.OptionRight == OptionRight.Put:
                        bid = self.Securities[contract].BidPrice
                        self.Debug("bid: " + str(bid))
                    else:
                        ask = self.Securities[contract].AskPrice
                        self.Debug("ask: " + str(ask))
                
                if bid > 0 and ask > 0:
                    self.contractsWithPrices += 1
                else:
                    self.contractsWithoutPrices += 1
                    #self.Debug(str(security.Symbol.Value))
                    #self.Debug("returned contracts: " + str([x.Value for x in contracts]))
        
        self.Debug("Number of Contracts with Prices: " + str(self.contractsWithPrices))
        self.Debug("Number of Contracts without Prices: " + str(self.contractsWithoutPrices))
        self.contractsWithPrices = 0
        self.contractsWithoutPrices = 0
        
    def MyCoarseFilterFunction(self, coarse):
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        filtered = [ x.Symbol for x in sortedByDollarVolume 
                      if x.Price > 5 and x.Price < 30 ][:5]
        return filtered
    
    def OnSecuritiesChanged(self, changes):
        
        for security in changes.AddedSecurities:
            if security.Symbol.SecurityType == SecurityType.Equity:
                self.OptionsFilter(security.Symbol) # call OptionsFilter to subscribe to the options contracts data
                self.Debug(str(security.Symbol) + " is an equity")
        
    def OptionsFilter(self, symbol):

        contracts = self.OptionChainProvider.GetOptionContractList(symbol, self.Time)
        underlyingPrice = self.Securities[symbol].Price
        
        puts = [i for i in contracts if i.ID.OptionRight == OptionRight.Put and
                                            underlyingPrice < i.ID.StrikePrice and
                                            self.DTE - self.ExpDelta < (i.ID.Date - self.Time).days < self.DTE + self.ExpDelta]
        
        calls = [i for i in contracts if i.ID.OptionRight == OptionRight.Call and
                                            underlyingPrice < i.ID.StrikePrice and
                                            self.DTE - self.ExpDelta < (i.ID.Date - self.Time).days < self.DTE + self.ExpDelta]
        
        if len(puts) > 0 and len(calls) > 0:
            
            sortedPuts = sorted(sorted(puts, key = lambda x: abs((x.ID.Date - self.Time).days - self.DTE)),
                                                     key = lambda x: underlyingPrice - x.ID.StrikePrice, reverse=True)[0]
           
            sortedCalls = sorted(sorted(calls, key = lambda x: abs((x.ID.Date - self.Time).days - self.DTE)),
                                                     key = lambda x: underlyingPrice - x.ID.StrikePrice, reverse=True)[0]
            
            atmContracts = [sortedPuts, sortedCalls]
            for contract in atmContracts:
                if contract not in self.contractsAdded:
                    self.contractsAdded.add(contract)
                    self.AddOptionContract(contract, Resolution.Minute)
            return atmContracts
        else:
            return str()