#region imports
from AlgorithmImports import *
#endregion

class DynamicEtfMomentumStrategy(QCAlgorithm):
    '''
    This algorithm implements a dynamic momentum-based strategy for a portfolio of ETFs.
    It incorporates relative strength for selection and absolute momentum for crash protection.
    The logic is as follows:
    1. On a monthly schedule, the algorithm calculates the 6-month performance of all ETFs in the universe.
    2. It ranks the ETFs by this performance and selects the top 3 performers.
    3. CRASH PROTECTION: It checks if the top performer's 6-month momentum is positive.
       - If YES (market is generally rising), it allocates the portfolio equally among the top 3 ETFs.
       - If NO (market is likely falling), it allocates 100% of the portfolio to a safe-haven
         asset (long-term US Treasury bonds, TLT) to protect capital.
    4. This process repeats monthly, ensuring the portfolio is always positioned in the strongest
       assets or protected in cash during downturns.
    '''
    def Initialize(self):
        """
        Initializes the algorithm, sets up the universe, and schedules the rebalancing logic.
        """
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2025, 1, 21)
        self.SetCash(40000)

        # 1. EXPANDED UNIVERSE: Add more diverse ETFs for better rotation opportunities.
        # User's list + broad market, international, and safe-haven assets.
        self.risk_on_symbols_str = ["CHAT", "FBCG", "FBOT", "FBTC", "FDIG", "FMET", "NUKZ", "QTUM",
                                    "SPY",  # S&P 500
                                    "QQQ",  # Nasdaq 100
                                    "VEA",  # Developed Markets ex-US
                                    "GLD"]  # Gold
        
        # Define the safe-haven asset for crash protection
        self.safe_haven_symbol_str = "TLT" # 20+ Year Treasury Bond ETF

        # Add data for all symbols
        self.all_symbols = [self.AddEquity(s, Resolution.Daily).Symbol for s in self.risk_on_symbols_str]
        self.safe_haven_symbol = self.AddEquity(self.safe_haven_symbol_str, Resolution.Daily).Symbol
        self.all_symbols.append(self.safe_haven_symbol)

        # --- Strategy Parameters ---
        self.lookback_period = 6 * 21 # Approx. 6 months in trading days
        self.top_n_etfs = 3           # Number of top ETFs to hold

        # Set a warm-up period to ensure we have enough data for our lookback calculation
        self.SetWarmUp(self.lookback_period)

        # Schedule the rebalancing logic to run once a month
        self.Schedule.On(self.DateRules.MonthStart(), self.TimeRules.AfterMarketOpen(self.risk_on_symbols_str[0], 5), self.MonthlyRebalance)
        
    def MonthlyRebalance(self):
        """
        This method is called once per month to execute the core strategy logic.
        """
        if self.IsWarmingUp:
            return

        # 2. ROBUST STRATEGY: Calculate momentum for each risk-on ETF
        momentum = {}
        history = self.History(self.all_symbols, self.lookback_period + 1, Resolution.Daily)
        if history.empty:
            return

        for symbol in self.risk_on_symbols_str:
            # Ensure we have history for the symbol
            if symbol not in history.index.levels[0]:
                continue
            
            # Get the price series for the symbol
            close_prices = history.loc[symbol]['close']
            
            # Calculate the percentage return over the lookback period.
            # We add a small value to the denominator to avoid division-by-zero errors.
            if len(close_prices) > self.lookback_period:
                momentum[symbol] = (close_prices[-1] / close_prices[0]) - 1

        if not momentum:
            return

        # Rank the ETFs by momentum in descending order
        ranked_etfs = sorted(momentum.items(), key=lambda x: x[1], reverse=True)

        # Select the top N performers
        top_performers = ranked_etfs[:self.top_n_etfs]
        self.Log(f"Top performers this month: {[etf[0] for etf in top_performers]}")
        
        # 3. PORTFOLIO BALANCING & CRASH PROTECTION
        # Get the momentum of the #1 ranked ETF to use as our market filter
        top_etf_symbol, top_etf_momentum = top_performers[0]
        
        # Absolute Momentum Rule (Crash Protection)
        if top_etf_momentum > 0:
            # If momentum is positive, invest in the top performers
            self.Log("Momentum is positive. Investing in top ETFs.")
            
            # Calculate equal weight for each of the top performers
            weight = 1.0 / self.top_n_etfs
            
            # Create a list of portfolio targets
            targets = [PortfolioTarget(etf[0], weight) for etf in top_performers]
            
            # Use SetHoldings to rebalance the portfolio to these new targets.
            # This will automatically liquidate any old positions that are no longer in the top N.
            self.SetHoldings(targets)

        else:
            # If momentum is negative, a market downturn is likely.
            # Liquidate all risk-on assets and move 100% to the safe-haven asset.
            self.Log("Momentum is negative. Moving to safe-haven asset (TLT).")
            self.Liquidate() # First, sell any existing risk-on positions
            self.SetHoldings(self.safe_haven_symbol, 1.0) # Then, allocate everything to TLT


    def OnData(self, data: Slice):
        """
        OnData is not used in this strategy, as all logic is handled
        in the scheduled MonthlyRebalance method.
        """
        pass