from AlgorithmImports import *

class MomentumAssetAllocationStrategy(QCAlgorithm):

    def Initialize(self):
        # Set the start date and initial capital for the backtest
        self.SetStartDate(2014, 1, 1)  # Start the backtest from January 1, 2014
        self.SetCash(100000)  # Initial capital of $100,000

        # Dictionary to store Rate of Change (ROC) indicators and SMAs (50-day and 200-day) for each ETF
        self.data = {}  # To store ROC (momentum) indicators
        self.sma50 = {}  # To store 50-day Simple Moving Average for regime filter
        self.sma200 = {}  # To store 200-day Simple Moving Average for regime filter
        period = 12 * 21  # Lookback period for ROC (12 months * 21 trading days per month = 252 days)

        # Warm-up period to gather enough historical data before making trading decisions
        self.SetWarmUp(period, Resolution.Daily)  # Ensure enough data is collected for the 12-month ROC calculation

        # Number of top ETFs to select based on momentum (we select the top 3)
        self.traded_count = 3
        
        # List of 10 sector ETFs (representing different sectors of the market)
        self.symbols = ["XLB", "XLE", "XLF", "XLI", "XLK", "XLP", "XLU", "XLV", "XLY", "VNQ"]

        # Initialize the indicators (ROC, SMA50, and SMA200) for each ETF
        for symbol in self.symbols:
            self.AddEquity(symbol, Resolution.Minute)  # Add equity data for each ETF at minute resolution
            self.data[symbol] = self.ROC(symbol, period, Resolution.Daily)  # 12-month ROC for momentum
            self.sma50[symbol] = self.SMA(symbol, 50, Resolution.Daily)  # 50-day SMA for regime filter
            self.sma200[symbol] = self.SMA(symbol, 200, Resolution.Daily)  # 200-day SMA for regime filter
        
        # Variable to track the last rebalance month to ensure the strategy rebalances only once per month
        self.recent_month = -1  # Initialized to an invalid month to force the first rebalance

    def OnData(self, data):
        # Exit if the algorithm is still warming up (collecting historical data for the indicators)
        if self.IsWarmingUp:
            return

        # Ensure that trading decisions are made after the market opens (at 9:31 AM)
        if not (self.Time.hour == 9 and self.Time.minute == 31):
            return  # Avoid trading at the market open to reduce the effect of early-day volatility

        self.Log(f"Market Open Time: {self.Time}")

        # Rebalance the portfolio once per month (check if the month has changed)
        if self.Time.month == self.recent_month:
            return  # Skip rebalancing if it's the same month as the last rebalance
        self.recent_month = self.Time.month  # Update the last rebalance month to the current month
        
        self.Log(f"New monthly rebalance...")

        # Prepare to select the top 3 ETFs based on momentum and the regime filter (SMA 50 > SMA 200)
        selected = {}

        # Loop through each ETF and check if the indicators (ROC, SMA50, and SMA200) are ready
        for symbol in self.symbols:
            if not data.ContainsKey(symbol):
                continue  # Skip the ETF if data for the symbol is not available
            
            roc = self.data[symbol]  # Retrieve the ROC (momentum) indicator for the ETF
            sma50 = self.sma50[symbol]  # Retrieve the 50-day SMA for the ETF
            sma200 = self.sma200[symbol]  # Retrieve the 200-day SMA for the ETF
            
            # Check if all indicators (ROC, SMA50, SMA200) have enough data and are ready
            if roc.IsReady and sma50.IsReady and sma200.IsReady:
                # Apply the regime filter: Only include ETFs where the 50-day SMA > 200-day SMA (bullish trend)
                if sma50.Current.Value > sma200.Current.Value:
                    selected[symbol] = roc  # Select the ETF if it passes the regime filter
                    self.Log(f"{symbol} passes the regime filter: 50-SMA > 200-SMA")
                else:
                    self.Log(f"{symbol} fails the regime filter: 50-SMA <= 200-SMA")

        # Sort the selected ETFs by their momentum (ROC) in descending order (highest momentum first)
        sorted_by_momentum = sorted(selected.items(), key=lambda x: x[1].Current.Value, reverse=True)
        self.Log(f"Number of assets passing the regime filter: {len(sorted_by_momentum)}")

        # List to store the selected top 3 ETFs based on momentum
        long = []
        
        # Only proceed if there are enough ETFs that passed the regime filter
        if len(sorted_by_momentum) >= self.traded_count:
            long = [x[0] for x in sorted_by_momentum][:self.traded_count]  # Select the top 3 momentum ETFs

        # Get the list of currently invested ETFs in the portfolio
        invested = [x.Key.Value for x in self.Portfolio if x.Value.Invested]
        
        # Liquidate any ETFs that are no longer part of the top 3 momentum performers
        for symbol in invested:
            if symbol not in long:
                self.Liquidate(symbol)  # Exit the position for ETFs that are no longer in the top 3

        self.Log(f"Selected long leg for next month: {long}")
        
        # **Dynamic Weighting**: Allocate capital proportionally to the strength of each ETF's momentum
        total_momentum = sum([x[1].Current.Value for x in sorted_by_momentum[:self.traded_count]])  # Sum of top 3 momentum values
        if total_momentum == 0:
            total_momentum = 1  # Avoid division by zero if momentum values are zero for all selected assets
        
        # Allocate portfolio weights based on each ETF's relative momentum strength
        for symbol, roc in sorted_by_momentum[:self.traded_count]:
            weight = roc.Current.Value / total_momentum  # Calculate weight based on momentum
            self.SetHoldings(symbol, weight)  # Set holdings for the ETF based on the calculated weight
            self.Log(f"Allocated {weight * 100:.2f}% to {symbol} based on momentum strength.")