##########################################################################
# The Dynamic 'Mini-Leap' Trader
# Author: Ikezi Kamanu
# Contributors: Leandro Maia
# ----------------------------------------------------
#
# Trade 'mini-LEAP' option contracts (120 - 365 DTE),
# based on daily universe selection.
#
# ----------------------------------------------------
# Entry:
#   Documentation pending.
#
# Position Management:
#   Documentation pending.
#
# Exit:
#   Documentation pending.
#
##########################################################################

from dateutil import parser

class DynamicMiniLeapRoller(QCAlgorithm):

    # ==============================================================
    # Initialize data, capital, scheduled routines, etc.
    # ==============================================================
    def Initialize(self):

        # Initializers (params, universe, routines, etc)
        # ----------------------------------------------------------
        self.InitializeAlgoParams()
        self.InitializeBacktestParams()
        self.SetupUnderlyingSecurity()
        self.InitializeUniverse()
        self.ScheduleRoutines()
        
    # ======================================================
    # Set algo params: ticker, thresholds
    # ======================================================
    def InitializeAlgoParams(self):
        self.ticker = "SPY"
        self.data = {}
        
        self.distFromPrice  = int(self.GetParameter("priceDistPct"))/100 # pick OTM strike that is x% higher than price   
        self.enterAtDTE     = int(self.GetParameter("enterDTE"))         # buy an option with x days till expiry 
        self.exitAtDTE      = int(self.GetParameter("exitDTE"))          # sell the option when x days til expiry
        self.roiTarget      = int(self.GetParameter("roiTargetPct"))/100 # sell when option has gained x% ROI
        self.pctOfAcct      = int(self.GetParameter("pctOfAcct"))/100    # trade with X% of your acct balance 
        self.stopLoss       = int(self.GetParameter("stopLossPct"))/100  # sell when option value has lost -x% 
        self.maxDaysInTrade = int(self.GetParameter("maxDaysInTrade"))   # sell when contract has beeen open for X days 

        self.MarketOpenMoreThanThirtyMins = False
        
    # ======================================================
    # Set backtest params: dates, cash, etc
    # ======================================================
    def InitializeBacktestParams(self):

        # set start/end date for backtest
        # --------------------------------------------
        self.SetStartDate(2020, 4, 1)   # Set Start Date 
        self.SetEndDate(2020, 5, 1)    # Set End Date

        # set starting balance for backtest
        # --------------------------------------------
        self.SetCash(100000)  
        
    # ======================================================
    # Configure underlying, set the custom intializer
    # ======================================================
    def SetupUnderlyingSecurity(self):

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


                        
    # ==============================================================
    # 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
        


    #######################################################
    ###############       Routines       ##################
    ########################################@##############

    # ==============================================================
    # Schedule Routines (similar to chron jobs)
    # ==============================================================
    def ScheduleRoutines(self):
          
        # 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 every 10 minutes
        # -----------------------------------------
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), \
                        self.TimeRules.Every(timedelta(minutes=10)), \
                        self.OnEveryTenMinsDuringMarket)
                
        # 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)

        self.Schedule.On(self.DateRules.EveryDay(),
                        self.TimeRules.AfterMarketOpen(self.SPY, 1),
                        self.OnOneMinuteAfterMarketOpen)
        
        self.Schedule.On(self.DateRules.EveryDay(),
                        self.TimeRules.BeforeMarketClose(self.SPY, 1),
                        self.OnOneMinuteBeforeMarketCloses)
                        

    # ==============================================================
    # Run this Logic 1 minute after market open
    # ==============================================================
    def OnOneMinuteAfterMarketOpen(self):
        self.MarketOpenMoreThanThirtyMins = False
        self.FlushPrevSelectStocks()
        self.UpdateSymbolDataAtMarketOpen()
    
    # ==============================================================
    # Run this Logic 1 minute before market close
    # ==============================================================
    def OnOneMinuteBeforeMarketCloses(self):
        self.MarketOpenMoreThanThirtyMins = False
        self.UpdateSymbolDataAtMarketClose()
                

                        
    # ==============================================================
    # Run this Logic 30 minutes after market open
    # ==============================================================
    def OnThirtyMinsAfterMarketOpen(self):

        self.MarketOpenMoreThanThirtyMins = True
        self.SelectStocksFromUniverse()

        # return 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
        
        self.ManageOpenPositions()
        #self.OpenLongCallIfNotInvested()


    # ==============================================================
    # Run this logic every 10 minutes while market is open
    # ==============================================================
    def OnEveryTenMinsDuringMarket(self):
        if self.MarketOpenMoreThanThirtyMins :
            self.ManageOpenPositions()
            self.OpenLongCallIfNotInvested()


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




    ###################################################################
    ##############   Order Logic (open/close/manage)    ###############
    ###################################################################

    # =======================================================================
    # Open a new option for the most recent stock in our universe selection
    # If the universe is empty, use our default stock
    # =======================================================================
    def OpenLongCallIfNotInvested(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) )
            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()

            if(len(self.selectedStocks) > 0 ):
                theSymbol = self.selectedStocks.pop()
                self.equity = self.AddEquity(theSymbol.Value, Resolution.Minute)
                self.equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
                self.symbol = self.equity.Symbol
                self.OpenLongCall(theSymbol, callStrike, expiration )             
            else:    
                return # universe is empty. do nothing
                # theSymbol = self.symbol # use previously used security if nothing in universe selection
            
             
            
    # ==============================================================
    # Open Long Call
    # ==============================================================
    def OpenLongCall(self, symbolArg, callStrike, expiration ):

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

        # buy call contract
        # -------------------
        orderMsg = "." #"Strike: "+str(callContract.StrikePrice) + " Stock @ " + str(callContract.UnderlyingLastPrice)
        self.SetHoldings(callContract, self.pctOfAcct, False, orderMsg)   
        self.Debug("[[ BUY ]] " + str(theOption) + " | "+  orderMsg + " | "+  self.Time.ctime() )

        # store the opening time in the object store.
        # -------------------------------------------
        # todo: Refactor this to use SymbolData. Currently using suggestion from here: 
        #       https://www.quantconnect.com/forum/discussion/9212/how-to-check-position-age/p1

        self.ObjectStore.Save(str(theOption), str(self.Time))
        
    # ==============================================================
    # Roll Long Call
    # ==============================================================
    # def RollLongCall(self, optionSymbolArg, callStrike, expiration ):
    #     underlyingSymbol = optionSymbolArg.underlyingSymbol
    #     self.Liquidate(optionSymbolArg, orderMsg)
    #     self.OpenLongCall(underlyingSymbol, callStrike, expiration )
      
      
    # ==============================================================
    # Manage positions: take profits, stop losses, rolling, etc.
    # ==============================================================
    def ManageOpenPositions(self):
        
        # check for open contracts and close them if warranted.
        # ------------------------------------------------------
        for symbol in self.Securities.Keys:
            if self.Securities[symbol].Invested:
                
                # Set debug values 
                # ---------------------------------------------------
                currTime    = self.Time
                stockPrice  = self.underlyingPrice
                profitPct   = round(self.Securities[symbol].Holdings.UnrealizedProfitPercent,2)
                daysTillExp = (self.Securities[symbol].Expiry - self.Time).days

                # if current contract is ITM, liquidate
                # This should be the first check, because we want to keep them OTM
                # -----------------------------------------------------------------
                if (self.underlyingPrice > self.Securities[symbol].StrikePrice):  
                    orderMsg = "[ITM] exit @ " + str(profitPct*100) + "% Profit || Stock @ $" + str(stockPrice)
                    self.Liquidate(symbol, orderMsg)    
                    return
    
                # if current contract has hit x% return, liquidate
                # ---------------------------------------------------
                ###
                ### todo: experiment with diff ROI roll targets here. 
                ###       for now, set the param value too high to reach
                ###       so this code never gets called. explore later
                ###
                elif( profitPct >= (self.roiTarget)):
                    orderMsg = "[TP] exit @ " + str(profitPct*100) + "% Profit || Stock @ $" + str(stockPrice)
                    self.Liquidate(symbol, orderMsg)
                    return

                # if current contract has hit x% loss, liquidate
                # ---------------------------------------------------
                elif( profitPct <= (-self.stopLoss)):
                    orderMsg = "[SL] exit @ " + str(profitPct*100) + "% Profit || Stock @ $" + str(stockPrice)
                    self.Liquidate(symbol, orderMsg)
                    return               

                # if current position has been open for X days, liquidate
                # --------------------------------------------------------
                elif self.ObjectStore.ContainsKey(str(symbol)):
                    date = self.ObjectStore.Read(str(symbol))
                    date = parser.parse(date)
                
                    if (self.Time - date).days >= self.maxDaysInTrade:
                        self.ObjectStore.Delete(str(symbol))
                        orderMsg = "[Duration] exit @ "+str(self.maxDaysInTrade)+" days @ " + str(profitPct*100) + "% Profit || Stock @ $" + str(stockPrice) + " || " + str(daysTillExp) + " DTE"
                        self.Liquidate(symbol, orderMsg)
                        
                    return    
                    
               
                # if current contract expiry is less than 'X' liquidate 
                # -----------------------------------------------------------
                elif ((self.Securities[symbol].Expiry - self.Time).days < self.exitAtDTE):  
                    orderMsg = "[DTE] exit @ " + str(profitPct*100) + "% Profit || Stock @ $" + str(stockPrice) + " || " + str(daysTillExp) + " DTE"
                    self.Liquidate(symbol, orderMsg)    
                
                    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
        
        
    ######################################################################
    ###########      Universe / Stock Selection Logic      ###############
    ######################################################################


    # ==============================================================
    # Initialize the universe, and universe selection criteria
    # ==============================================================
    def InitializeUniverse(self):
        
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.SPY = self.AddEquity('SPY', Resolution.Minute).Symbol
        
        self.day = -1
        self.num_coarse = 20
        self.min_stock_price = 10

        # set params used for filter criteria
        # ------------------------------------
        self.min_days_after_earnings = 10
        self.max_days_after_earnings = 80
        self.ema_period = 40 # 8 on 5 min timeframe
        self.sma_period = 275 # 55 on 5 min timeframe
        self.bb_period = 600 # 20 on 30 min timeframe
        self.bb_k = 2
        self.gap_distance = 0.02 # 2%
        self.data = {}
        self.selectedStocks = []
        
        
    # ==============================================================
    # Universe Coarse selection logic.
    # ---------------------------------
    # Find stocks with fundamentals with price > 5, 
    # rank by dollar volume, and then pick the top X
    # ==============================================================
    def CoarseSelectionFunction(self, coarse):
        
        if self.day == self.Time.day:
            return Universe.Unchanged

        self.day = self.Time.day

        # drop stocks which have no fundamental data or have too low prices
        selected = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > self.min_stock_price)]
        
        # rank the stocks by dollar volume 
        # -----------------------------------
        filtered = sorted(selected, key=lambda x: x.DollarVolume, reverse=True)
        return [ x.Symbol for x in filtered[:self.num_coarse]]
    
    # ========================================================================
    # Universe Fine selection logic 
    # ---------------------------------
    # Find common stocks where the share class is not a depository receipt,
    # and have neither had recent earnings reports nor any coming soon.
    #
    # todo: explore other ways to account for future earnings. 
    # ========================================================================
    def FineSelectionFunction(self, fine):
        
        filtered = [x for x in fine if x.SecurityReference.IsPrimaryShare
                    and x.SecurityReference.SecurityType == "ST00000001"
                    and x.SecurityReference.IsDepositaryReceipt == 0
                    and x.CompanyReference.IsLimitedPartnership == 0
                    and x.EarningReports.FileDate < self.Time - timedelta(days=self.min_days_after_earnings)
                    and x.EarningReports.FileDate > self.Time - timedelta(days=self.max_days_after_earnings)]

        return [x.Symbol for x in filtered]
        
    # ====================================================================
    # Flush stocks previously selected from universe 
    # ====================================================================
    def FlushPrevSelectStocks(self):
        self.selectedStocks.clear()

    # =====================================================================
    # Select new stocks from universe that pass technical indicator filters 
    # =====================================================================
    def SelectStocksFromUniverse(self):

        for symbol in self.data.keys():
            if (self.data[symbol].GapUp) and (self.data[symbol].EMA > self.data[symbol].SMA) and \
               (self.data[symbol].BB.UpperBand.Current.Value < self.Securities[symbol].Close):

                # if stock passes criteria (ie: there is signal), add to our 'selected stocks'
                # ----------------------------------------------------------------------------
                self.selectedStocks.append(symbol)
                self.Debug("   + " + str(symbol).split(" ")[0] + "\t Added on " + self.Time.ctime())
    
    
    # ============================================================================
    # When securities are added or removed, update our internal symbol dictionary 
    # ============================================================================
    def OnSecuritiesChanged(self, changes):
        
        for security in changes.RemovedSecurities:
            if security.Symbol in self.data:
                del self.data[security.Symbol]
        
        for security in changes.AddedSecurities:
            if (security.Symbol not in self.data)  and \
               (security.Symbol.SecurityType == SecurityType.Equity): 
                self.data[security.Symbol] = SymbolData(security.Symbol, self.ema_period, self.sma_period, self.bb_period, self.bb_k, self)
                
    
    # ============================================================================
    # Update state of SymbolData at market open 
    # ============================================================================
    def UpdateSymbolDataAtMarketOpen(self):
        for symbol in self.data.keys():
            gap = (self.Securities[symbol].Close - self.data[symbol].LastClose) / self.data[symbol].LastClose
            if gap > self.gap_distance:
                self.data[symbol].GapUp = True
            else:
                self.data[symbol].GapUp = False
            
    # ============================================================================
    # Update state of SymbolData at market close 
    # ============================================================================
    def UpdateSymbolDataAtMarketClose(self):
        for symbol in self.data.keys():
            self.data[symbol].LastClose = self.Securities[symbol].Close
            

