#region imports
from AlgorithmImports import *
#endregion


##########################
# 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)
        self.ker        = KaufmanEfficiencyRatio(lookbackInDays)
        self.atr        = AverageTrueRange(14)
        
        # Track Price 
        # -----------
        self.lastPrice          = 0
        self.entryPrice         = 0
        self.lastDailyBar       = None
        self.lastHourlyBar      = None
        
        # Messages
        # ---------
        self.ExitMessage        = ""
        self.EntryMessage       = ""

        # Stop Loss State
        # -------------------------------
        self.initialStopLoss    = 0
        self.trailStopATRCoef   = 3
        self.ResetStopLosses()

    def SeedHistory(self,history):
        for row in history.loc[self.symbol].itertuples():
            self.momp.Update(row.Index, row.close)
            self.ker.Update(row.Index, row.close)
            
            
    ## ====================================
    def OnSymbolHourlyData(self, tradeBar):
        if(tradeBar.Period == timedelta(seconds=3600)):
            if(self.lastHourlyBar != tradeBar):
                self.lastPrice      = tradeBar.Close
                self.lastHourlyBar   = tradeBar
                

    ## ====================================
    def OnSymbolDailyData(self, tradeBar):
        if(tradeBar.Period == timedelta(days=1)):
            if(self.lastDailyBar != tradeBar):
                self.atr.Update(tradeBar)
                self.lastPrice      = tradeBar.Close
                self.lastDailyBar   = tradeBar
                self.PlotCharts()

    @property 
    def PerfMetricValue(self):        
        if( (self.perfMetric == MetricsEnum.MOMENTUM_PCT) and self.momp.IsReady):
            if (self.momp.Current.Value > 0):
                return self.momp.Current.Value

        elif( (self.perfMetric == MetricsEnum.MOMP_KER) and self.ker.IsReady and self.momp.IsReady):
            if(self.momp.Current.Value > 0):
                # calculate a new composite metric that combines both
                return (self.ker.Current.Value * self.momp.Current.Value)
                
            
            # if( self.ker.Current.Value > 0.4 and self.momp.Current.Value > 0):
            #     return self.ker.Current.Value
        
        # elif( (self.perfMetric == MetricsEnum.KER) and self.ker.IsReady and self.momp.IsReady):
            return self.ker.Current.Value if self.momp.Current.Value > 0 else -abs(self.ker.Current.Value)
    
        
        return float('-inf')

    # Trailing Stop Exit    
    # =====================================
    def TrailingExitSignalFired(self):
        if( not self.atr.IsReady ):
            return False   

        # If trailing stop is NOT set, get last price, and set it
        # --------------------------------------------------------
        if( not self.trailStopActivated ):
            self.highestPrice       = self.lastPrice
            self.trailingStopLoss   = self.lastPrice - (self.atr.Current.Value * self.trailStopATRCoef)

            # Make sure we are never below 2% of entry price
            self.trailingStopLoss   = max (self.trailingStopLoss, self.lowestStopLoss)

            self.trailStopActivated = True
            
            # Recursively call this function to check for stops
            # again, now that the trailing stop has been activated
            # and the stop loss value has been updated.
            # --------------------------------------------------
            return self.TrailingExitSignalFired()            
            
        # If trailing stop loss is activated, check if price closed below it.
        # If it did, then exit. If not, update the trailing stop loss. 
        # -------------------------------------------------------------------
        else: 
            if self.PriceIsBelowTrailingStop():              
                return True

            else:
                # Udpate trailing stop loss if price has gone up
                # -----------------------------------------------
                if self.lastPrice > self.highestPrice:
                    self.highestPrice = self.lastPrice
                    newTrailingStopLoss = self.highestPrice -  (self.atr.Current.Value * self.trailStopATRCoef)
                    
                    # Make sure we are never below 2% of entry price
                    self.trailingStopLoss    = max (self.trailingStopLoss, newTrailingStopLoss, self.lowestStopLoss)
                    
                    if( self.symbol.Value == "BOOT"): #ADMA
                        self.algo.Log(f"[{self.symbol}][Trailing Stop Updated] ${self.trailingStopLoss:.2f} (Price: ${self.lastPrice:.2f})")
                    
                    # check again just in case price ends up below the new trailing stop level
                    if self.PriceIsBelowTrailingStop():
                        return True

        return False 

    # ------------------------------------------------
    def PriceIsBelowTrailingStop(self):
        if( self.lastPrice < self.trailingStopLoss ):
            # self.algo.LogUtil.Info(f"[{self.symbol}][Trailing Stop Triggered] : ${self.lastPrice:.2f}")
            self.ExitMessage = "Trailing Stop Triggered"
            # self.PlotCharts()
            return True

    ## Logic to run immediately after a new position is opened.
    ## ---------------------------------------------------------
    def OnPositionOpened(self):            
    
        # self.algo.LogUtil.Info(f"Bought {self.symbol.Value}] @ approx ${self.lastPrice:.2f}") 
        self.EntryMessage   = ""
        self.entryPrice = self.lastPrice
        self.algo.WarmUpIndicator(self.symbol, self.atr)
        self.algo.AddEquity(self.symbol.Value, Resolution.Daily)
        
        # For trailing stop
        # ------------------
        self.SetInitialStops()
        return
    
    
    
    ## Logic to run immediately after a position is closed
    ## ---------------------------------------------------------
    def OnPositionClosed(self):            
        self.PlotCharts()
        # self.algo.Log(f"Sold {self.symbol.Value}] @ approx ${self.lastPrice:.2f}")         
        self.ExitMessage  = "No Exit Message"
        self.ResetStopLosses()
        
    # ========================================================================
    def ResetStopLosses(self):
        self.trailStopActivated   = False
        self.initialStopLoss      = 0
        self.trailActivationPrice = 0
        self.trailingStopLoss     = 0
        
    # ========================================================================
    # Set initial stop and activation level. Called after new position opened.
    # ========================================================================
    def SetInitialStops(self):
        ## TODO: Use onOrderEvent to set this, because the actual price may be different
        self.entryPrice         = self.lastPrice
        self.initialStopLoss    = self.lastPrice - (self.atr.Current.Value * self.trailStopATRCoef)        

        # Make sure we are never below 2% of entry price
        self.initialStopLoss    = max(self.initialStopLoss, self.entryPrice * 0.98)
        self.lowestStopLoss     = self.initialStopLoss    


    # ====================
    def PlotCharts(self):    
    
        return 
        # if( self.symbol.Value != "SLAB"):
        #     return
        
        # Plot Price on a single chart for this symbol 
        # ------------------------------------------------
        self.algo.Plot(f"{self.symbol.Value}-charts", "Price", self.lastPrice)        
        
        # Stop losses & activation levels
        # ------------------------------------
        # self.algo.Plot(f"{self.symbol}-charts", "Initial Stop", self.initialStopLoss)            
        # self.algo.Plot(f"{self.symbol}-charts", "Acivation Pt", self.trailActivationPrice)            
        self.algo.Plot(f"{self.symbol}-charts", "TrailingStop", self.trailingStopLoss)            
                
        return 

