################################################################################
# "If Only I Had..." - The COVID $ZM Call Roller
# -----------------------------------------------
# A thought experiment on riding Zoom's explosive COVID rally through systematic 
# options trading. This algorithm simulates buying slightly OTM call options on 
# Zoom during the 2020 pandemic surge and rolling them forward for maximum gains.
# 
# Entry: Trading days between 1/1/20-9/30/20, buy ZM calls 120 DTE, strike 15% OTM
# Exit: Liquidate when ITM or at 30 DTE
#
# Not investment advice - just a FOMO-inspired backtesting experiment!
################################################################################

from AlgorithmImports import *
class MiniLeapRoller(QCAlgorithm):

    # ================================================
    # Initialize data, capital, schedulers, etc.
    # ================================================
    def Initialize(self):
        
        # set params: ticker, entry, exit thresholds
        # -----------------------------------------------------
        self.ticker = "ZM"
        self.distFromPrice = int(self.GetParameter("priceDist"))   # pick strike that is this % dist. from price   
        self.enterAtDTE    = int(self.GetParameter("enterDTE"))    # buy option with this many days till expiry 
        self.exitAtDTE     = int(self.GetParameter("exitDTE"))     # sell when this many days left tille expiry
        
        
        # set start/end date for backtest
        # --------------------------------------------
        self.SetStartDate(2020, 1, 1)  # Set Start Date 
        self.SetEndDate(2020, 9, 30) # Set End Date
        
        # set starting balance for backtest
        # --------------------------------------------
        self.SetCash(2000)  

        # add the underlying asset
        # ---------------------------------
        self.equity = self.AddEquity(self.ticker, Resolution.Minute)
        self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.symbol = self.equity.Symbol
        self.forceInitialized = False


        # set custom security intializer
        # -------------------------------
        self.SetSecurityInitializer(self.InitializeSecurities)
        
        # schedule routine to run 30 minutes after every market open
        # --------------------------------------------------------------
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), \
                        self.TimeRules.AfterMarketOpen(self.symbol, 30), \
                        self.OnThirtyMinsIntoMarketOpen)
                  
        # schedule routine to run 30 minutes before every market close
        # --------------------------------------------------------------
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), \
                        self.TimeRules.BeforeMarketClose(self.symbol, 30), \
                        self.OnThirtyMinsBeforeMarketClose)
        
        # set the warmup period
        # ----------------------
        self.SetWarmUp(200)


    # ==============================================================
    # Initialize the security
    # ==============================================================
    def InitializeSecurities(self, security):
        
        # intialize securities with last known price, 
        # so that we can immediately trade the security
        # ------------------------------------------------
        bar = self.GetLastKnownPrice(security)
        security.SetMarketPrice(bar)
        
    
    # ================================================
    # OnData Event handler
    # ================================================
    def OnData(self, data):
        
        # exit if we dont have data
        # ------------------------------
        if self.Securities[self.symbol] is None:
            return
        
        if not self.Securities[self.symbol].HasData:
             return

        # exit if there are no bars or indicators aren't ready
        # ----------------------------------------------------------
        if not data.Bars.ContainsKey(self.symbol): 
            return

        else:    
            
            # get/store  current price of underlying
            # -----------------------------
            self.underlyingPrice = data.Bars[self.symbol].Close


    # ================================================
    # Run 30 minutes after market open
    # ================================================
    def OnThirtyMinsIntoMarketOpen(self):
        
        # exit if we are still warming up
        # ----------------------------------
        if(self.IsWarmingUp):
            return

        # exit if we dont have data
        # --------------------------
        if not self.Securities[self.symbol].HasData:
            return

        # otherwise, if we have no holdings, open a position.
        # --------------------------------------------------
        if not self.Portfolio.Invested:

            
            # set strikes and expiration
            # ------------------------------
            callStrike = self.underlyingPrice * (1 + (self.distFromPrice/100) )
            expiration = self.Time + timedelta(days=self.enterAtDTE)
            
            ### --------------------------------------
            ### todo: use deltas to pick strikes 

            # retrive closest call contracts
            # -------------------------------
            callContract = self.GetLongCall(self.symbol, callStrike, expiration)
            
            # subscribe to data for those contracts
            # -----------------------------------------
            self.AddOptionContract(callContract, Resolution.Minute)

            # buy call 
            # --------------
            self.SetHoldings(callContract, 1)            

            self.Debug (f"{self.Time}   [+]---  BUY  {str(callContract)} || Stock @ {str(self.underlyingPrice)}") 
            
    # ==============================================
    # Run 30 minutes before market close
    # ==============================================
    def OnThirtyMinsBeforeMarketClose(self):
        # exit if we are still warming up
        # ----------------------------------
        if(self.IsWarmingUp):
            return

        # exit if we dont have data
        # --------------------------
        if not self.Securities[self.symbol].HasData:
            return

        # otherwise, check for holdings and close them position.
        # --------------------------------------------------
        for symbol in self.Securities.Keys:
            if self.Securities[symbol].Invested:
                
                # if current contract is ITM, liquidate) 
                # ----------------------------------------
                if (self.underlyingPrice > self.Securities[symbol].StrikePrice):  
                    self.Debug (f"{self.Time}   [-]---  SELL {symbol} |  ITM  | Stock @ {str(self.underlyingPrice)}")
                    self.Debug ("-")
                    self.Liquidate()
                    return
                    
               
                # if current contract is < 30 DTE, liquidate) 
                # ----------------------------------------------
                elif ((self.Securities[symbol].Expiry - self.Time).days < self.exitAtDTE):  
                    self.Debug (f"{self.Time}   [-]---  SELL {symbol} |  <30 DTE  | {(self.Securities[symbol].Expiry - self.Time).days} DTE")
                    self.Debug ("-")
                    self.Liquidate()   
                    return
                
                
    # ========================================================
    # Get Long Call, given a symbol, strike and expiration
    # ========================================================
    def GetLongCall(self, symbolArg, callStrikeArg, expirationArg):

        contracts = self.OptionChainProvider.GetOptionContractList(symbolArg, self.Time)
        
        # get calls
        # -------------
        calls = [symbolArg for symbolArg in contracts if symbolArg.ID.OptionRight == OptionRight.Call]

        # get expiration closest to desired expiration
        # ---------------------------------------
        callsSortedByExpiration = sorted(calls, key=lambda p: abs(p.ID.Date - expirationArg), reverse=False)
        closestExpirationDate = callsSortedByExpiration[0].ID.Date

        # get all contracts with closest expiration
        # ------------------------------------------------
        callsFilteredByExpiration = [contract for contract in callsSortedByExpiration if contract.ID.Date == closestExpirationDate]
        
        # sort contracts to find ones near our desired strikes
        # ------------------------------------------------------
        callsSortedByStrike = sorted(callsFilteredByExpiration, key=lambda p: abs(p.ID.StrikePrice - callStrikeArg), reverse=False)
        
        # pick contract closest to desired strike and expiration
        # ------------------------------------------------
        callOptionContract = callsSortedByStrike[0]
        
        return callOptionContract