"""
hedged_rf — Hedged Random Forest for Time-Series Forecasting
=============================================================

A Python implementation of the Hedged Random Forest (HRF) 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 optimizing non-equal (and even
negative) weights for individual trees, minimizing the mean-squared forecast error.

Key components
--------------
- EWMA estimation of the tree-error mean vector and covariance matrix
- Linear shrinkage toward structured targets (CVC for covariance, constant-mean for μ)
- Numba-accelerated parallel computation of HAC variance estimates
- Convex optimization (CVXPY / OSQP) for the gross-exposure-constrained weight problem

Quick start
-----------
>>> from sklearn.ensemble import RandomForestRegressor
>>> from hedged_rf import fit_hrf, extract_tree_predictions, predict_hrf
>>>
>>> rf = RandomForestRegressor(n_estimators=500, random_state=42)
>>> rf.fit(X_train, y_train)
>>>
>>> # Build residual matrix R  (T × n_trees)
>>> R = extract_tree_predictions(rf, X_train, y_train)
>>>
>>> # Estimate HRF weights
>>> w = fit_hrf(R, lambda_param=0.15, H=6, kappa=2.0, verbose=True)
>>>
>>> # Generate forecasts
>>> y_pred = predict_hrf(rf, X_test, w)

Public API
----------
fit_hrf                   Main entry point — returns optimal weight vector
predict_hrf               Generate predictions using HRF weights
extract_tree_predictions  Helper to build the (T × N) residual matrix from a fitted RF
estimate_mu_complete      Stand-alone EWMA + shrinkage mean estimator
estimate_sigma_complete   Stand-alone EWMA + shrinkage covariance estimator
solve_hrf_optimization    Stand-alone convex solver for the weight problem
"""

from .core import (
    fit_hrf,
    predict_hrf,
    estimate_mu_ewma,
    estimate_mu_shrinkage,
    estimate_mu_complete,
    estimate_sigma_ewma,
    estimate_sigma_shrinkage,
    estimate_sigma_complete,
    create_cvc_target,
    solve_hrf_optimization,
    estimate_autocovariance,
    estimate_autocovariance_bivariate,
    estimate_nu_i,
    estimate_nu_ij,
    ENABLE_NUMBA,
)

import numpy as np
from typing import Tuple


def extract_tree_predictions(
    rf,
    X: np.ndarray,
    y: np.ndarray,
) -> np.ndarray:
    """
    Build the (T × N) residual matrix from a fitted scikit-learn RandomForest.

    Parameters
    ----------
    rf : sklearn.ensemble.RandomForestRegressor
        A *fitted* random forest regressor.
    X : array-like of shape (T, d)
        Feature matrix used during training.
    y : array-like of shape (T,)
        Target vector used during training.

    Returns
    -------
    R : np.ndarray of shape (T, N)
        Residual matrix where R[t, j] = y[t] - tree_j.predict(X[t]).
    """
    X = np.asarray(X)
    y = np.asarray(y).ravel()
    predictions = np.column_stack([tree.predict(X) for tree in rf.estimators_])
    return y[:, None] - predictions


__all__ = [
    "fit_hrf",
    "predict_hrf",
    "extract_tree_predictions",
    "estimate_mu_ewma",
    "estimate_mu_shrinkage",
    "estimate_mu_complete",
    "estimate_sigma_ewma",
    "estimate_sigma_shrinkage",
    "estimate_sigma_complete",
    "create_cvc_target",
    "solve_hrf_optimization",
    "estimate_autocovariance",
    "estimate_autocovariance_bivariate",
    "estimate_nu_i",
    "estimate_nu_ij",
    "ENABLE_NUMBA",
]

__version__ = "1.0.0"
__author__ = "Ezequiel Grillo"