class SymbolData(object):
    def __init__(self, symbol, ema, sma, bb, k, algorithm):
        self.Symbol = symbol
        self.LastClose = 0
        self.GapUp = False
        self.EMA = ExponentialMovingAverage(ema)
        self.SMA = SimpleMovingAverage(sma)
        self.BB = BollingerBands(bb, k, MovingAverageType.Exponential)
        
        algorithm.RegisterIndicator(symbol, self.EMA, Resolution.Minute, Field.Close)
        algorithm.RegisterIndicator(symbol, self.SMA, Resolution.Minute, Field.Close)
        algorithm.RegisterIndicator(symbol, self.BB, Resolution.Minute, Field.Close)
        
        # Logic for 'manual' warmup . 
        # --------------------------------------
        
        # Check for daily data. we need at least one day
        # ------------------------------------------------
        history = algorithm.History(symbol, 1, Resolution.Daily)
        if history.empty or 'close' not in history.columns:
            return
        
        for index, row in history.loc[symbol].iterrows():
            self.LastClose = row['close']
            

        # Check for minute data. we need at least the MAX of our indicator periods
        # --------------------------------------------------------------------------
        history = algorithm.History(symbol, max(ema, sma, bb), Resolution.Minute)
        if history.empty or 'close' not in history.columns:
            return
        for index, row in history.loc[symbol].iterrows():
            self.EMA.Update(index, row['close'])
            self.SMA.Update(index, row['close'])
            self.BB.Update(index, row['close'])