"""
Hedged Random Forest (HRF) - Production-Ready Implementation

This module implements the Hedged Random Forest method as described in:
Beck, E., & Wolf, M. (2025). "Forecasting inflation with the hedged random forest."
SNB Working Papers 07/2025, Swiss National Bank.

The HRF improves upon the standard Random Forest by using optimized non-equal weights
(including negative weights) for individual trees, minimizing the mean-squared error
of forecasts.

Key Features:
    - EWMA (Exponentially Weighted Moving Average) estimation
    - Linear shrinkage for robust covariance estimation
    - Numba-accelerated parallel computation
    - Convex optimization via CVXPY

Author: [Ezequiel Grillo]
Version: 1.0.0
Date: January 2026
"""

import numpy as np
import cvxpy as cp
from scipy import linalg
from numba import jit, prange
from typing import Tuple, Optional
import warnings

warnings.filterwarnings('ignore')

# ==============================================================================
# GLOBAL CONFIGURATION
# ==============================================================================

ENABLE_NUMBA = True
"""bool: Global flag to enable/disable Numba JIT compilation."""


# ==============================================================================
# NUMBA-ACCELERATED CORE FUNCTIONS
# ==============================================================================

@jit(nopython=True, cache=True)
def _autocovariance_numba(series: np.ndarray, mean: float, max_lag: int) -> np.ndarray:
    """
    Compute autocovariance function using Numba JIT compilation.
    
    Calculates γ(h) = (1/T) Σₜ (xₜ - μ)(xₜ₊ₕ - μ) for h = 0, 1, ..., max_lag.
    
    Args:
        series: Time series data (1D array).
        mean: Pre-computed mean of the series.
        max_lag: Maximum lag to compute.
        
    Returns:
        Array of autocovariances [γ(0), γ(1), ..., γ(max_lag)].
    """
    T = len(series)
    acov = np.zeros(max_lag + 1)
    centered = series - mean
    
    # Variance (lag 0)
    for t in range(T):
        acov[0] += centered[t] ** 2
    acov[0] /= T
    
    # Autocovariances (lag > 0)
    for h in range(1, max_lag + 1):
        for t in range(T - h):
            acov[h] += centered[t] * centered[t + h]
        acov[h] /= T
    
    return acov


@jit(nopython=True, cache=True)
def _autocovariance_bivariate_numba(series_i: np.ndarray, series_j: np.ndarray,
                                     mean_i: float, mean_j: float, 
                                     max_lag: int) -> np.ndarray:
    """
    Compute autocovariance of the product of two centered series.
    
    Calculates the autocovariance function of (xᵢₜ - μᵢ)(xⱼₜ - μⱼ).
    
    Args:
        series_i: First time series.
        series_j: Second time series.
        mean_i: Mean of first series.
        mean_j: Mean of second series.
        max_lag: Maximum lag to compute.
        
    Returns:
        Array of autocovariances for the product series.
    """
    T = len(series_i)
    product = (series_i - mean_i) * (series_j - mean_j)
    
    # Compute mean manually (Numba-compatible)
    mean_product = 0.0
    for t in range(T):
        mean_product += product[t]
    mean_product /= T
    
    psi = np.zeros(max_lag + 1)
    centered_product = product - mean_product
    
    # Variance
    for t in range(T):
        psi[0] += (product[t] - mean_product) ** 2
    psi[0] /= T
    
    # Autocovariances
    for h in range(1, max_lag + 1):
        for t in range(T - h):
            psi[h] += centered_product[t] * centered_product[t + h]
        psi[h] /= T
    
    return psi


@jit(nopython=True, cache=True)
def _compute_column_means(R: np.ndarray) -> np.ndarray:
    """
    Compute column-wise means of a matrix.
    
    Args:
        R: Input matrix (T × N).
        
    Returns:
        Array of column means (N,).
    """
    T, N = R.shape
    means = np.empty(N)
    for j in range(N):
        means[j] = np.mean(R[:, j])
    return means


