#region imports
from AlgorithmImports import *
#endregion
class TestIndexOptionAlgorithm(QCAlgorithm):

   def Initialize(self):
      # Backtesting parameters
      self.SetStartDate(2021, 10, 11)
      self.SetEndDate(2021, 10, 16)
      self.SetCash(1000000)
      
      # Index Ticker Symbol
      self.ticker = "SPX"
      # Time Resolution
      self.timeResolution = Resolution.Minute   # Resolution.Minute .Hour .Daily
            
      # Set brokerage model and margin account
      self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
      
      # Days to Expiration
      self.dte = 0
      # Number of strikes to retrieve from the option chain universe (nStrikes on each side of ATM)
      self.nStrikes = 10
            
      # Add the underlying Index
      index = self.AddIndex(self.ticker, self.timeResolution)
      index.SetDataNormalizationMode(DataNormalizationMode.Raw)
      self.underlyingSymbol = index.Symbol
      
      # Keep track of the option contract subscriptions
      self.optionContractsSubscriptions = []
      
      # Set Security Initializer (This does not seem to solve the issue with the benchmark below)
      self.SetSecurityInitializer(self.security_initializer)

      # -----------------------------------------------------------------------------
      # Scheduled function: every day, 25 minutes after the market open
      # -----------------------------------------------------------------------------
      self.Schedule.On(self.DateRules.EveryDay(self.underlyingSymbol)
                       , self.TimeRules.AfterMarketOpen(self.underlyingSymbol, 25)
                       , Action(self.openPosition)
                       )

   def security_initializer(self, security):
       # * SecurityType.Index & SecurityType.IndexOption
        security.SetDataNormalizationMode(DataNormalizationMode.Raw)
        security.SetMarketPrice(self.GetLastKnownPrice(security))
            
   def InitialFilter(self, underlyingsymbol, symbol_list, min_strike_rank, max_strike_rank, min_expiry, max_expiry):
        ''' This method is an initial filter of option contracts
            based on the range of strike price and the expiration date 
            https://www.quantconnect.com/tutorials/applied-options/iron-condor'''
        if len(symbol_list) == 0 : return None
        # fitler the contracts based on the expiry range
        contract_list = [i for i in symbol_list if min_expiry <= (i.ID.Date - self.Time).days <= max_expiry]
        if not contract_list: return None
        # find the strike price of ATM option
        atm_strike = sorted(contract_list,
                            key = lambda x: abs(x.ID.StrikePrice - self.Securities[underlyingsymbol].Price))[0].ID.StrikePrice
        strike_list = sorted(set([i.ID.StrikePrice for i in contract_list]))
        # find the index of ATM strike in the sorted strike list
        atm_strike_rank = strike_list.index(atm_strike)
        try: 
            min_strike = strike_list[atm_strike_rank + min_strike_rank + 1]
            max_strike = strike_list[atm_strike_rank + max_strike_rank - 1]

        except:
            min_strike = strike_list[0]
            max_strike = strike_list[-1]
            
        # skip weekly options
        filtered_contracts = [i for i in contract_list if i.ID.StrikePrice >= min_strike \
                                                        and i.ID.StrikePrice <= max_strike]

        return filtered_contracts
   
   def openPosition(self):
   
      self.Debug("Entering method openPosition")
      
      # Get the option contracts
      contracts = self.OptionChainProvider.GetOptionContractList(self.underlyingSymbol, self.Time)
      contracts = self.InitialFilter(self.underlyingSymbol, contracts, -self.nStrikes, self.nStrikes, 0, self.dte)

      # Exit if we got no chains
      if not contracts:
         self.Debug(" -> No chains!")
         return
   
      # Log the number of contracts that were found
      self.Debug(" -> Found " + str(len(contracts)) + " contracts in the option chain!")
      
      # Do not open any new positions if we have already one open
      if self.Portfolio.Invested:
         return
   
      # Get the furthest expiry date
      expiry = sorted(contracts, key = lambda x: x.ID.Date, reverse = True)[0].ID.Date
      
      # Sort all Put contracts (with the given expiry date) by the strike price in reverse order
      puts = sorted([contract for contract in contracts 
                        if contract.ID.Date == expiry 
                           and contract.ID.OptionRight == OptionRight.Put 
                     ]
                    , key = lambda x: x.ID.StrikePrice
                    , reverse = True
                    )
                    
      # Get the ATM put
      contract = puts[0]
                     
      # Subscribe to the option contract data feed
      if not contract in self.optionContractsSubscriptions:
         self.AddOptionContract(contract, self.timeResolution)
         self.optionContractsSubscriptions.append(contract)
       
      # Sell the Put  
      self.MarketOrder(contract, -1, True)
      # Keep track of the sold contract
      self.put = contract

   def OnOrderEvent(self, orderEvent):
   
      self.Debug(orderEvent)

   
   def OnData(self, slice):
      if self.Portfolio.Invested:
         # Get the security
         security = self.Securities[self.put]
         # Get the price (Looks like self.Securities[x].Price is calculated as the mid-price from self.GetLastKnownPrice(x))
         price = security.Price
         lastKnown = self.GetLastKnownPrice(security)
         # Find out at which time the price was last retreived
         lastPriceDttm = lastKnown.Time
         lastKnownPrice = 0.5*(lastKnown.Bid.Close + lastKnown.Ask.Close)
         # Print the Price every 15 minutes (avoid clogging the log)
         if self.Time.minute % 15 == 0:
            self.Log(f"Contract: {self.put}   Price: {price} ({lastKnownPrice}) @ {lastPriceDttm}")