# region imports
from AlgorithmImports import *
# endregion

class SPXExample(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 6, 1)
        self.SetEndDate(2023, 6, 1)
        self.SetCash(100000)
        self.AmountToSpendOnOptions = 5000
        self.EnableZeroDayOptions = True
        self.LogOptions = True

        # SPX weekly options, including 0dte
        self.spx = self.add_index("SPX").Symbol
        self.SPXoption = self.AddIndexOption(self.spx, "SPXW")
        self.SPXoption.SetFilter(self.OptionFilter)

        self.NeedToBuy = True
        self.CallOptionType = 1
        self.PutOptionType = 2

    def OnData(self, data: Slice):
        # Buy once.  Buy a call and a put.
        if self.NeedToBuy:
            self.NeedToBuy = False
            # Buy call by Delta, closest to 0.70
            self.BuyOptionByDelta(self.CallOptionType, self.SPXoption, 0.70)
            # Buy put by Delta, closest to 0.50
            self.BuyOptionByDelta(self.PutOptionType, self.SPXoption, 0.50)
            # Buy call by Price, closest to 2.00
            self.BuyOptionByPrice(self.CallOptionType, self.SPXoption, 2)
            # Buy put by Price, closest to 2.00
            self.BuyOptionByPrice(self.PutOptionType, self.SPXoption, 2)


    def OptionFilter(self, universe):
        # Note that we're limiting expirations to within 0-7 days here.  If you need later this should be adjusted.
        return universe.IncludeWeeklys().Strikes(-10, 10).Expiration(0, 7)

    def BuyOptionByPrice(self, OptionType, underlyingOptions, price):
        self.BuyOption(OptionType, underlyingOptions, price, False)
    def BuyOptionByDelta(self, OptionType, underlyingOptions, delta):
        self.BuyOption(OptionType, underlyingOptions, delta, True)

    def BuyOption(self, OptionType, underlyingOptions, deltaprice, BuyOptionsByDelta = True):
        chain = self.CurrentSlice.OptionChains.get(underlyingOptions.Symbol)
        if chain is None:
            self.Log("Chain error!")
            return

        # First grab expiry dates
        expiry_dates = [contract.Expiry.date() for contract in chain]
        expiry_dates = sorted(set(expiry_dates))
        
        # If you explicitly don't want an option that expires today you can pick another date.
        # Here we just grab the next expiry if 0dte is allowed, otherwise we filter for dates after today.
        next_expiry = None
        if self.EnableZeroDayOptions:
            next_expiry = expiry_dates[0]
        else:
            expiries = [date for date in expiry_dates if date > self.Time.date()]
            if not expiries:
                return
            next_expiry = expiries[0]
        
        
        # Filter contracts for the next expiry date and are call options
        all_options = None
        desired_delta = deltaprice
        if OptionType == self.CallOptionType:
            # Call
            all_options = [contract for contract in chain if contract.Expiry.date() == next_expiry and contract.Right == OptionRight.Call]
        else:
            # Put
            all_options = [contract for contract in chain if contract.Expiry.date() == next_expiry and contract.Right == OptionRight.Put]
            # reverse delta for puts
            if desired_delta > 0:
                desired_delta = -desired_delta
        if not all_options:
            return
        
        # Select the option to BUY with delta closest to the desired delta
        # Result is put into self.selected_option
        if BuyOptionsByDelta:
            self.SelectOptionByDelta(all_options, desired_delta)
        else:
            self.SelectOptionByPrice(all_options, deltaprice)

        
        # Find the last prices for the options.
        option_price = None
        if self.selected_option.Symbol in self.CurrentSlice.Bars:
            option_contract_data = self.CurrentSlice.Bars[self.selected_option.Symbol]
            option_price = option_contract_data.Close

        # Now select a number of contracts based on how much we want to invest
        numcontracts = math.floor(self.AmountToSpendOnOptions/(option_price*100))
        numcontracts = max(1,numcontracts)
        numcontracts = min(300,numcontracts)
        self.MarketOrder(self.selected_option.Symbol, numcontracts)

        # Log option details.
        if self.LogOptions:
            self.LogOption(self.CallOptionType, option_price)


    def SelectOptionByDelta(self, all_options, desired_delta):
        closest_option = None
        closest_delta_diff = float('inf')
        for option in all_options:
            delta_diff = abs(option.Greeks.Delta - desired_delta)
            #self.Log(f"Delta: {option.Greeks.Delta}, diff: {delta_diff}")
            if delta_diff < closest_delta_diff:
                closest_delta_diff = delta_diff
                closest_option = option
                #self.Log("selected")
        self.selected_option = closest_option

    def SelectOptionByPrice(self, all_options, desired_price):
        closest_option = None
        closest_price_diff = float('inf')

        for option in all_options:
            mid_price = (option.BidPrice + option.AskPrice) / 2
            price_diff = abs(mid_price - desired_price)
            if price_diff < closest_price_diff:
                closest_price_diff = price_diff
                closest_option = option
            elif price_diff == closest_price_diff:
                # If price difference is the same, select the lowest strike call or highest strike put
                if (closest_option.Right == OptionRight.Call and option.Strike < closest_option.Strike):
                    closest_option = option
                elif (closest_option.Right == OptionRight.Put and option.Strike > closest_option.Strike):
                    closest_option = option
        self.selected_option = closest_option

    def LogOption(self, optiontype, lastoptionprice):
        # The last option selected should be in self.selected_option so we'll log details of that.
        # If we wanted to find the actual price we bought at, we probably would need to put that in an OnOrderEvent function.
        opt = "CALL"
        if optiontype == self.PutOptionType:
            opt = "PUT"
        strike_price = self.selected_option.Strike
        delta = self.selected_option.Greeks.Delta
        expiry_date = self.selected_option.Expiry
        self.Log(f"{opt} Strike: {strike_price}, Delta: {delta}, Expiry: {expiry_date}, Last Price: {lastoptionprice}")