from AlgorithmImports import *


class ETFMomentumCorrelationHedgeAlgorithm(QCAlgorithm):
    """ETF Asset Momentum with a Correlation-Filtered Selective Short Hedge.

    Monthly strategy: rank 13 cross-asset ETFs by average 3/6/9/12-month
    momentum, hold top-4 long at 25% each. When the 20-day average pairwise
    correlation exceeds the 250-day average pairwise correlation (momentum-
    favorable regime), add a 30% short of the worst-ranked ETF.
    """

    def initialize(self) -> None:
        self.set_start_date(2007, 7, 1)
        self.set_end_date(2026, 6, 29)
        self.set_cash(100_000)
        self.settings.automatic_indicator_warm_up = True
        self.settings.seed_initial_prices = True
        # Define some parameters.
        mom_period_months = [3, 6, 9, 12]
        self._top_n = 4
        self._short_weight = 0
        self._corr_short_window = 20
        self._corr_long_window = 250
        # Add the ETFs and their momentum indicators.
        for ticker in ["SPY", "IWM", "EFA", "EEM", "IYR", "QQQ", "LQD", "IEF", "TIP", "GLD", "USO", "DBC", "FXE"]:
            equity = self.add_equity(ticker, Resolution.DAILY)
            equity.rocps = [self.rocp(equity.symbol, period*21, Resolution.DAILY) for period in mom_period_months]
        # Add a Scheduled Event to rebalance the portfolio monthly.
        self.schedule.on(self.date_rules.month_start("SPY"), self.time_rules.at(8, 0), self._rebalance)

    def _rebalance(self) -> None:
        # Ensure all momentum indicators are warmed up.
        for security in self.securities.values():
            if not all(rocp.is_ready for rocp in security.rocps):
                return
        # Calculate the momentum score of each stock.
        # Momentum score = average of 3/6/9/12-month total returns
        score_by_security = {
            security: sum(rocp.current.value for rocp in security.rocps) / len(security.rocps)
            for security in self.securities.values()
        }
        # Select the stocks to long and short.
        sorted_by_score = sorted(score_by_security, key=lambda s: score_by_security[s])
        long_symbols = sorted_by_score[-self._top_n:]
        short_candidate = sorted_by_score[0]
        # Calculate the correlation regime. Momentum is favorable when short avg > long avg.
        avg_corr_short, avg_corr_long = self._compute_correlations()
        if avg_corr_short is None or avg_corr_long is None:
            return
        momentum_favorable = avg_corr_short > avg_corr_long
        # Create portfolio targets. 
        # Always long with equal-weight. Short the worst ETF at 30% if favorable.
        targets = [PortfolioTarget(s, 1/self._top_n) for s in long_symbols]
        if momentum_favorable:
            targets.append(PortfolioTarget(short_candidate, -self._short_weight))
        # Place the trades to rebalance the portfolio.
        self.set_holdings(targets, liquidate_existing_holdings=True)

    def _compute_correlations(self) -> tuple:
        """Return (avg_20d_corr, avg_250d_corr) of pairwise daily-return correlations."""
        # Get the trailing price history.
        history = self.history(self.securities.keys(), self._corr_long_window + 1, Resolution.DAILY)
        if history.empty:
            return None, None
        # Calculate the daily returns.
        returns = history["close"].unstack(level=0).pct_change().dropna(how="all")
        if len(returns) < self._corr_long_window:
            return None, None
        # Calcualte the average pair-wise correlations across long and short windows.
        avg_corr_short = self._avg_pairwise_correlation(returns.tail(self._corr_short_window))
        avg_corr_long = self._avg_pairwise_correlation(returns.tail(self._corr_long_window))
        if np.isnan(avg_corr_short) or np.isnan(avg_corr_long):
            return None, None
        return avg_corr_short, avg_corr_long

    def _avg_pairwise_correlation(self, returns_df: pd.DataFrame) -> float:
        """Mean of the upper-triangle pairwise Pearson correlations."""
        corr_matrix = returns_df.corr()
        n = corr_matrix.shape[0]
        mask = np.triu(np.ones((n, n), dtype=bool), k=1)
        values = corr_matrix.values[mask]
        return float(np.nanmean(values))