from AlgorithmImports import *

# Goal:
#   - Option chain selection

class TSLAOptionChainFilter(QCAlgorithm): 
    def Initialize(self):
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2023, 1, 10)
        self.SetCash(5000)

        # Time Resolution controller 
        self.timeResolution = Resolution.Hour

        # Set TSLA equity 
        equity = self.AddEquity("TSLA", self.timeResolution)
        self.underlyingSymbol = equity.Symbol

        # Set the RAW price 
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw)

        # Set TSLA options in the Universe 
        option = self.AddOption("TSLA", self.timeResolution)
        self.symbol = option.Symbol

        # Set the model for better pricing model
        option.PriceModel = OptionPriceModels.BjerksundStensland()

        # Set broker 
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Cash)

        ## Set the security initializer to call SetMarketPrice
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))
        
        # Point of interest contracts 
        self.contractStrikes = []


    def OnData(self, slice: Slice) -> None:
        if not self.Portfolio.Invested:
            atm_strike = int(slice[self.underlyingSymbol].Price)
            min_strike = atm_strike - 22
            max_strike = atm_strike + 22
            min_expiry = self.Time + timedelta(days=20)
            max_expiry = self.Time + timedelta(days=25)

            options = self.GetOption(slice, self.underlyingSymbol, OptionRight.Call, min_expiry, max_expiry, min_strike, max_strike)
            
            expiry = min([symbol.ID.Date for symbol in options])
            filtered_options = [symbol for symbol in options if symbol.ID.Date == expiry]

            strikes = [x.ID.StrikePrice for x in filtered_options]
            self.Log( "Price: " + str(atm_strike) + " | " + " Expiry: " + str(expiry.date()) + " | " + str(strikes))

    def GetOption(self, slice, underlying, option_right, min_expiry, max_expiry, min_strike, max_strike):

        ## Get list of Options Contracts for a specific time
        chainProvider = self.OptionChainProvider 
        contracts = chainProvider.GetOptionContractList(underlying, self.Time)  
        
        if len(contracts) == 0: return
       
        filtered_options = [x for x in contracts if x.ID.OptionRight == option_right and \
                                                     x.ID.StrikePrice >= min_strike and \
                                                     x.ID.StrikePrice <= max_strike and \
                                                     x.ID.Date >= min_expiry and \
                                                     x.ID.Date <= max_expiry]
           
        if len(filtered_options) == 0: 
            return
        
        sorted_options = sorted(filtered_options, key = lambda x: x.ID.StrikePrice)

        return sorted_options        


        # filtered = [x for x in chain if x.ID.Date.date() >= (self.Time.date() + timedelta(days = min_expiration)) ]

        # # Pick the closest expiration day
        # expiry = min([symbol.ID.Date for symbol in filtered])

        # # Get current equity price and round it 
        # atm_strike = int(self.Securities[self.underlyingSymbol].Price)

        # # Filter based on 
        # filtered_symbols = [symbol for symbol in filtered if symbol.ID.Date == expiry and 
        #                     symbol.ID.StrikePrice >= atm_strike + min_strike and 
        #                     symbol.ID.StrikePrice <= atm_strike + max_strike]
        
        # contract_symbol = sorted(filtered_symbols, key=lambda symbol: symbol.ID.StrikePrice)

        # option_strikes = sorted([x.ID.StrikePrice for x in contract_symbol])

        # if self.contractStrikes != option_strikes:
        #     self.contractStrikes = option_strikes
        #     self.Log("price: " + str(atm_strike) + " | " + str(option_strikes))

    def OnOrderEvent(self, orderEvent):
        if not orderEvent.Status == OrderStatus.Filled:
            return 
        message = f"Symbol {orderEvent.Symbol} |" + \
                    f"Fill Price {orderEvent.FillPrice} "+ \
                    f"Action {str(orderEvent.Direction)}"
        self.Log(f"{self.Time} {message}")