from AlgorithmImports import *
import numpy as np
import math

class SystematicPortfolioRotation(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 1, 1)
        self.SetCash(100000)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)

        # Stage 1: Universe Selection
        self.risk_on_tickers = ["SPY", "QQQ", "USMV", "EFAV", "EEMV", "GLD", "TLT"]
        self.defensive_allocation = {"SHY": 0.50, "BIL": 0.25, "GLD": 0.25}
        
        self.symbols = []
        # Add Equities to the universe
        for ticker in self.risk_on_tickers + list(self.defensive_allocation.keys()):
            symbol = self.AddEquity(ticker, Resolution.Daily).Symbol
            if symbol not in self.symbols:
                self.symbols.append(symbol)

        # Proxies for Dual Momentum Regime Filter
        self.market_proxy = self.Symbol("SPY")
        self.risk_free_proxy = self.Symbol("BIL")

        self.lookback_days = 252 # Approx. 12 months of trading days

        # Stage 4 Setup: Quarterly Rebalancing
        self.rebalance_months = [1, 4, 7, 10]
        self.Schedule.On(
            self.DateRules.MonthStart("SPY"),
            self.TimeRules.AfterMarketOpen("SPY", 30),
            self.ExecuteRebalance
        )

    def ExecuteRebalance(self):
        # Enforce quarterly schedule
        if self.Time.month not in self.rebalance_months:
            return

        # Fetch historical daily close prices
        history = self.History(self.symbols, self.lookback_days, Resolution.Daily)
        if history.empty:
            return

        # Calculate 12-month momentum (returns)
        returns = {}
        for symbol in self.symbols:
            if symbol in history.index.levels[0]:
                prices = history.loc[symbol]['close']
                if not prices.empty and len(prices) > 0:
                    returns[symbol] = (prices.iloc[-1] / prices.iloc[0]) - 1

        if self.market_proxy not in returns or self.risk_free_proxy not in returns:
            return

        # Stage 2: Regime Detection (Absolute Momentum Proxy)
        market_return = returns[self.market_proxy]
        rf_return = returns[self.risk_free_proxy]
        
        # If Market > T-Bills, we are in a Bullish/Risk-On state
        is_risk_on = market_return > rf_return

        target_weights = {}

        # Stage 3: Allocation
        if not is_risk_on:
            self.Debug(f"{self.Time}: Bearish Regime Detected. Shifting to Defensive Basket.")
            for ticker, weight in self.defensive_allocation.items():
                target_weights[self.Symbol(ticker)] = weight
        else:
            self.Debug(f"{self.Time}: Bullish Regime Detected. Applying Median Selection.")
            ro_returns = {s: returns[s] for s in returns if s.Value in self.risk_on_tickers}
            
            # Sort risk-on assets by 12-month return
            sorted_assets = sorted(ro_returns.items(), key=lambda x: x[1])
            
            # The Median Anomaly: Exclude extreme winners/losers, select the middle
            n = len(sorted_assets)
            mid_idx = n // 2
            
            # Select the middle 3 assets to represent the median cluster
            median_symbols = [sorted_assets[i][0] for i in range(mid_idx - 1, mid_idx + 2)]
            
            # Equal weight the median assets (Proxy for Hierarchical Risk Parity)
            weight_per_asset = 1.0 / len(median_symbols)
            for symbol in median_symbols:
                target_weights[symbol] = weight_per_asset

        # Stage 4: Execution with Turnover Control (5% Threshold)
        for symbol, target_weight in target_weights.items():
            current_weight = 0
            if self.Portfolio.TotalPortfolioValue > 0:
                current_weight = self.Portfolio[symbol].HoldingsValue / self.Portfolio.TotalPortfolioValue
            
            # Only trade if the required allocation change is greater than 5%
            if abs(current_weight - target_weight) >= 0.05:
                self.SetHoldings(symbol, target_weight)

        # Liquidate assets no longer in the target portfolio
        for holding in self.Portfolio.Values:
            if holding.Symbol not in target_weights and holding.Invested:
                self.Liquidate(holding.Symbol)