#region imports
from AlgorithmImports import *
#endregion
from datetime import timedelta

class IronCondorAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2023, 3, 1)
        self.SetEndDate(2023, 3, 31)
        self.SetCash(100000) 
        self.equity = self.AddEquity("SPY", Resolution.Minute)
        
        self.InitOptionsAndGreeks(self.equity)

        # use the underlying equity SPY as the benchmark
        self.SetBenchmark(self.equity.Symbol)

        self.InitOptionsAndGreeks(self.equity)

        self.creditReceived = 0
        self.Leg1Cr = 0
        self.Leg2Db = 0
        self.Leg3Cr = 0
        self.Leg4Db = 0
        
    ## Initialize Options settings, chain filters, pricing models, etc
    ## ====================================================================
    def InitOptionsAndGreeks(self, theEquity):

        ## 1. Specify the data normalization mode (must be 'Raw' for options)
        theEquity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        ## 2. Set Warmup period of at least 30 days
        self.SetWarmup(30, Resolution.Daily)

        ## 3. Set the security initializer to call SetMarketPrice
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))

        ## 4. Subscribe to the option feed for the symbol
        theOptionSubscription = self.AddOption(theEquity.Symbol)

        ## 5. set the pricing model, to calculate Greeks and volatility
        theOptionSubscription.PriceModel = OptionPriceModels.CrankNicolsonFD()  # both European & American, automatically
                
        ## 6. Set the function to filter out strikes and expiry dates from the option chain
        theOptionSubscription.SetFilter(self.OptionsFilterFunction)
        
        
    def OnData(self, data):
        self.OpenIronCondor()

        if self.Portfolio.Invested:
            self.Log("Open P&L: " + str(self.Portfolio.TotalUnrealizedProfit))

        # Check unrealized P&L against opening credit
        if self.Portfolio.TotalUnrealizedProfit < -(3* self.creditReceived * 100):
            self.Liquidate(tag="Stop Loss - 3x credit received")


    def OpenIronCondor(self):
        
        if self.Time.hour == 9 and self.Time.minute == 33 and not self.Portfolio.Invested:
            self.SellAnOTMPut5Delta()
            self.BuyAnOTMPut1Delta()
            self.SellAnOTMCall5Delta()
            self.BuyAnOTMCallDelta1Delta()
            self.creditReceived = (self.Leg1Cr - self.Leg2Db + self.Leg3Cr - self.Leg4Db)
            self.Log("Leg 1 Credit: " + str(self.Leg1Cr))
            self.Log("Leg 2 Debit: " + str(self.Leg2Db))
            self.Log("Leg 3 Credit: " + str(self.Leg3Cr))
            self.Log("Leg 4 Debit: " + str(self.Leg4Db))
            self.Log("Credit collected: " + str(self.creditReceived * 100))
            self.Log("Stop loss: " + str(-(3 * self.creditReceived * 100)))

    
    def SellAnOTMPut5Delta(self):
        
        ## Sell a 5 delta put expiring in today
        putContract = self.SelectContractByDelta(self.equity.Symbol, .05, 0, OptionRight.Put)
        
        ## construct an order message -- good for debugging and order rrecords
        orderMessage = f"Stock @ ${self.CurrentSlice[self.equity.Symbol].Close} | " + \
                       f"Sell {putContract.Symbol} "+ \
                       f"({round(putContract.Greeks.Delta, 2)} Delta)"
                       
        self.Debug(f"{self.Time} {orderMessage}")
        self.Order(putContract.Symbol, -1, False, orderMessage)
        self.Leg1Cr = putContract.AskPrice

    def BuyAnOTMPut1Delta(self):
        
        ## Buy a 1 delta put expiring in today
        putContract = self.SelectContractByDelta(self.equity.Symbol, .01, 0, OptionRight.Put)
        
        ## construct an order message -- good for debugging and order rrecords
        orderMessage = f"Stock @ ${self.CurrentSlice[self.equity.Symbol].Close} | " + \
                       f"Buy {putContract.Symbol} "+ \
                       f"({round(putContract.Greeks.Delta, 2)} Delta)"
                       
        self.Debug(f"{self.Time} {orderMessage}")
        self.Order(putContract.Symbol, 1, False, orderMessage)
        self.Leg2Db = putContract.BidPrice

    def SellAnOTMCall5Delta(self):
        
        ## Sell a 5 delta call expiring in today
        callContract = self.SelectContractByDelta(self.equity.Symbol, .05, 0, OptionRight.Call)
        
        ## construct an order message -- good for debugging and order rrecords
        orderMessage = f"Stock @ ${self.CurrentSlice[self.equity.Symbol].Close} | " + \
                       f"Sell {callContract.Symbol} "+ \
                       f"({round(callContract.Greeks.Delta, 2)} Delta)"
                       
        self.Debug(f"{self.Time} {orderMessage}")
        self.Order(callContract.Symbol, -1, False, orderMessage)  
        self.Leg3Cr = callContract.AskPrice  

    def BuyAnOTMCallDelta1Delta(self):
        
        ## Buy a 1 delta call expiring in today
        callContract = self.SelectContractByDelta(self.equity.Symbol, .01, 0, OptionRight.Call)
        
        ## construct an order message -- good for debugging and order rrecords
        orderMessage = f"Stock @ ${self.CurrentSlice[self.equity.Symbol].Close} | " + \
                       f"Buy {callContract.Symbol} "+ \
                       f"({round(callContract.Greeks.Delta, 2)} Delta)"
                       
        self.Debug(f"{self.Time} {orderMessage}")
        self.Order(callContract.Symbol, 1, False, orderMessage)
        self.Leg4Db = callContract.BidPrice


    ## Get an options contract that matches the specified criteria:
    ## Underlying symbol, delta, days till expiration, Option right (put or call)
    ## ============================================================================
    def SelectContractByDelta(self, symbolArg, strikeDeltaArg, expiryDTE, optionRightArg= OptionRight.Call):

        canonicalSymbol = self.AddOption(symbolArg)
        theOptionChain  = self.CurrentSlice.OptionChains[canonicalSymbol.Symbol]
        theExpiryDate   = self.Time + timedelta(days=expiryDTE)
        
        ## Filter the Call/Put options contracts
        filteredContracts = [x for x in theOptionChain if x.Right == optionRightArg] 

        ## Sort the contracts according to their closeness to our desired expiry
        contractsSortedByExpiration = sorted(filteredContracts, key=lambda p: abs(p.Expiry - theExpiryDate), reverse=False)
        closestExpirationDate = contractsSortedByExpiration[0].Expiry                                        
                                            
        ## Get all contracts for selected expiration
        contractsMatchingExpiryDTE = [contract for contract in contractsSortedByExpiration if contract.Expiry == closestExpirationDate]
    
        ## Get the contract with the contract with the closest delta
        closestContract = min(contractsMatchingExpiryDTE, key=lambda x: abs(abs(x.Greeks.Delta)-strikeDeltaArg))

        return closestContract


    ## The options filter function.
    ## Filter the options chain so we only have relevant strikes & expiration dates. 
    ## =============================================================================
    def OptionsFilterFunction(self, optionsContractsChain):

        strikeCount  = 100 # no of strikes around underyling price => for universe selection
        minExpiryDTE = 0  # min num of days to expiration => for uni selection
        maxExpiryDTE = 0  # max num of days to expiration => for uni selection
        
        return optionsContractsChain.IncludeWeeklys()\
                                    .Strikes(-strikeCount, strikeCount)\
                                    .Expiration(timedelta(minExpiryDTE), timedelta(maxExpiryDTE))

    
    # Liquidate before options are exercised
    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if order.Type == OrderType.OptionExercise:
            self.Liquidate()