from AlgorithmImports import *

class VirtualYellowGiraffe(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 17)
        self.SetEndDate(2019, 3, 17)
        self.SetCash(50000) 
        self.equity = self.AddEquity('SPY', Resolution.Minute)
    
        self.InitOptionsAndGreeks(self.equity)

        self.delta_call = 0.15
        self.delta_put = 0.30
        self.days_to_exp = 40

        self.has_stock = 0
        self.has_options = 0

        self.close_at_DTE = 0 #    0 to disable
        self.close_at_profit = 0.5 #  takeprofit at residual value of 25% = 75% profit | set to 0 to disable

        self.SetBenchmark ('SPY')

        self.strikeCount  = 100 # no of strikes around underyling price => for universe selection
        self.minExpiryDTE = 30  # min num of days to expiration => for uni selection
        self.maxExpiryDTE = 45  # max num of days to expiration => for uni selection
        

    ## 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 leasr 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):

        #if not self.Portfolio.Invested:
        #    self.has_stock = 0
        #    self.has_options = 0
        
        ## If we're done warming up, and not invested, Sell a put. 
        if (not self.IsWarmingUp): 
            if (self.has_stock == 0) and (self.has_options == 0):
                self.SellAnOTMPut()
                
            elif (self.has_stock == 1) and (self.has_options == 0):
                self.SellAnOTMCall()

            if self.has_options == 1:
                #checking if <21 DTE
                if self.close_at_DTE > 0: #check if enabled
                    if (self.exp_date - self.Time).days < self.close_at_DTE:
                        self.Liquidate(self.ticket.Symbol,"Closing at 21 DTE")
                        self.Debug("Liquidating at 21DTE")
                        self.has_options = 0
                        self.takeProfitTicket.Cancel("Liquidating at 21DTE")
                

      
    ## Sell an OTM Put Option.
    ## Use Delta to select a put contract to sell
    ## ==================================================================
    def SellAnOTMPut(self):
        
        ## Sell a 20 delta put expiring in 2 weeks (14 days)
        putContract = self.SelectContractByDelta(self.equity.Symbol, self.delta_put, self.days_to_exp, OptionRight.Put)
        
        ## construct an order message -- good for debugging and order rrecords

        if self.CurrentSlice.ContainsKey(self.equity.Symbol):
            orderMessage = f"Stock @ ${self.CurrentSlice[self.equity.Symbol].Close} |" + \
                        f"Sell PUT {putContract.Symbol} "+ \
                        f"({round(putContract.Greeks.Delta,2)} Delta)"
                        
            self.Debug(f"{self.Time} {orderMessage}")
            self.ticket = self.MarketOrder(putContract.Symbol, -1, False, orderMessage  )
            self.type = 'PUT'  
            
   
    def SellAnOTMCall(self):
        
        ## Sell a 20 delta call expiring in 2 weeks (14 days)
        callContract = self.SelectContractByDelta(self.equity.Symbol, self.delta_call, self.days_to_exp, OptionRight.Call)
        
        ## construct an order message -- good for debugging and order rrecords

        if self.CurrentSlice.ContainsKey(self.equity.Symbol):
            orderMessage = f"Stock @ ${self.CurrentSlice[self.equity.Symbol].Close} |" + \
                        f"Sell CALL {callContract.Symbol} "+ \
                        f"({round(callContract.Greeks.Delta,2)} Delta)"
                        
            self.Debug(f"{self.Time} {orderMessage}")
            self.ticket = self.MarketOrder(callContract.Symbol, -1, False, orderMessage  ) 
            self.type = 'CALL'  
           
    ## 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)
        if (self.CurrentSlice.OptionChains.values().__len__()>0):
            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  = self.strikeCount # no of strikes around underyling price => for universe selection
        minExpiryDTE = self.minExpiryDTE  # min num of days to expiration => for uni selection
        maxExpiryDTE = self.maxExpiryDTE  # max num of days to expiration => for uni selection
        
        return optionsContractsChain.IncludeWeeklys()\
                                    .Strikes(-strikeCount, strikeCount)\
                                    .Expiration(timedelta(minExpiryDTE), timedelta(maxExpiryDTE))
    
    def OnOrderEvent(self, orderEvent):
        #self.Debug(str(orderEvent))
        if orderEvent.Status == 3: #filled
            if (orderEvent.Symbol.Value == 'SPY'): #we have the stock
                self.has_options = 0
                if (orderEvent.FillQuantity > 0): #buying stock, assigned a PUT
                    self.has_stock = 1
                elif (orderEvent.FillQuantity < 0): #selling stock, assigned a CALL        
                    self.has_stock = 0
            elif (orderEvent.Symbol.Value != 'SPY'): #we have an option
                if orderEvent.IsAssignment or orderEvent.Message[0:3] == 'OTM': #assignment or option expired worthless
                    self.has_options = 0
                    self.type = ''
                    if self.close_at_profit > 0:
                        self.takeProfitTicket.Cancel()
                else:
                    if not(orderEvent.get_LimitPrice() is None): # we have a limit order fill, need to close the initial order
                        #self.Liquidate(self.ticket.Symbol) 
                        self.has_options = 0
                        self.type = ''
                        self.Debug('Take profit hit')
                    else: #it's one of the CALL/PUT market orders
                        self.has_options = 1
                        self.exp_date = orderEvent.Symbol.ID.Date
                        if self.close_at_profit > 0:
                            self.takeProfitTicket = self.LimitOrder(orderEvent.Symbol, 1 , round(orderEvent.FillPrice * self.close_at_profit, 2)) 
           
            


        #if orderEvent.IsAssignment != True and orderEvent.Status == 3 and orderEvent.Symbol.Value != 'SPY':
        #    self.has_options = 1
        
    def OnAssignmentOrderEvent(self, assignmentEvent):
        #if assignmentEvent.Status == 3:
        #    if (assignmentEvent.FillQuantity > 0) and (assignmentEvent.Symbol.Value == 'SPY'): #buying, assigned a PUT
        #        self.has_options = 0
        #        self.has_stock = 1
        #    elif (assignmentEvent.FillQuantity < 0) and (assignmentEvent.Symbol.Value == 'SPY'): #selling, assigned a CALL
        #        self.has_options = 0
        #        self.has_stock = 0
        self.Debug(str(assignmentEvent))