#region imports
from AlgorithmImports import *
#endregion
# ######################################################
# ## Code Seperated for readiability
# ######################################################

class OptionsUtil():

    def __init__(self, algo, theEquity):
        self.algo = algo
        self.InitOptionsAndGreeks(theEquity)
        
    ## 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.algo.SetWarmup(30, Resolution.Daily)

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

        ## 4. Subscribe to the option feed for the symbol
        theOptionSubscription = self.algo.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)

    ## Buy an OTM Call Option.
    ## Use Delta to select a call contract to buy
    ## ---------------------------------------------------------------------------
    def BuyAnOTMCall(self, theSymbol):
        
        ## Buy a Call expiring 
        callDelta  = float(self.algo.GetParameter("callDelta"))/100
        callDTE    = int(self.algo.GetParameter("callDTE")) 

        callContract = self.SelectContractByDelta(theSymbol, callDelta, callDTE, OptionRight.Call)
        
        # construct an order message -- good for debugging and order rrecords
        # ------------------------------------------------------------------------------    
        #  if( callContract is not None ): # Might need this.... 
        orderMessage = f"Stock @ ${self.algo.CurrentSlice[theSymbol].Close} |" + \
                       f"Buy {callContract.Symbol} "+ \
                       f"({round(callContract.Greeks.Delta,2)} Delta)"
                       
        self.algo.Debug(f"{self.algo.Time} {orderMessage}")
        self.algo.Order(callContract.Symbol, 1, False, orderMessage  )   
           
           
           
    ## Sell an OTM Put Option.
    ## Use Delta to select a put contract to sell
    ## ---------------------------------------------------------------------------
    def SellAnOTMPut(self, theSymbol):
        
        ## Sell a Put expiring in 2 weeks (14 days)
        putDelta  = float(self.algo.GetParameter("putDelta"))/100
        putDTE    = int(self.algo.GetParameter("putDTE")) 

        putContract = self.SelectContractByDelta(theSymbol, putDelta, putDTE, OptionRight.Put)
        
        ## construct an order message -- good for debugging and order rrecords
        orderMessage = f"Stock @ ${self.algo.CurrentSlice[theSymbol].Close} |" + \
                       f"Sell {putContract.Symbol} "+ \
                       f"({round(putContract.Greeks.Delta,2)} Delta)"
                       
        self.algo.Debug(f"{self.algo.Time} {orderMessage}")
        self.algo.Order(putContract.Symbol, -1, False, orderMessage  )   
           
   
    ## 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.algo.AddOption(symbolArg)
        if(canonicalSymbol.Symbol not in self.algo.CurrentSlice.OptionChains):
            self.algo.Log(f"{self.algo.Time} [Error] Option Chain not found for {canonicalSymbol.Symbol}")
            return
        
        theOptionChain  = self.algo.CurrentSlice.OptionChains[canonicalSymbol.Symbol]
        theExpiryDate   = self.algo.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  = 200 # no of strikes around underyling price => for universe selection
        minExpiryDTE = 55  # min num of days to expiration => for uni selection
        maxExpiryDTE = 65  # max num of days to expiration => for uni selection
        
        return optionsContractsChain.IncludeWeeklys()\
                                    .Strikes(-strikeCount, strikeCount)\
                                    .Expiration(timedelta(minExpiryDTE), timedelta(maxExpiryDTE))

#region imports
from AlgorithmImports import *
#endregion
############################################################
# Long SPY shares with .5% allocated to 30 delta protective puts at 60 DTE, rolled at 30 days.
# Reference (reddit discussion): https://bit.ly/3zJHUIj
############################################################
from datetime import timedelta
from OptionsUtil import *

