##########################################################################
# The 'Mini-Leap' Roller
# Author: Ikezi Kamanu
# -------------------------------------------------
#
# Entry:
#   30 Mins before market close, if the portfolio is empty, 
#   open as many single OTM call options as possible, with
#   strikes that are 10% above the underlying, expiring in 120 days.
#
# Exit:
#   30 Mins after market open, for every call in the  portfolio, 
#   if the call is ITM or there are less than 30 days till expiry, 
#   liquidate the position. They will be re-opened (rolled) before close. 
#
##########################################################################

class MiniLeapRoller(QCAlgorithm):

    # ==============================================================
    # Initialize data, capital, scheduled routines, etc.
    # ==============================================================
    def Initialize(self):
        
        # set params: ticker, entry, exit thresholds
        # -----------------------------------------------------
        self.ticker = "ZM"
        self.distFromPrice = int(self.GetParameter("priceDist"))   # pick strike that is this % above 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.OnThirtyMinsAfterMarketOpen)
                  
        # 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):
        
        # if we have data
        # ------------------------------
        if (self.Securities[self.symbol] is not None) and \
           (self.Securities[self.symbol].HasData) and \
           (data.Bars.ContainsKey(self.symbol)):

            # keep track of the current close price 
            # ------------------------------------------
            self.underlyingPrice = data.Bars[self.symbol].Close
        

    # ==============================================================
    # Run this 30 minutes before market close
    # ==============================================================
    def OnThirtyMinsBeforeMarketClose(self):

        # exit if we are still warming up
        # or if we do not have any data  
        # ----------------------------------
        if(self.IsWarmingUp) or \
          (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: consider using ATR for strike selection 
            ### todo: consider using delta for strike selection
            ### todo: encapsulate strike and expiry selection in GetLongCall()

            # 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 ("\r+-------------")
            self.Debug (f"| {self.Time}   [+]---  BUY  {str(callContract)} || Stock @ {str(self.underlyingPrice)}") 
            
            
    # ==============================================================
    # Run this 30 minutes after market open
    # ==============================================================
    def OnThirtyMinsAfterMarketOpen(self):

        # exit if we are still warming up
        # or if we do not have any data  
        # ----------------------------------
        if(self.IsWarmingUp) or \
          (not self.Securities[self.symbol].HasData):
            return

        # check for open contracts and close them if warranted.
        # ------------------------------------------------------
        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 expiry is less than 'X' 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, desired strike and expiration
    # ==============================================================
    def GetLongCall(self, symbolArg, callStrikeArg, expirationArg):

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

        # sort contracts by expiry dates and select 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 for selected expiration
        # ------------------------------------------------
        callsFilteredByExpiration = [contract for contract in callsSortedByExpiration if contract.ID.Date == closestExpirationDate]
        
        # sort contracts and select the one closest to desired strike
        # -----------------------------------------------------------
        callsSortedByStrike = sorted(callsFilteredByExpiration, key=lambda p: abs(p.ID.StrikePrice - callStrikeArg), reverse=False)
        callOptionContract = callsSortedByStrike[0]
        
        return callOptionContract