##################################
# 
# SymbolData Class
#
##################################
class SymbolData():
    
    ## Constructor
    ## -----------
    def __init__(self, theSymbol, algo):

        ## Algo / Symbol / Price reference
        self.algo                 = algo
        self.symbol               = theSymbol
        self.lastDailyClose       = None
        self.lastClosePrice       = None
        
        ## Initialize our Indicators and rolling windows
        self.ema                  = ExponentialMovingAverage(self.algo.dailyEMAPeriod)
        self.momentum             = MomentumPercent(self.algo.hourlyMomPeriod)
        self.lastDailyCloseWindow = RollingWindow[float](2)  
        self.emaWindow            = RollingWindow[float](2)  
    
        ## These will hold our 'messages' used in our order notes
        self.ClosePositionMessage = ""
        self.OpenPositionMessage  = ""

    ## Seed Daily indicators with history.
    ## -----------------------------------
    def SeedDailyIndicators(self, dailyHistory):

        # Loop over the history data and update our indicators
        if dailyHistory.empty or 'close' not in dailyHistory.columns:
            # self.algo.Log(f"No Daily history for {self.symbol}")
            return
        
        else:
            for timeIndex, dailyBar in dailyHistory.loc[self.symbol].iterrows():
                if(self.ema is not None): 
                    self.ema.Update(timeIndex, dailyBar['close'])
                    self.lastDailyClose = dailyBar['close']
                    self.timeOflastDailyClose = timeIndex
                    
                    self.lastDailyCloseWindow.Add( dailyBar['close'] )
                    self.emaWindow.Add( self.ema.Current.Value )

    ## Seed intraday indicators with history
    ## These indicators might be have eitehr hourly or minute resolution
    ## -----------------------------------------------------------------
    def SeedIntradayIndicators(self, hourlyHistory=None, minuteHistory=None):
        
        # Loop over the history data and update our indicators
        if hourlyHistory is not None:    
            if hourlyHistory.empty or 'close' not in hourlyHistory.columns:
                self.algo.Log(f"Missing hourly history for {self.symbol}")
            else:
                for timeIndex, hourlyBar in hourlyHistory.loc[self.symbol].iterrows():
                    if(self.momentum is not None): 
                        self.momentum.Update(timeIndex, hourlyBar['close'])

        if minuteHistory is not None:    
            if  minuteHistory.empty or 'close' not in minuteHistory.columns:
                self.algo.Log(f"Missing minute history for {self.symbol}")
            else:
                for timeIndex, hourlyBar in hourlyHistory.loc[self.symbol].iterrows():
                    if(self.someIndicator is not None): 
                        self.someIndicator.Update(timeIndex, hourlyBar['close'])
            return
        
    ## Daily screening criteria. Called by the main algorithm. 
    ## Returns true if daily screening conditions are met. 
    ## Replace with your own criteria.
    ## -------------------------------------------------------
    def DailyScreeningCriteriaMet(self):
    
        ## Price is above ema.
        if(self.lastDailyCloseWindow[0]  > self.emaWindow[0]):
                return True
                
        return False

    ## Intraday screening criteria. Called by the main 
    ## algorithm. Returns true if entry conditions are met. 
    ## Replace with your own criteria.
    ## ----------------------------------------------------
    def IntradayScreeningCriteriaMet(self):

        ## If we have positive momentum
        if (self.momentum.Current.Value > 0):

            ## Informative message that will be submitted as order notes
            self.OpenPositionMessage = f"OPEN (${round(self.lastDailyCloseWindow[0],3)} > EMA: {round(self.emaWindow[0],3)})"
            return True
            
        return False

    ## Trade Exit criteria. Called by the main algorithm.
    ## Returns true if exit conditions are met. 
    ## Replace with your own criteria.
    ## ---------------------------------------------------
    def ExitCriteriaMet(self):

        ## Exit If price is below ema
        if( self.lastClosePrice < self.ema.Current.Value ):
            self.ClosePositionMessage = f"CLOSE (${round(self.lastClosePrice,3)} < EMA: {round(self.ema.Current.Value,3)})"
            return True
        return False

    ## Returns true if our daily indicators are ready.
    ## Called from the main algo
    ## -----------------------------------------------
    def DailyIndicatorsAreReady(self):
        return (self.ema.IsReady and self.lastDailyCloseWindow.IsReady)

    ## Returns true if our intraday indicators are ready.
    ## Called from the main algo
    ## --------------------------------------------------
    def IntradayIndicatorsAreReady(self):
        return (self.momentum.IsReady)
        
    ## Called by the main algorithm right after 
    ## a position is opened for this symbol.
    ## ----------------------------------------
    def OnPositionOpened(self):
        
        ## Register & warmup the indicators we need to track for exits
        self.algo.RegisterIndicator(self.symbol, self.ema, timedelta(1)) 
        self.algo.WarmUpIndicator(self.symbol, self.ema, Resolution.Daily)
        
    ## Called by the main algorithm right after 
    ## a position is closed for this symbol.
    ## ----------------------------------------
    def OnPositionClosed(self):
        # cleanup
        pass