###############################
# Perf Enum  
############################### 
class MetricsEnum(Enum):
    
    MOMENTUM_PCT  = "MOMENTUM PCT"
    MOMP_KER      = "MOMP_KER"
    SHARPE        = "SHARPE"
    DRAWDOWN      = "DRAWDOWN"
    RET_DD_RATIO  = "RET/DD"
    THARP_EXP     = "THARP EXP"


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

#region imports
from AlgorithmImports import *
#endregion
# https://www.quantconnect.com/forum/discussion/12347/creating-our-own-index-fund/p1
##########################################################################################
#
# QQQ Index Rebalancer
# ------------------------
#
# Inspired by:
# https://www.quantconnect.com/forum/discussion/12347/creating-our-own-index-fund/p1
#
# Values for External Params
# ---------------------------------------
# maxHoldings           = 5     # Number of positions to hold, max
# lookbackInDays        = 50    # 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) 
# 
# Todo
# ----------------------------------------
# Add Hourly Stop Tracked ATR StopLoss
#
# Queue position 
#   - instead of trying to setHolding right away, queue till next data slice
#   - next data slice may be next day, so lets subecribe to minute data on the spot. 
#
# Plot ker distribution
#
##########################################################################################

# import matplotlib.pyplot as plt
from SymbolData import *
class ETFUniverse(QCAlgorithm):

    ## Main entry point for the algo     
    ## ==================================================================================
    def Initialize(self):
        self.InitAssets()
        self.InitExternalParams()
        self.InitBacktestParams()
        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(2010, 1, 1)   # Set Start Date
        # self.SetEndDate(2017, 1, 1)     # Set end Date

        self.SetCash(100000)             # Set Strategy Cash
        self.EnableAutomaticIndicatorWarmUp = True 
        self.SetWarmup(self.lookbackInDays+1, Resolution.Daily)
        self.SetBenchmark("SPY")

    # ====================================================================
    # Initialize external parameters. Called from Initialize(). 
    # ====================================================================
    def InitExternalParams(self):
        self.maxExposurePct         = float(self.GetParameter("maxExposurePct"))
        self.maxHoldings            = int(self.GetParameter("maxHoldings"))
        self.perfMetricIndex        = int(self.GetParameter("perfMetric"))
        self.rebalancePeriodIndex   = int(self.GetParameter("rebalancePeriodIndex"))
        self.lookbackInDays         = int(self.GetParameter("lookbackInDays"))
        self.useETFWeights          = bool(int(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 = []
        self.queuedPositions    = []

        # Values related to our alpha (rebalancing period and preferred perf metric 
        self.PerfMetrics      = [MetricsEnum.MOMENTUM_PCT, MetricsEnum.MOMP_KER ]
        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 efficiency ratio, take top 1/3rd 
            # ----------------------------------------------------
            # numHoldings = int(len(self.symDataDict)/3) 
            # sortedByEfficencyRatio = sorted(self.symDataDict.items(), key=lambda x: x[1].EfficiencyRatio, reverse=True)[:self.numHoldings]
            # Sort by perfMetric
            
            # debug
            perf_values = [data.PerfMetricValue if data.PerfMetricValue != float('-inf') else 0.0 for _, data in self.symDataDict.items()]

            self.StoreResearchData(perf_values)            
            filtered_data = [(symbol, data) for symbol, data in self.symDataDict.items() if data.PerfMetricValue != float('-inf')]
            sortedByPerfMetric = sorted(filtered_data, key=lambda x: x[1].PerfMetricValue, reverse=True)[:self.maxHoldings]
            
            # debug
            # chosenSymbolMetrics = [(x[0].Value,x[1].PerfMetricValue)  for x in sortedByPerfMetric]

            # Get Symbol object (the key of dict)
            self.SelectedSymbols    = [x[0] for x in sortedByPerfMetric]
            self.SelectedSymbolData = [x[1] for x in sortedByPerfMetric]
            
            
            #### TODO: 
            #### Remove Symboldata from self.symDataDict that didnt make the cut
            # self.CleanUpSymbolData(self.SelectedSymbols)
            
            # self.Plot("Symbols Found", 'Line', (len(self.SelectedSymbols)))	
            # self.Debug(f"!!!Symbols Found: {self.Time.strftime('%y-%m-%d')}{[x.Value for x in self.SelectedSymbols]}")
            

            return self.SelectedSymbols 

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

    # ====================================
    def LiquidatePositions(self):
        
        for symbol,symbolData in self.symDataDict.items():
            
            if(self.Securities.ContainsKey(symbol) and (symbol in self.Portfolio) and self.Portfolio[symbol].Invested):
                symbolData.OnPositionClosed()

            self.RemoveSecurity(symbol.Value) 
        
        # self.liquidationEvents = 0 if (not hasattr(self, "liquidationEvents")) else self.liquidationEvents +1
        # self.Plot("Liquidation Event", 'Line', self.liquidationEvents)	
        self.Liquidate()

        


    # ====================================
    def OnData(self,dataSlice):
        
        # Debug: Plot how many holdings we have
        # -------------------------------------
        # numHoldings = len([x.Key for x in self.Portfolio if x.Value.Invested])
        # self.Plot("Num of Positions", 'Line', numHoldings)	
        
        # Process queued positions from previous cycle
        # --------------------------------------------
        self.ProcessQueuedPositions()
        

        # Send new slice data to each symboldata object
        # ---------------------------------------------- 
        for symbol,symbolData in self.symDataDict.items():
            if(symbol in dataSlice) and (self.Portfolio[symbol].Invested):

                # If it's a daily bar....
                if(symbol in dataSlice and dataSlice[symbol] is not None and dataSlice[symbol].Period == timedelta(days=1)):                    
                    # Calll the daily data handler in symbolData
                    symbolData.OnSymbolDailyData(dataSlice[symbol])

                # else If it's an hourly bar....
                elif(symbol in dataSlice and dataSlice[symbol] is not None and dataSlice[symbol].Period == timedelta(seconds=3600)):
                    symbolData.OnSymbolHourlyData(dataSlice[symbol])

                    # Check trailing stop, queue for liqudation if exit signal fired
                    if symbolData.TrailingExitSignalFired():
                        self.QueuePosition(symbol, 0)
                
        
        # Handle Universe repopulation
        # ----------------------------
        if (self.universeRepopulated):
            # self.Liquidate()
            self.LiquidatePositions()
            self.weights = {}

            
            if(self.useETFWeights):
                # calculate sum of the etf weights for each selected symbol
                weightsSum = sum(self.symDataDict[symbol].etfWeight if self.symDataDict[symbol].etfWeight is not None else 0 for symbol in self.SelectedSymbols)
            
                ## Debug: Check how many of these have 0 etf weight
                # zero_count = sum(1 for symbol in self.SelectedSymbols if self.symDataDict[symbol].etfWeight is None)
                # total_count = len(self.SelectedSymbols)
                # percentage = (zero_count / total_count) * 100


            # self.Debug(f"!!!Symbols to trade: {self.Time.strftime('%y-%m-%d')}{[x.Value for x in self.SelectedSymbols]}")
            # # self.Debug(f"{[x.etfWeight for x in [symdata for sym,symdata in self.symDataDict]]}")
            # # self.Debug(f"{[x.etfWeight for sym, x in self.symDataDict.items()]}")
            # self.Debug(f"{[f'{sym.Value} {x.etfWeight}' for sym, x in self.symDataDict.items() if sym in self.SelectedSymbols]}")

            for symbol in self.SelectedSymbols:

                # Either use ETF Weights or equal weights
                if(self.useETFWeights and weightsSum > 0):
                    try:
                        symbolWeight = self.symDataDict[symbol].etfWeight / weightsSum # respect weighting
                    except:
                        symbolWeight = 1 / len(self.SelectedSymbols) # Equally weighted    
                else:
                    symbolWeight = 1 / len(self.SelectedSymbols) # Equally weighted


                if(symbol in self.CurrentSlice):
                    if True or self.Securities[symbol].IsTradable:
                        self.AddEquity(symbol.Value, Resolution.Hour)
                        self.QueuePosition(symbol,symbolWeight)



                    # 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
    
    def QueuePosition(self, symbol, symbolWeight):
        # if we havent already queued this position with a weight, 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 we are atttempting to queue a symbol that already has a 
        # trade queued with weight 0 (meaning an exit before rebalancing)
        # then remove the previous entry
        self.queuedPositions = [(sym, weight) for sym, weight in self.queuedPositions if not (sym == symbol and symbolWeight == 0)]

        if symbol not in [x[0] for x in self.queuedPositions]:
            self.queuedPositions.append((symbol, symbolWeight))




    ## Loop through symboldata and remove what we dont need
    ## -------------------------------------------------
    def CleanUpSymbolData(self, selectedSymbols):
        # delete symbolData we dont need any more. 
        symbolsToRemove = [symbol for symbol in self.symDataDict 
                            if self.Securities.ContainsKey(symbol) 
                            and not self.Portfolio[symbol].Invested
                            and symbol not in selectedSymbols]

        for symbol in symbolsToRemove:
            del self.symDataDict[symbol]
            
    ## Loop through queued positions and open new trades  
    ## -------------------------------------------------
    def ProcessQueuedPositions(self):

        if len(self.queuedPositions) == 0:
            return

        tradeEvents = 0
        # Iterate through a copy of self.queuedPositions, 
        # since we modify self.queuedPositions in this loop.
        # ---------------------------------------------------
        for symbol,weight in [*self.queuedPositions]:

            if self.CurrentSlice.ContainsKey(symbol) and self.CurrentSlice[symbol] is not None:                    
                                
                openOrders = self.Transactions.GetOpenOrders()
                pendingPositionExists = any(order.Symbol == symbol for order in openOrders)

                if not pendingPositionExists:
                    orderNote = f"{weight:3.2f} of {symbol.Value}"
                    
                    if weight == 0:
                        orderNote = f"Possible Trailing Exit for {symbol.Value}" 

                    self.Log(f"{weight:3.2f} of {symbol.Value}")
                    adjustedWeight = weight * self.maxExposurePct
                    self.SetHoldings(symbol,weight, tag=f"{orderNote}")
                    tradeEvents = tradeEvents + 1

                    if weight > 0:
                        self.symDataDict[symbol].OnPositionOpened()

                # Remove occurences of symbol from queue
                self.queuedPositions = [t for t in self.queuedPositions if t[0] != symbol]

        
        self.Plot("Trade Dequeues", 'Line', tradeEvents)	


            
            # else:
            #     self.Log(f"[ERROR] Could not process {symbol} in queue: No Data in current slice")            
    
    # =================================================================
    # 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')


    # This is used in the research notebook to plot values.
    def StoreResearchData(self, perf_values):

        # save the data to object store for late retrieval
        # Convert the dictionary to a JSON string
        # this is 
        json_data = json.dumps(perf_values)
        save_successful = self.ObjectStore.Save("my_key", json_data)