@jit(nopython=True, parallel=True, cache=True)
def _compute_all_nu_ij_numba(R: np.ndarray, lambda_param: float, 
                             H: int) -> np.ndarray:
    """
    Compute variance estimates νᵢⱼ for all pairs (i,j) in parallel.
    
    Calculates νᵢⱼ = Var(σ̂ᵢⱼ) using EWMA-based estimation with HAC standard errors.
    Formula: νᵢⱼ = [λ²/(1-(1-λ)²)] · [ψ₀ + 2Σₕ₌₁ᴴ (1-λ)ʰ ψₕ]
    
    Args:
        R: Residual matrix (T × N).
        lambda_param: EWMA decay parameter λ ∈ (0,1).
        H: Bandwidth for autocovariance estimation.
        
    Returns:
        Matrix of variance estimates (N × N).
        
    Note:
        This function is parallelized over the outer loop (i index).
        Each thread writes to independent memory locations.
    """
    T, N = R.shape
    nu_matrix = np.zeros((N, N))
    
    factor = lambda_param ** 2 / (1 - (1 - lambda_param) ** 2)
    
    # Pre-compute EWMA weights
    weights = np.empty(H)
    for h in range(H):
        weights[h] = (1 - lambda_param) ** (h + 1)
    
    # Center all series once
    col_means = _compute_column_means(R)
    R_centered = np.empty_like(R)
    for j in range(N):
        R_centered[:, j] = R[:, j] - col_means[j]
    
    # Parallel computation over pairs (i,j)
    for i in prange(N):
        for j in range(N):
            product = R_centered[:, i] * R_centered[:, j]
            mean_prod = np.mean(product)
            centered_prod = product - mean_prod
            
            # Compute autocovariance at lag 0
            psi_0 = 0.0
            for t in range(T):
                psi_0 += (product[t] - mean_prod) ** 2
            psi_0 /= T
            
            sum_term = psi_0
            
            # Add weighted autocovariances
            for h in range(1, H + 1):
                psi_h = 0.0
                for t in range(T - h):
                    psi_h += centered_prod[t] * centered_prod[t + h]
                psi_h /= T
                sum_term += 2 * weights[h - 1] * psi_h
            
            nu_matrix[i, j] = factor * sum_term
    
    return nu_matrix


@jit(nopython=True, parallel=True, cache=True)
def _compute_all_nu_i_numba(R: np.ndarray, lambda_param: float, 
                            H: int) -> float:
    """
    Compute variance estimates νᵢ for all series i in parallel.
    
    Calculates νᵢ = Var(μ̂ᵢ) and returns their sum Σᵢ νᵢ.
    Formula: νᵢ = [λ²/(1-(1-λ)²)] · [ψ₀ + 2Σₕ₌₁ᴴ (1-λ)ʰ ψₕ]
    
    Args:
        R: Residual matrix (T × N).
        lambda_param: EWMA decay parameter λ ∈ (0,1).
        H: Bandwidth for autocovariance estimation.
        
    Returns:
        Sum of variance estimates Σᵢ νᵢ.
        
    Note:
        Mean is computed manually within parallel loop to avoid race conditions.
    """
    T, N = R.shape
    factor = lambda_param ** 2 / (1 - (1 - lambda_param) ** 2)
    nu_array = np.empty(N)
    
    for i in prange(N):
        series = R[:, i]
        
        # Compute mean manually (thread-safe)
        mean = 0.0
        for t in range(T):
            mean += series[t]
        mean /= T
        
        # Compute autocovariances
        psi = _autocovariance_numba(series, mean, H)
        
        # Compute variance estimate
        sum_term = psi[0]
        for h in range(1, H + 1):
            sum_term += 2 * ((1 - lambda_param) ** h) * psi[h]
        
        nu_array[i] = factor * sum_term
    
    return nu_array.sum()


# ==============================================================================
# AUTOCOVARIANCE ESTIMATION
# ==============================================================================

def estimate_autocovariance(series: np.ndarray, max_lag: int) -> np.ndarray:
    """
    Estimate autocovariance function with optional Numba acceleration.
    
    Args:
        series: Time series data.
        max_lag: Maximum lag to compute.
        
    Returns:
        Array of autocovariances [γ(0), γ(1), ..., γ(max_lag)].
    """
    if ENABLE_NUMBA:
        mean = series.mean()
        return _autocovariance_numba(series, mean, max_lag)
    else:
        T = len(series)
        mean = series.mean()
        acov = np.zeros(max_lag + 1)
        for h in range(max_lag + 1):
            if h == 0:
                acov[h] = np.mean((series - mean) ** 2)
            else:
                acov[h] = np.mean((series[:-h] - mean) * (series[h:] - mean))
        return acov


