from AlgorithmImports import *
from datetime import datetime, timedelta

class MyAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 1, 5)  # Set your desired start date
        self.SetEndDate(2024, 1, 10)    # Set your desired end date
        self.SetCash(100000)            # Set your initial capital

        # Add SPX and SPY indices
        spx = self.AddIndex("SPX", Resolution.Minute).Symbol
        spy = self.AddEquity("SPY", Resolution.Minute).Symbol
        self.optionSymbols = []
        self.optionSymbols.append(Symbol.CreateCanonicalOption(spx, "SPXW", Market.USA, "?SPXW"))
        self.optionSymbols.append(self.AddOption(spy).Symbol)

        # Schedule the SelectOptionStrikes method to run every 5 minutes
        self.Schedule.On(self.DateRules.EveryDay('SPY'), self.TimeRules.Every(timedelta(minutes=60)), self.SelectOptionStrikes)

    def OnSecuritiesChanged(self, changes):
        self.Plot('SPX', 'Added', len(changes.AddedSecurities))
        self.Plot('SPX', 'Removed', len(changes.RemovedSecurities))
        self.Plot("Total Securities", len(self.ActiveSecurities.Keys))

    def SelectOptionStrikes(self):

        if not self.IsMarketOpen(self.optionSymbols[0]):
            for optionSymbol in self.optionSymbols:
                for activeSymbol in self.ActiveSecurities.Keys:
                    if activeSymbol.SecurityType in [SecurityType.Option, SecurityType.IndexOption] and activeSymbol.Underlying == optionSymbol.Underlying and not self.Portfolio[activeSymbol].Invested:
                        self.RemoveOptionContract(activeSymbol)

        if self.CurrentSlice.ContainsKey(self.optionSymbols[0].Underlying):
            # Define the range around the spot price
            strikeRange = 20

            # Loop through SPX and SPY symbols
            for optionSymbol in self.optionSymbols:
                # Get the spot price of the specified index
                spotPrice = self.Securities[optionSymbol.Underlying].Price
                self.Log(f'{optionSymbol.Underlying} spotPrice {spotPrice}')
                
                # Get the option chain for the specified index
                optionChain = self.OptionChainProvider.GetOptionContractList(optionSymbol, self.Time)
                optionChain = [x for x in optionChain if x.ID.Date.date() == self.Time.date()] if optionSymbol.Underlying == 'SPX' else [x for x in optionChain if x.ID.Date.date() == (self.Time.date() + timedelta(days=1))]
                self.Log(f'{optionSymbol.Underlying} length of optionChain {len(optionChain)}')

                # Select 20 strikes around the spot price with expiration date equal to the current date
                selectedStrikes = sorted([x for x in optionChain],key=lambda x: abs(x.ID.StrikePrice - spotPrice))[:40]
                selectedStrikes = sorted(selectedStrikes, key=lambda x: x.ID.StrikePrice)
                if len(selectedStrikes) > 0:
                    self.Log(f'{optionSymbol.Underlying} length of selectedStrikes {len(selectedStrikes)} expiry: {selectedStrikes[0].ID.Date.date()} low {selectedStrikes[0].ID.StrikePrice} high {selectedStrikes[-1].ID.StrikePrice}')

                    # Remove strikes that are no longer within the selected range and not invested
                    for activeSymbol in self.ActiveSecurities.Keys:
                        if activeSymbol.Underlying == optionSymbol.Underlying and activeSymbol not in selectedStrikes and not self.Portfolio[activeSymbol].Invested:
                            self.RemoveOptionContract(activeSymbol)

                    # Add new strikes to the universe
                    for strike in selectedStrikes:
                        if strike not in self.ActiveSecurities.Keys:
                            self.AddOptionContract(strike, Resolution.Minute)

            # Log the total number of securities
            self.Debug(f"Total Securities: {len(self.ActiveSecurities.Keys)}")