class LongSPYOTMPut(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2006, 12, 15)
        self.SetCash(1000000) 
        self.equity = self.AddEquity("SPY", Resolution.Minute)
        self.SPYSymbol = self.equity.Symbol
        self.InitParameters()
        self.OptionsUtil = OptionsUtil(self, self.equity)

    ## Initialize parameters (periods, days-till-expiration, etc)
    ## ------------------------------------------------------------
    def InitParameters(self):

        ## Params for our SPY Put contract 
        self.putExpiryDate     = None   # var to track the expiry date
        self.putInitialDTE     = 60 #int(self.ParamManager.GetParameter("putInitialDTE"))          # Enter at 60 days to exp
        self.putExitDTECoeff   = .5 #float(self.ParamManager.GetParameter("putExitDTECoeff"))      # Exit at .5 (half)-way to exp
        self.putStrikeDelta    = .030 #float(self.ParamManager.GetParameter("putStrikeDelta"))/100   # -30% delta
        self.spendPctOnShares  = .995 #float(self.ParamManager.GetParameter("spendPctOnShares"))/100 # 95% on shares
        
        ## schedule routine to run 30 minutes after every market open
        self.Schedule.On(self.DateRules.EveryDay(self.SPYSymbol), \
                         self.TimeRules.AfterMarketOpen(self.SPYSymbol, 30), \
                         self.DailyAtMarketOpen)

    ## Every morning at market open, Check for entries / exits. 
    ## ------------------------------------------------------------
    def DailyAtMarketOpen(self):

        ## If algo is done warming up 
        if ((not self.IsWarmingUp) and self.CurrentSlice.ContainsKey("SPY")):

            # check the number of puts in the portfolio                
            putsInPortfolio = len([x for x in self.Portfolio if (x.Value.Symbol.HasUnderlying and x.Value.Invested)])

            ## If we have no investments or no puts in our portfolio, (re)allocate shares & puts
            if (not self.Portfolio.Invested or (putsInPortfolio == 0)):
                self.SetSharesHoldings() ## Set shares holdings to X% (may involve buying or selling)
                self.BuyOTMPuts()        ## Buy OTM Puts with whatever is left
                
            ## If we have holdings, check to see if we are past our exit DTE for the puts
            elif( self.Portfolio.Invested ):
                currentDTE = (self.putExpiryDate - self.Time).days
                if ( currentDTE <= (self.putInitialDTE * self.putExitDTECoeff) ):
                    for x in self.Portfolio:
                        if x.Value.Invested:
                            assetLabel  = "Puts  " if (x.Value.Symbol.HasUnderlying) else "Shares"
                            assetChange = round(self.Portfolio[x.Value.Symbol].UnrealizedProfitPercent,2)
                            profitLabel = "Profit" if (assetChange > 0) else "loss" 
                            if(x.Value.Symbol.HasUnderlying):
                                self.Liquidate(x.Value.Symbol, tag=f" {currentDTE} DTE. Sold {assetLabel} [ {assetChange}% {profitLabel} ]") 

    ## Allocate X% of available capital to SPY shares
    ## ------------------------------------------------------------
    def SetSharesHoldings(self):
        if(self.CurrentSlice.ContainsKey(self.SPYSymbol) and (self.CurrentSlice[self.SPYSymbol] is not None)): 
            currentSpyPrice = self.CurrentSlice[self.SPYSymbol].Price 
            approxSpendAmt  = math.floor((self.Portfolio.Cash*self.spendPctOnShares) / currentSpyPrice) * currentSpyPrice
            self.SetHoldings("SPY",self.spendPctOnShares, tag=f"Set SPY to ~{self.spendPctOnShares*100}% of Equity. ")
        else:
            self.Log(f"${self.Time} [Warning] Could not adjust SPY shares held. Slice doesnt contian key")

    ## Buy OTM Put Options with all available cash
    ## ------------------------------------------------------------
    def BuyOTMPuts(self):
        
        ## Buy a Put at the specified delta, with the specified expiration date
        putContract   = self.OptionsUtil.SelectContractByDelta(self.SPYSymbol, self.putStrikeDelta, \
                                                               self.putInitialDTE, OptionRight.Put)
        
        if( putContract is None ):
            # self.Log(f"{self.Time} [Error] Could not retrieve Put Contract from chain")
            return
        
        ## Calculate our affordable quantity        
        affordableQty = math.floor(self.Portfolio.Cash / ( putContract.AskPrice * 100 ))

        ## construct an order message -- good for debugging and order rrecords
        orderMessage = f"Buy Puts with ~{round((1-self.spendPctOnShares),2)*100}% of Equity. (${round(self.Portfolio.Cash,2)}) "+\
                       f"- {putContract.Strike} Strike | {(putContract.Expiry-self.Time).days} DTE"+ \
                       f"({round(putContract.Greeks.Delta,2)} Delta)"
                       
        if( affordableQty > 0 ):
            self.Order(putContract.Symbol, affordableQty, False, orderMessage  )   
            self.putExpiryDate = putContract.Expiry

        else:
            self.Log(f"${self.Time} [Warning] {self.Portfolio.Cash} is not enough cash to buy Puts ")
            self.Liquidate()