def estimate_autocovariance_bivariate(series_i: np.ndarray, series_j: np.ndarray,
                                      max_lag: int) -> np.ndarray:
    """
    Estimate autocovariance of product of two series with optional Numba acceleration.
    
    Args:
        series_i: First time series.
        series_j: Second time series.
        max_lag: Maximum lag to compute.
        
    Returns:
        Array of autocovariances for the product series.
    """
    if ENABLE_NUMBA:
        mean_i = series_i.mean()
        mean_j = series_j.mean()
        return _autocovariance_bivariate_numba(series_i, series_j, mean_i, mean_j, max_lag)
    else:
        product = series_i * series_j
        mean_product = product.mean()
        psi = np.zeros(max_lag + 1)
        for h in range(max_lag + 1):
            if h == 0:
                psi[h] = np.var(product, ddof=0)
            else:
                psi[h] = np.mean((product[:-h] - mean_product) * (product[h:] - mean_product))
        return psi


# ==============================================================================
# MEAN VECTOR ESTIMATION
# ==============================================================================

def estimate_mu_ewma(R: np.ndarray, lambda_param: float = 0.15) -> np.ndarray:
    """
    Estimate mean vector using EWMA (Exponentially Weighted Moving Average).
    
    Computes μ̂ = Σₜ wₜ xₜ where wₜ = λ(1-λ)^(T-t) (normalized).
    
    Args:
        R: Residual matrix (T × N).
        lambda_param: EWMA decay parameter. Default: 0.15 (for monthly data).
            - Smaller λ → more weight on distant past
            - Larger λ → more weight on recent observations
            
    Returns:
        EWMA mean vector (N,).
    """
    T, N = R.shape
    
    # Compute exponential weights
    t = np.arange(1, T + 1)
    weights = lambda_param * (1 - lambda_param) ** (T - t)
    weights /= weights.sum()
    
    # Weighted average
    mu_hat = weights @ R
    
    return mu_hat


def estimate_nu_i(series: np.ndarray, lambda_param: float = 0.15, 
                  H: int = 6) -> float:
    """
    Estimate variance of EWMA mean estimator for a single series.
    
    Computes νᵢ = Var(μ̂ᵢ) using HAC standard errors.
    
    Args:
        series: Time series data.
        lambda_param: EWMA decay parameter.
        H: Bandwidth for autocovariance estimation.
        
    Returns:
        Variance estimate νᵢ.
    """
    psi = estimate_autocovariance(series, max_lag=H)
    
    factor = lambda_param ** 2 / (1 - (1 - lambda_param) ** 2)
    
    sum_term = psi[0]
    for h in range(1, H + 1):
        sum_term += 2 * ((1 - lambda_param) ** h) * psi[h]
    
    return factor * sum_term


def estimate_mu_shrinkage(R: np.ndarray, mu_hat_ewma: np.ndarray, 
                          lambda_param: float = 0.15, H: int = 6,
                          use_fast: bool = True) -> Tuple[np.ndarray, float]:
    """
    Apply linear shrinkage to EWMA mean estimator.
    
    Computes μ† = α·μ* + (1-α)·μ̂ where:
        - μ* = constant mean target (grand mean)
        - α = ν/(ν+γ) is the shrinkage intensity
        - ν = estimation variance
        - γ = distance to target
    
    Args:
        R: Residual matrix (T × N).
        mu_hat_ewma: EWMA mean estimate.
        lambda_param: EWMA decay parameter.
        H: Bandwidth for autocovariance estimation.
        use_fast: Use Numba-accelerated parallel computation.
        
    Returns:
        Tuple of (shrunk mean vector, shrinkage intensity α).
    """
    T, N = R.shape
    mu_star = mu_hat_ewma.mean()
    
    # Estimate variance
    if use_fast and ENABLE_NUMBA:
        nu = _compute_all_nu_i_numba(R, lambda_param, H)
    else:
        nu = sum(estimate_nu_i(R[:, i], lambda_param, H) for i in range(N))
    
    # Compute distance to target
    gamma = np.sum((mu_star - mu_hat_ewma) ** 2)
    
    # Shrinkage intensity (clipped to [0,1])
    alpha = np.clip(nu / (nu + gamma), 0, 1)
    
    # Linear combination
    mu_shrunk = alpha * mu_star + (1 - alpha) * mu_hat_ewma
    
    return mu_shrunk, alpha