##########################################################################
# Scheduled Intraday Universe Screening
# ---------------------------------------------
# FOR EDUCATIONAL PURPOSES ONLY. DO NOT DEPLOY.
#
# Entry:
# ------
# Daily: At midnight, screen for stocks trading above daily EMA 
# Intraday: In the afternoon, screen those daily stocks for postive momentum
# Open positions for the top 'X' stocks with highest positive momentum
#
# Exit:
# -----
# Exit when price falls below EMA.
# Optionally: exit at End of day if EoDExit flag is set. 
#
# ................................................................
# Copyright(c) 2021 Quantish.io - Granted to the public domain
# Do not remove this copyright notice | info@quantish.io
#########################################################################

from SymbolData import *

class EMAMOMUniverse(QCAlgorithm):
    
    def Initialize(self):
        self.InitBacktestParams()
        self.InitAssets()
        self.InitAlgoParams()
        self.InitUniverse()
        self.ScheduleRoutines()
        
    def InitBacktestParams(self):
        self.SetStartDate(2020, 1, 1)
        self.SetEndDate(2021, 1, 1)
        self.SetCash(100000)

    def InitAssets(self):
        self.AddEquity("SPY", Resolution.Hour) # benchmark 
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))

    def InitAlgoParams(self):
        self.dailyEMAPeriod  = 10 
        self.hourlyMomPeriod = 4
        self.minsAfterOpen   = 300
        self.useEoDExit      = 0

        self.maxCoarseSelections    = 30
        self.maxFineSelections      = 10
        self.maxIntradaySelections  = 5
        self.maxOpenPositions       = 5

    def InitUniverse(self):
        ## Init universe configuration, selectors 
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseUniverseSelection, self.FineUniverseSelection) 
        self.EnableAutomaticIndicatorWarmUp = True 
        
        ## Init vars for tracking universe state
        self.symDataDict = { }
        self.screenedDailyStocks = []
        self.screenedIntradayStocks = []
        self.queuedPositions        = []


    ## Schedule screening and liquidation routines, as needed.
    ## -------------------------------------------------------
    def ScheduleRoutines(self):
        ## Intraday selection
        self.Schedule.On(self.DateRules.EveryDay(), 
                         self.TimeRules.AfterMarketOpen("SPY", self.minsAfterOpen), 
                         self.IntraDaySelection)         
        
        ## End of Day Liquidation
        if(self.useEoDExit):
            self.Schedule.On(self.DateRules.EveryDay(), 
                             self.TimeRules.BeforeMarketClose("SPY", 2), 
                             self.LiquidateAtEoD)         

    ## ----------------------
    def LiquidateAtEoD(self):
         self.Liquidate(tag="EoD liquidatation")
        
    ## Process our queued positions, check for exits for held positions
    ## ----------------------------------------------------------------
    def OnData(self, dataSlice):
        
        self.ProcessQueuedPositions()        

        for symbol in dataSlice.Keys:
            if symbol in self.symDataDict:
                symbolData = self.symDataDict[symbol] 

                if( (symbol in dataSlice) and (dataSlice[symbol] is not None)):
                    symbolData.lastClosePrice = dataSlice[symbol].Close
                    if( symbolData.ExitCriteriaMet() ):
                        self.Liquidate(symbol, tag=symbolData.ClosePositionMessage)
                        # del self.symDataDict[symbol]
                        self.RemoveSecurity(symbol)
    
    
    ## Check if we are already holding the max # of open positions.
    ## ------------------------------------------------------------
    def PortfolioAtCapacity(self):
        numHoldings = len([x.Key for x in self.Portfolio if x.Value.Invested])
        return ( numHoldings >= self.maxOpenPositions )
        
        
    ## Coarse universe selection. Replace with your own universe filters.
    ## ------------------------------------------------------------------
    def CoarseUniverseSelection(self, universe): 
        
        if (self.PortfolioAtCapacity()):        
            return []
        else:
            coarseuniverse = sorted(universe, key=lambda c: c.DollarVolume, reverse=True)  
            coarseuniverse = [c for c in coarseuniverse if c.Price > 50][:self.maxCoarseSelections]
            return [x.Symbol for x in coarseuniverse]

    ## Fine universe selection. Replace with your own universe filters.
    ## ----------------------------------------------------------------
    def FineUniverseSelection(self, universe):
        if (self.PortfolioAtCapacity()):        
            return []    
        else:
            fineUniverse = [x for x in universe if x.SecurityReference.IsPrimaryShare
                                    and x.SecurityReference.SecurityType == "ST00000001"
                                    and x.SecurityReference.IsDepositaryReceipt == 0
                                    and x.CompanyReference.IsLimitedPartnership == 0]
            
            
            ## Fetch the stocks that match our daily screening criteria
            screenedStocks = self.GetDailyScreenedStocks(fineUniverse)
            self.screenedDailyStocks = screenedStocks[:self.maxFineSelections]  
            
            ## NOTE:
            ## If you plan to do intraday selection, then 
            ## return a blank array otherwise, comment out this line
            ## ...........................................................
            return []
            
            ## NOTE: 
            ## If there is no need for intraday selection, then 
            ## uncomment the below line to return daily screened stocks
            ## ...........................................................
            ## return self.screenedDailyStocks
    
        

    ## Intraday universe selection. Replace with your own Criteria.
    ## ------------------------------------------------------------
    def IntraDaySelection(self):

        if (not self.PortfolioAtCapacity()):

            ## Fetch the stocks that meet intraday screening criteria
            screenedStocks      = self.GetIntraDayScreenedStocks(self.screenedDailyStocks)
            
            ## Get the symboldata for the screened stocks, and rank by momentum  
            screenedStockData   = [ self.symDataDict[symbol] for symbol in screenedStocks]    
            screenedDataSorted = sorted(screenedStockData, key=lambda x: x.momentum, reverse=True)  
            screenedSymbolsSorted = [ stockData.symbol for stockData in screenedDataSorted ] 
            
            self.screenedIntradayStocks = screenedSymbolsSorted[:self.maxIntradaySelections]
    
            for stock in self.screenedIntradayStocks:
                self.AddSecurity(SecurityType.Equity, stock, Resolution.Minute)
    
        self.screenedDailyStocks = []
    
    ## Screen the given array of stocks for those matching *Daily* 
    ## screening criteria, and return them. 
    ##
    ## Seeds the stock symboldata class with daily history, and then
    ## calls the class's DailyScreeningCriteriaMet() method, where
    ## the actual daily screening logic lives (eg: indicator checks).
    ## --------------------------------------------------------------
    def GetDailyScreenedStocks(self, stocksToScreen):

        screenedStocks = []
        
        for stock in stocksToScreen:  
            symbol = stock.Symbol

            ## If we are already invested in this, skip it
            if (symbol in self.Portfolio) and (self.Portfolio[symbol].Invested):
                continue
            else:
                ## Store data for this symbol in our dictionary, seed it with some history     
                if symbol not in self.symDataDict:
                    self.symDataDict[symbol] = SymbolData(symbol, self) 
                    
                symbolData = self.symDataDict[symbol]
                
                ## we need at least 2 values for EMA for our 
                ## first signal, so we get the required history + 1 day 
                dailyHistory = self.History(symbol, self.dailyEMAPeriod+1, Resolution.Daily)
                
                ## Seed daily indicators so they can be calculated
                symbolData.SeedDailyIndicators(dailyHistory)
                    
                ## If the daily screening criteria is met, we return it
                if symbolData.DailyIndicatorsAreReady(): 
                    if symbolData.DailyScreeningCriteriaMet():
                        screenedStocks.append(symbol)
                    else:
                        ## if the criteria isnt met, we dont need this symboldata
                        del self.symDataDict[symbol]

        return screenedStocks
        
    ## Screen the given array of stocks for those matching *Intraday* 
    ## screening criteria, and return them. 
    ##
    ## Seeds the stock symboldata class with intraday history, and then
    ## calls the class's IntradayScreeningCriteriaMet() method, where
    ## the actual intraday screening logic lives (eg indicator checks).
    ## --------------------------------------------------------------
    def GetIntraDayScreenedStocks(self, stocksToScreen):

        screenedStockSymbols = []
        
        ## loop through stocks and seed their indicators
        for symbol in stocksToScreen:  

            ## If we are already invested in this, skip it
            if (symbol in self.Portfolio) and (self.Portfolio[symbol].Invested):
                continue
            else:
                if( symbol in self.symDataDict ):
                    symbolData = self.symDataDict[symbol]
                    history = self.History(symbol, self.hourlyMomPeriod, Resolution.Hour)
    
                    symbolData.SeedIntradayIndicators(history)
                    if( symbolData.IntradayIndicatorsAreReady() ):
                        if ( symbolData.IntradayScreeningCriteriaMet() ):
                            screenedStockSymbols.append( symbolData.symbol )
                        else:
                            ## if the criteria isnt met, we dont need this symboldata
                            del self.symDataDict[symbol]                    
                else:
                    self.Log(f"- - - - No symdata for {symbol}")
                    
        return screenedStockSymbols
      
      
    ## Called when we add/remove a security from the algo  
    ## --------------------------------------------------
    def OnSecuritiesChanged(self, changes):
        
        ## The trade actually takes place here, when the symbol
        ## passes all our screenings, and we call AddSecurity
        if (not self.PortfolioAtCapacity()):
            for security in changes.AddedSecurities:
                if(security.Symbol != "SPY"):
                    
                    # if we havent already queued this position, queue it. 
                    ## we queue it instead of opening the position, because
                    ## we dont yet have data for this symbol. We will get 
                    ## data for it in the next call to OnData, where we do
                    ## open the position
                    if( security.Symbol not in self.queuedPositions ):
                        self.queuedPositions.append(security.Symbol)
                        
        for security in changes.RemovedSecurities:
            if(security.Symbol != "SPY"):            
                if security.Symbol in self.symDataDict:
                    symbol = security.Symbol
                    symbolData = self.symDataDict[symbol]
                    symbolData.OnPositionClosed()
                    
                    ## remove this sumbol from our local cache
                    del self.symDataDict[symbol] 
                    self.screenedDailyStocks = [ x for x in self.screenedDailyStocks if x != symbol]                    
                    self.screenedIntradayStocks = [ x for x in self.screenedIntradayStocks if x != symbol]                    
        
        
    ## Loop through queued positions and open new trades  
    ## -------------------------------------------------
    def ProcessQueuedPositions(self):
        for symbol in list(self.queuedPositions):
            if self.CurrentSlice.ContainsKey(symbol) and self.CurrentSlice[symbol] is not None:
                symbolData = self.symDataDict[symbol]
                
                ## extra check to make sure we arent going above capacity
                if(not self.PortfolioAtCapacity()):
                    self.SetHoldings(symbol, 1/self.maxOpenPositions, tag=symbolData.OpenPositionMessage)
                    symbolData.OnPositionOpened()
                self.queuedPositions.remove(symbol)