##########################################################################################
#
# QQQ Index Rebalancer
# ------------------------
#
# Inspired by:
# https://www.quantconnect.com/forum/discussion/12347/creating-our-own-index-fund/p1
#
# Values for External Params
# ---------------------------------------
# maxHoldings           = 10    # Number of positions to hold, max
# lookbackInDays        = 30    # look at performance over last 30 days    
# perfMetricIndex       = 0     # '0' corresponds to Momentum Pct         
# rebalancePeriodIndex  = 0     # '0' corresponds to Monthly rebalancing
# useETFWeights = 0     = 0     # Honor the ETF's weighting (even if a few holdings) 
#
##########################################################################################

class ETFUniverse(QCAlgorithm):

    ## Main entry point for the algo     
    ## ==================================================================================
    def Initialize(self):
        self.InitAssets()
        self.InitBacktestParams()
        self.InitExternalParams()
        self.InitAlgoParams()
        self.ScheduleRoutines()

    ## Init assets: Symbol, broker model, universes, etc. Called from Initialize(). 
    ## ==================================================================================
    def InitAssets(self):
        self.ticker = "QQQ"
        self.etfSymbol = self.AddEquity(self.ticker, Resolution.Hour).Symbol
        self.AddUniverse(self.Universe.ETF(self.etfSymbol, self.UniverseSettings, self.ETFConstituentsFilter))
        self.UniverseSettings.Resolution = Resolution.Hour
        self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))

    ## Set backtest params: dates, cash, etc. Called from Initialize().
    ## ==================================================================
    def InitBacktestParams(self):
        self.SetStartDate(2012, 1, 1)   # Set Start Date
        # self.SetEndDate(2019, 1, 1)     # Set end Date
        self.SetCash(10000)             # Set Strategy Cash
        self.EnableAutomaticIndicatorWarmUp = True 
        self.SetWarmup(31, Resolution.Daily)

    # ====================================================================
    # Initialize external parameters. Called from Initialize(). 
    # ====================================================================
    def InitExternalParams(self):
        self.maxHoldings           = int(self.GetParameter("maxHoldings"))
        self.perfMetricIndex       = int(self.GetParameter("perfMetricIndex"))
        self.rebalancePeriodIndex  = int(self.GetParameter("rebalancePeriodIndex"))
        self.lookbackInDays        = int(self.GetParameter("lookbackInDays"))
        self.useETFWeights         = bool(self.GetParameter("useETFWeights") == 1)
        
    # ==================================================================================
    # Set algo params: Symbol, broker model, ticker, etc. Called from Initialize(). 
    # ==================================================================================
    def InitAlgoParams(self):

        # Flags to track and trigger rebalancing state
        self.timeToRebalance = True
        self.universeRepopulated = False

        # State vars
        self.weights            = {}
        self.symDataDict        = {}
        self.SelectedSymbols    = []
        self.SelectedSymbolData = []

        # Values related to our alpha (rebalancing period and preferred perf metric 
        self.PerfMetrics      = [MetricsEnum.MOMENTUM_PCT, MetricsEnum.SHARPE]
        self.rankPerfMetric   = self.PerfMetrics[self.perfMetricIndex]

        self.RebalancePeriods = [IntervalEnum.MONTHLY, IntervalEnum.WEEKLY, IntervalEnum.DAILY]
        self.rebalancePeriod  = self.RebalancePeriods[self.rebalancePeriodIndex]

    # ==================================
    def ScheduleRoutines(self):
        if( self.rebalancePeriod == IntervalEnum.MONTHLY ):
            self.Schedule.On( self.DateRules.MonthStart(self.etfSymbol),
                              self.TimeRules.AfterMarketOpen(self.etfSymbol, 31),
                              self.SetRebalanceFlag )

        elif( self.rebalancePeriod == IntervalEnum.WEEKLY ):
            self.Schedule.On( self.DateRules.WeekStart(self.etfSymbol),
                              self.TimeRules.AfterMarketOpen(self.etfSymbol, 31),
                              self.SetRebalanceFlag )

    
    # ====================================
    def ETFConstituentsFilter(self, constituents):
        if( self.timeToRebalance ):

            self.weights = {c.Symbol: c.Weight for c in constituents}
            
            ## set flags
            self.universeRepopulated = True
            self.timeToRebalance = False

            for symbol in self.weights:
                if symbol not in self.symDataDict:
                    
                    # symbol, algo, etfWeight, lookbackInDays, MetricsEnum.MOMENTUM_PCT):
                    self.symDataDict[symbol] = SymbolData(self, symbol, self.weights[symbol], \
                                                          self.lookbackInDays, self.rankPerfMetric)
                
                symbolData = self.symDataDict[symbol]
                
                history = self.History(symbol, 31, Resolution.Daily)
                symbolData.SeedHistory(history)    
            
            # Sort by momentum
            sortedByMomentum = sorted(self.symDataDict.items(), key=lambda x: x[1].PerfMetricValue, reverse=False)[:10]
                        
            
            # Get Symbol object (the key of dict)
            self.SelectedSymbols    = [x[0] for x in sortedByMomentum]
            self.SelectedSymbolData = [x[1] for x in sortedByMomentum]
            
            
            #### TODO: 
            #### Remove Symboldata from self.symDataDict that didnt make the cut
            
            
            return self.SelectedSymbols 

        else: 
            return []    
            
            
    # ====================================
    def SetRebalanceFlag(self):
        self.timeToRebalance = True

    # ====================================
    def OnData(self,data):
        
        if (self.universeRepopulated):
            self.Liquidate()
            self.weights = {}

            
            # calculate sum of weights
            weightsSum = sum(self.symDataDict[symbol].etfWeight for symbol in self.SelectedSymbols)
            
            for symbol in self.SelectedSymbols:

                # Either use ETF Weights or equal weights
                if(self.useETFWeights):
                    symbolWeight = self.symDataDict[symbol].etfWeight / weightsSum # respect weighting
                else:
                    symbolWeight = 1 / len(self.SelectedSymbols) # Equally weighted
                
                self.SetHoldings(symbol, symbolWeight) 
    
            self.universeRepopulated = False    


            # for symbol, weight in self.weights.items():
            #     if symbol in self.ActiveSecurities:
            #         self.SetHoldings(symbol, weight)  # Market cap weighted
            #         # self.SetHoldings(symbol, 1 / len(self.weights))  # Equally weighted
    
    
    # =================================================================
    # Periodically check if a security is no longer in the ETF
    # =================================================================
    # def RemoveDelistedSymbols(self, changes):
    #     for investedSymbol in [x.Key for x in self.Portfolio if x.Value.Invested]:
    #         if( investedSymbol not in self.weights.keys() ):
    #             self.Liquidate(symbol, 'No longer in universe')



##########################
# Symbol Data Class 
##########################
class SymbolData():
    
    def __init__(self, algo, symbol, etfWeight, lookbackInDays, perfMetric):
        self.algo       = algo
        self.symbol     = symbol   
        self.etfWeight  = etfWeight
        self.perfMetric = perfMetric
        self.momp       = MomentumPercent(lookbackInDays)

    def SeedHistory(self,history):
        for row in history.loc[self.symbol].itertuples():
            self.momp.Update(row.Index, row.close)

    @property 
    def PerfMetricValue(self):
        if( (self.perfMetric == MetricsEnum.MOMENTUM_PCT) and self.momp.IsReady):
            return self.momp.Current.Value
        else:
            return float('-inf')
    
###############################
# Perf Enum  
############################### 
class MetricsEnum(Enum):
    
    MOMENTUM_PCT  = "MOMENTUM PCT"
    SHARPE        = "SHARPE"
    DRAWDOWN      = "DRAWDOWN"
    RET_DD_RATIO  = "RET/DD"
    THARP_EXP     = "THARP EXP"

###############################
# Interval Enum  
############################### 
class IntervalEnum(Enum):
    
    MONTHLY  = "MONTHLY"
    WEEKLY   = "WEEKLY"
    DAILY    = "DAILY"