def estimate_mu_complete(R: np.ndarray, lambda_param: float = 0.15, 
                        H: int = 6, verbose: bool = False,
                        use_fast: bool = True) -> np.ndarray:
    """
    Complete pipeline for mean vector estimation.
    
    Performs EWMA estimation followed by linear shrinkage.
    
    Args:
        R: Residual matrix (T × N).
        lambda_param: EWMA decay parameter.
        H: Bandwidth for autocovariance estimation.
        verbose: Print estimation details.
        use_fast: Use Numba-accelerated computation.
        
    Returns:
        Final mean vector estimate.
    """
    mu_ewma = estimate_mu_ewma(R, lambda_param)
    mu_final, alpha = estimate_mu_shrinkage(R, mu_ewma, lambda_param, H, use_fast)
    
    if verbose:
        print(f"  μ estimation:")
        print(f"    EWMA mean: {mu_ewma.mean():.6e}")
        print(f"    Shrinkage α: {alpha:.4f}")
        print(f"    Final mean: {mu_final.mean():.6e}")
    
    return mu_final


# ==============================================================================
# COVARIANCE MATRIX ESTIMATION
# ==============================================================================

def estimate_sigma_ewma(R: np.ndarray, lambda_param: float = 0.15) -> np.ndarray:
    """
    Estimate covariance matrix using EWMA.
    
    Computes Σ̂ = Σₜ wₜ (xₜ - x̄)(xₜ - x̄)' where wₜ are exponential weights.
    
    Args:
        R: Residual matrix (T × N).
        lambda_param: EWMA decay parameter.
        
    Returns:
        EWMA covariance matrix (N × N).
    """
    T, N = R.shape
    
    # Compute exponential weights
    t = np.arange(1, T + 1)
    weights = lambda_param * (1 - lambda_param) ** (T - t)
    weights /= weights.sum()
    
    # Center data
    x_bar = R.mean(axis=0)
    R_centered = R - x_bar
    
    # Weighted covariance
    R_weighted = R_centered * np.sqrt(weights[:, np.newaxis])
    Sigma_hat = R_weighted.T @ R_weighted
    
    return Sigma_hat


def create_cvc_target(Sigma_hat: np.ndarray) -> np.ndarray:
    """
    Create Constant Variance-Covariance (CVC) shrinkage target.
    
    Constructs a structured target matrix F with:
        - Constant diagonal: f₁ = average variance
        - Constant off-diagonal: f₂ = average covariance
    
    Args:
        Sigma_hat: Sample covariance matrix (N × N).
        
    Returns:
        CVC target matrix F (N × N).
    """
    N = Sigma_hat.shape[0]
    
    # Average variance
    f1 = np.trace(Sigma_hat) / N
    
    # Average covariance
    mask = np.triu(np.ones((N, N), dtype=bool), k=1)
    f2 = Sigma_hat[mask].sum() / (N * (N - 1) / 2)
    
    # Construct target
    F = np.full((N, N), f2)
    np.fill_diagonal(F, f1)
    
    return F


def estimate_nu_ij(R_i: np.ndarray, R_j: np.ndarray, 
                   lambda_param: float = 0.15, H: int = 6) -> float:
    """
    Estimate variance of EWMA covariance estimator for a pair of series.
    
    Computes νᵢⱼ = Var(σ̂ᵢⱼ) using HAC standard errors.
    
    Args:
        R_i: First time series.
        R_j: Second time series.
        lambda_param: EWMA decay parameter.
        H: Bandwidth for autocovariance estimation.
        
    Returns:
        Variance estimate νᵢⱼ.
    """
    y_i = R_i - R_i.mean()
    y_j = R_j - R_j.mean()
    
    psi = estimate_autocovariance_bivariate(y_i, y_j, max_lag=H)
    
    factor = lambda_param ** 2 / (1 - (1 - lambda_param) ** 2)
    
    sum_term = psi[0]
    for h in range(1, H + 1):
        sum_term += 2 * ((1 - lambda_param) ** h) * psi[h]
    
    return factor * sum_term


def estimate_sigma_shrinkage(R: np.ndarray, Sigma_hat: np.ndarray, F: np.ndarray,
                             lambda_param: float = 0.15, H: int = 6,
                             verbose: bool = False, 
                             use_fast: bool = True) -> Tuple[np.ndarray, float]:
    """
    Apply linear shrinkage to EWMA covariance estimator.
    
    Computes Σ† = α·F + (1-α)·Σ̂ where:
        - F = CVC target matrix
        - α = ν/(ν+γ) is the shrinkage intensity
        - ν = total estimation variance
        - γ = total distance to target
    
    Args:
        R: Residual matrix (T × N).
        Sigma_hat: EWMA covariance estimate.
        F: CVC target matrix.
        lambda_param: EWMA decay parameter.
        H: Bandwidth for autocovariance estimation.
        verbose: Print estimation details.
        use_fast: Use Numba-accelerated parallel computation.
        
    Returns:
        Tuple of (shrunk covariance matrix, shrinkage intensity α).
    """
    T, N = R.shape
    
    if verbose:
        print(f"    Computing shrinkage for {N}×{N} matrix...")
    
    # Estimate total variance
    if use_fast and ENABLE_NUMBA:
        nu_matrix = _compute_all_nu_ij_numba(R, lambda_param, H)
        nu = nu_matrix.sum()
    else:
        nu = 0
        for i in range(N):
            for j in range(N):
                nu += estimate_nu_ij(R[:, i], R[:, j], lambda_param, H)
    
    # Compute total distance to target
    gamma = np.sum((F - Sigma_hat) ** 2)
    
    # Shrinkage intensity (clipped to [0,1])
    alpha = np.clip(nu / (nu + gamma), 0, 1)
    
    # Linear combination
    Sigma_shrunk = alpha * F + (1 - alpha) * Sigma_hat
    
    if verbose:
        print(f"    ν: {nu:.6e}, γ: {gamma:.6e}, α: {alpha:.4f}")
    
    return Sigma_shrunk, alpha


def estimate_sigma_complete(R: np.ndarray, lambda_param: float = 0.15,
                           H: int = 6, verbose: bool = False,
                           use_fast: bool = True) -> np.ndarray:
    """
    Complete pipeline for covariance matrix estimation.
    
    Performs EWMA estimation followed by linear shrinkage to CVC target.
    
    Args:
        R: Residual matrix (T × N).
        lambda_param: EWMA decay parameter.
        H: Bandwidth for autocovariance estimation.
        verbose: Print estimation details.
        use_fast: Use Numba-accelerated computation.
        
    Returns:
        Final covariance matrix estimate.
    """
    Sigma_ewma = estimate_sigma_ewma(R, lambda_param)
    F = create_cvc_target(Sigma_ewma)
    Sigma_final, alpha = estimate_sigma_shrinkage(
        R, Sigma_ewma, F, lambda_param, H, verbose, use_fast=use_fast
    )
    
    if verbose:
        print(f"  Σ estimation:")
        print(f"    EWMA condition number: {np.linalg.cond(Sigma_ewma):.2f}")
        print(f"    Shrinkage α: {alpha:.4f}")
        print(f"    Final condition number: {np.linalg.cond(Sigma_final):.2f}")
        
        eigvals = linalg.eigvalsh(Sigma_final)
        print(f"    Min eigenvalue: {eigvals.min():.6e}")
        status = "✓ Positive definite" if eigvals.min() > 0 else "⚠️ Not positive definite"
        print(f"    {status}")
    
    return Sigma_final


# ==============================================================================
# HRF OPTIMIZATION
# ==============================================================================

def solve_hrf_optimization(mu: np.ndarray, Sigma: np.ndarray, 
                          kappa: float = 2.0, 
                          verbose: bool = False) -> np.ndarray:
    """
    Solve HRF weight optimization problem using convex programming.
    
    Solves:
        min_w  (w'μ)² + w'Σw
        s.t.   w'1 = 1
               ||w||₁ ≤ κ
    
    Args:
        mu: Mean vector (N,).
        Sigma: Covariance matrix (N × N).
        kappa: L₁ constraint parameter (gross-exposure limit).
            - kappa = 1.0: only non-negative weights allowed
            - kappa = 2.0: recommended default
            - kappa → ∞: no constraint on negative weights
        verbose: Print optimization details.
        
    Returns:
        Optimal weight vector (N,).
        
    Raises:
        RuntimeError: If optimization fails to converge.
    """
    N = len(mu)
    Q = Sigma + np.outer(mu, mu)
    
    # Define optimization problem
    w = cp.Variable(N)
    objective = cp.Minimize(cp.quad_form(w, Q))
    constraints = [
        cp.sum(w) == 1.0,
        cp.norm(w, 1) <= kappa
    ]
    
    problem = cp.Problem(objective, constraints)
    problem.solve(solver=cp.OSQP, verbose=False)
    
    if verbose:
        print(f"  HRF Optimization:")
        print(f"    Status: {problem.status}")
        print(f"    Objective value: {problem.value:.6e}")
    
    if problem.status not in [cp.OPTIMAL, cp.OPTIMAL_INACCURATE]:
        print(f"⚠️ Warning: Optimization status = {problem.status}")
    
    w_opt = w.value
    
    if verbose:
        sum_w = np.sum(w_opt)
        l1_norm = np.sum(np.abs(w_opt))
        n_negative = np.sum(w_opt < 0)
        print(f"    Σw = {sum_w:.10f} (should be 1.0)")
        print(f"    ||w||₁ = {l1_norm:.6f} (should be ≤ {kappa})")
        print(f"    Negative weights: {n_negative} ({100 * n_negative / N:.1f}%)")
    
    return w_opt


# ==============================================================================
# MAIN PIPELINE
# ==============================================================================

def fit_hrf(R: np.ndarray, 
            lambda_param: float = 0.15,
            H: int = 6,
            kappa: float = 2.0,
            verbose: bool = False,
            use_fast: bool = True) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Fit Hedged Random Forest model to residual matrix.
    
    This is the main entry point for HRF estimation. It performs:
        1. EWMA + shrinkage estimation of mean vector μ
        2. EWMA + shrinkage estimation of covariance matrix Σ
        3. Convex optimization to find optimal weights w
    
    Args:
        R: Residual matrix (T × N) where:
            - T = number of time observations
            - N = number of trees in the random forest
            - R[t, j] = residual of tree j at time t
        lambda_param: EWMA decay parameter ∈ (0,1).
            Recommended values:
            - 0.06 for daily data
            - 0.15 for monthly data (default)
            - 0.25 for quarterly data
        H: Bandwidth for autocovariance estimation.
            Typical range: [3, 12]. Default: 6.
        kappa: Gross-exposure constraint for L₁ norm.
            - 1.0: only non-negative weights
            - 2.0: recommended default (allows moderate negative weights)
            - Higher values allow more extreme positions
        verbose: Print detailed estimation progress.
        use_fast: Enable Numba-accelerated parallel computation.
            Set to False only for debugging.
            
    Returns:
        w_hrf: Optimal weight vector (N,).
            
    Example:
        >>> # Train Random Forest and extract residuals
        >>> rf = RandomForestRegressor(n_estimators=500)
        >>> rf.fit(X_train, y_train)
        >>> R = extract_residuals(rf, X_train, y_train)
        >>> 
        >>> # Fit HRF
        >>> w = fit_hrf(R, verbose=True)
        >>> 
        >>> # Make predictions
        >>> pred_hrf = predict_hrf(rf, X_test, w)
    """
    if verbose:
        mode = "OPTIMIZED (Numba parallel)" if use_fast else "ORIGINAL (sequential)"
        print(f"\nFitting HRF model:")
        print(f"  Residual matrix: {R.shape}")
        print(f"  Parameters: λ={lambda_param}, H={H}, κ={kappa}")
        print(f"  Mode: {mode}")
    
    # Estimate mean vector
    mu = estimate_mu_complete(R, lambda_param, H, verbose, use_fast)
    
    # Estimate covariance matrix
    Sigma = estimate_sigma_complete(R, lambda_param, H, verbose, use_fast)
    
    # Solve optimization problem
    w_hrf = solve_hrf_optimization(mu, Sigma, kappa, verbose)
    
    if verbose:
        print(f"✓ HRF model fitted successfully\n")
    
    return w_hrf


def predict_hrf(rf, X: np.ndarray, w: np.ndarray) -> np.ndarray:
    """
    Generate predictions using HRF weights.

    Args:
        rf: Fitted scikit-learn RandomForestRegressor.
        X: Feature matrix (n_samples, n_features).
        w: Optimal weight vector from fit_hrf (n_trees,).

    Returns:
        Predicted values (n_samples,).
    """
    tree_preds = np.column_stack([tree.predict(X) for tree in rf.estimators_])
    return tree_preds @ w


