# Core numerical computing
import numpy as np
import pandas as pd
import warnings
warnings.filterwarnings('ignore')

# For financial data
import yfinance as yf       # Loading financial data
import empyrical as emp     # Performance metrics

# Book data (pip install "git+https://github.com/dppalomar/pob.git#subdirectory=python")
from pob_python import SP500_stocks_2015to2020, SP500_index_2015to2020

# Plotting
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import seaborn as sns
sns.set_theme(style="darkgrid")
import plotly.io as pio
pio.renderers.default = "notebook"
pio.templates.default = "plotly_white"

# Optimization
import cvxpy as cp                  # interface for convex optimization solvers
from qpsolvers import solve_qp

def backtest(
        portfolio_funcs: dict[str, callable],
        dataset: dict[str, pd.DataFrame],
        lookback: int,
        rebalance_every: int = 1,
        optimize_every: int = 1,
        cost_bps: float = 0,
        verbose: bool = True
) -> tuple[pd.DataFrame, dict[str, pd.DataFrame]]:
    """
    Backtest portfolio strategies over historical data.

    At each time step from ``lookback`` to ``T``, the function:

    1. Re-optimizes portfolio weights every ``optimize_every`` periods by
       calling the user-supplied portfolio function.
    2. Rebalances to the latest optimized weights every ``rebalance_every``
       periods, computing turnover-based transaction costs.
    3. Computes the portfolio return for the period and drifts the weights
       according to realised asset returns.

    Parameters
    ----------
    portfolio_funcs : dict[str, callable]
        Dictionary mapping strategy names to portfolio weight functions.
        Each function receives a dict with keys ``'prices'``
        (``pd.DataFrame``) and ``'index'`` (``pd.DataFrame | None``) for the
        current lookback window, and must return a 1-D ``np.ndarray`` of
        portfolio weights.
    dataset : dict[str, pd.DataFrame]
        Must contain the key ``'prices'`` (T × N asset price DataFrame).
        May optionally contain ``'index'`` (T × 1 benchmark price DataFrame).
    lookback : int
        Number of past periods fed to each portfolio function for estimation.
    rebalance_every : int, default 1
        Number of periods between rebalancing to the target weights.
    optimize_every : int, default 1
        Number of periods between re-optimizing the target weights.
        Must be a multiple of ``rebalance_every``.
    cost_bps : float, default 0
        One-way transaction cost in basis points applied to turnover.
    verbose : bool, default True
        If True, print progress and re-optimization messages.

    Returns
    -------
    portf_rets_df : pd.DataFrame
        DataFrame of portfolio returns (one column per strategy), indexed
        from ``lookback`` to ``T - 1``.
    ws : dict[str, pd.DataFrame]
        Dictionary mapping each strategy name to a DataFrame of portfolio
        weights over the same period.

    Raises
    ------
    ValueError
        If ``optimize_every`` is not a multiple of ``rebalance_every``.
    """
    if optimize_every % rebalance_every != 0:
        raise ValueError(
            f"optimize_every ({optimize_every}) must be a multiple of "
            f"rebalance_every ({rebalance_every})."
        )

    stock_prices = dataset['prices']
    index_prices = dataset.get('index')
    T, N = stock_prices.shape
    portf_rets = {}
    ws = {}

    for portfolio_name, portfolio_func in portfolio_funcs.items():
        if verbose:
            print(f"Backtesting portfolio '{portfolio_name}'...")

        current_w = np.zeros(N)
        target_w = np.zeros(N)
        w = pd.DataFrame(index=stock_prices.index, columns=stock_prices.columns)
        portf_ret = pd.Series(index=stock_prices.index, dtype=float)
        portf_ret.name = portfolio_name

        for t in range(lookback, T):
            if (t - lookback) % rebalance_every == 0:
                # Re-optimize target weights if necessary
                if (t - lookback) % optimize_every == 0:
                    target_w = portfolio_func({
                        'prices': stock_prices.iloc[t - lookback:t],
                        'index': index_prices.iloc[t - lookback:t] if index_prices is not None else None
                    })
                    if verbose:
                        print(f"  Reoptimizing at t={t} "
                              f"(cardinality={np.sum(target_w > 1e-5)})")

                # Rebalance to target weights
                turnover = np.abs(target_w - current_w).sum()
                current_w = target_w.copy()
                transaction_cost = turnover * cost_bps / 10_000
            else:
                transaction_cost = 0.0

            # Store weights
            w.iloc[t] = current_w

            # Period return
            Xt = stock_prices.iloc[t] / stock_prices.iloc[t - 1] - 1
            portf_ret.iloc[t] = (Xt * current_w).sum() - transaction_cost

            # Drift weights with asset returns
            current_w = current_w * (1 + Xt)
            w_sum = current_w.sum()
            if abs(w_sum) > 1e-12:
                current_w = current_w / w_sum

        portf_rets[portfolio_name] = portf_ret[lookback:]
        ws[portfolio_name] = w[lookback:]

    portf_rets_df = pd.DataFrame(portf_rets)
    return portf_rets_df, ws

# Download S&P 500 index and SPY ETF data
SP500_SPY_prices = yf.download(["^GSPC", "SPY"], start="2007-01-01")['Close']
SP500_SPY_prices.columns = ['S&P 500', 'SPY']

# Plot
SP500_SPY_prices.plot(figsize=(12, 6), logy=True, linewidth=1.0)
plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
plt.gca().xaxis.set_major_locator(mdates.YearLocator())
plt.gca().xaxis.set_minor_locator(mdates.MonthLocator())
plt.title("Tracking of S&P 500 index")
plt.ylabel("dollars")
plt.xlabel(None)
plt.legend(title=None)
plt.tight_layout()
plt.show()

def sparse_tracking_via_MM(
        X: pd.DataFrame,
        r: pd.Series,
        lambda_val: float,
        w0: np.ndarray | None = None,
        max_iter: int = 5,
        thres: float = 1e-9,
        p: float = 1e-3,
        u: float = 1,
        reg: float = 1e-7,
        rtol: float = 1e-3,
        return_obj: bool = False,
        verbose: bool = False,
) -> np.ndarray | tuple[np.ndarray, list[float]]:
    """
    Sparse index tracking via the MM (majorization-minimization) algorithm.

    Solves the non-convex problem

        min_w  (1/T)||r - Xw||^2 + lambda * sum log(1 + |w_i|/p) / log(1 + u/p)
        s.t.   w >= 0,  1'w = 1

    by iteratively majorizing the log penalty with a weighted-L1 surrogate and
    solving the resulting QP subproblem.

    Parameters
    ----------
    X : pd.DataFrame
        T x N DataFrame of asset returns.
    r : pd.Series
        Length-T Series of benchmark (index) returns.
    lambda_val : float
        Regularization parameter controlling sparsity.
    w0 : np.ndarray or None, default None
        Initial weight vector of length N.  If None, uses 1/N uniform.
    max_iter : int, default 5
        Maximum number of MM iterations.
    thres : float, default 1e-9
        Weights below this value are zeroed after convergence.
    p : float, default 1e-3
        Shape parameter of the log penalty (controls approximation to L0).
    u : float, default 1
        Scale parameter of the log penalty.
    reg : float, default 1e-7
        Tikhonov regularization added to the QP cost matrix, expressed as a
        fraction of the mean diagonal of X'X.  Doubled automatically on
        QP solver failure.
    rtol : float, default 1e-3
        Relative tolerance on the weight vector for early stopping.
    return_obj : bool, default False
        If True, also return the sequence of objective values.
    verbose : bool, default False
        If True, print progress messages.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N summing to 1.
    obj_values : list[float]
        Objective value at each iteration.  Only returned when
        ``return_obj=True``.
    """
    X_mat = X.values
    r_vec = r.values.reshape(-1, 1)
    T, N = X_mat.shape

    # Precompute Gram matrix and cross term
    XtX = X_mat.T @ X_mat
    Xtr = X_mat.T @ r_vec  # (N, 1)

    w = np.ones(N) / N if w0 is None else w0.copy()
    log_denom = np.log(1 + u / p)

    def _objective(w):
        res = r_vec - X_mat @ w.reshape(-1, 1)
        tracking = (1 / T) * np.dot(res.ravel(), res.ravel())
        penalty = lambda_val * np.sum(np.log(1 + np.abs(w) / p) / log_denom)
        return tracking + penalty

    obj_values = [_objective(w)] if return_obj else []

    # QP components that don't change across iterations
    P = (2 / T) * XtX
    P = 0.5 * (P + P.T)  # ensure exact symmetry
    A_eq = np.ones((1, N))
    b_eq = np.array([1.0])
    G_ub = -np.eye(N)
    h_ub = np.zeros(N)

    current_reg = reg

    for k in range(max_iter):
        # MM surrogate weights
        d = 1.0 / (log_denom * (p + np.abs(w)))  # (N,)
        w_prev = w.copy()

        # # Solve the convex subproblem using CVXPY
        # w_var = cp.Variable(N)
        # objective = cp.Minimize((1/T) * cp.sum_squares(r_vec - X_mat @ w_var) +
        #                        lambda_val * cp.sum(cp.multiply(d, cp.abs(w_var))))
        # constraints = [cp.sum(w_var) == 1, w_var >= 0]
        # problem = cp.Problem(objective, constraints)

        # Build and solve the QP subproblem
        P_reg = P + current_reg * np.mean(np.diag(P)) * np.eye(N)
        q = (-(2 / T) * Xtr).ravel() + lambda_val * d  # must be 1-D

        w_new = solve_qp(P_reg, q, G_ub, h_ub, A_eq, b_eq, solver="quadprog")

        if w_new is None:
            if verbose:
                print(f"  Iter {k}: QP solver failed, doubling regularization.")
            current_reg *= 2
            continue

        w = w_new

        if return_obj:
            obj_values.append(_objective(w))

        if verbose:
            print(f"  Iter {k}: ||dw||/||w|| = "
                  f"{np.linalg.norm(w - w_prev) / max(np.linalg.norm(w_prev), 1e-12):.2e}")

        if np.linalg.norm(w - w_prev) / max(np.linalg.norm(w_prev), 1e-12) < rtol:
            break

    # Threshold and renormalize
    w[w < thres] = 0.0
    w_sum = w.sum()
    if w_sum > 0:
        w /= w_sum

    if return_obj:
        return w, obj_values
    return w

X = SP500_stocks_2015to2020.iloc[:, :10].pct_change().dropna()
r = SP500_index_2015to2020.pct_change().dropna()

sparse_tracking_via_MM(X.iloc[:, :10], r, lambda_val=5e-7)

def compute_alpha(
        X: pd.DataFrame,
        r: pd.Series,
) -> dict[str, pd.DataFrame | np.ndarray]:
    """
    Compute cumulative return ratios and transformed returns.

    Given asset returns X (T x N) and benchmark returns r (T x 1), define

    .. math::

        \\alpha_{t,n} = \\prod_{s=1}^{t}(1 + X_{s,n}) \\;/\\;
                        \\prod_{s=1}^{t}(1 + r_s)

    and the transformed returns
    :math:`\\tilde X_{t,n} = X_{t,n} \\, \\alpha_{t-1,n}` with
    :math:`\\alpha_{0,n} = 1`.

    Parameters
    ----------
    X : pd.DataFrame
        T x N DataFrame of asset returns.
    r : pd.Series
        Length-T Series of benchmark returns.

    Returns
    -------
    dict[str, pd.DataFrame | np.ndarray]
        ``'alpha'`` : T x N cumulative return ratios.
        ``'alpha_lagged'`` : T x N lagged cumulative return ratios
        (first row = 1).
        ``'last_alpha_lagged'`` : length-N array of final-row lagged alpha.
        ``'X_tilde'`` : T x N transformed returns.
    """
    X_cum = (1 + X).cumprod()
    r_cum = (1 + r).cumprod()

    alpha = X_cum.div(r_cum, axis=0)

    alpha_lagged = alpha.shift(1)
    alpha_lagged.iloc[0] = 1.0

    X_tilde = X * alpha_lagged

    return {
        'alpha': alpha,
        'alpha_lagged': alpha_lagged,
        'last_alpha_lagged': alpha_lagged.iloc[-1].values,
        'X_tilde': X_tilde,
    }

def _prepare_returns(
        dataset: dict[str, pd.DataFrame],
) -> tuple[pd.DataFrame, pd.Series]:
    """
    Compute asset and benchmark returns, aligned on common valid dates.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.

    Returns
    -------
    X : pd.DataFrame
        T x N asset returns (NaN rows dropped).
    r : pd.Series
        Length-T benchmark returns.
    """
    X = dataset['prices'].pct_change().dropna()
    r = dataset['index'].pct_change().dropna().squeeze()
    common_idx = X.index.intersection(r.index)
    return X.loc[common_idx], r.loc[common_idx]


def tracking_using_plain_returns_wfixed(
        dataset: dict[str, pd.DataFrame],
        lambda_val: float = 5e-7,
) -> np.ndarray:
    """
    Index tracking with fixed portfolio weights.

    Solves the sparse tracking problem on plain asset returns,
    yielding a weight vector that is held constant over time.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    lambda_val : float, default 5e-7
        Sparsity regularization parameter.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N.
    """
    X, r = _prepare_returns(dataset)
    return sparse_tracking_via_MM(X, r, lambda_val=lambda_val, max_iter=10)


def tracking_using_plain_returns_wt(
        dataset: dict[str, pd.DataFrame],
        lambda_val: float = 5e-7,
) -> np.ndarray:
    """
    Index tracking with time-varying portfolio weights.

    Optimizes over alpha-transformed returns so that the resulting
    portfolio drifts naturally with the index. The initial weights
    ``w0`` are scaled by the final lagged alpha and renormalized.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    lambda_val : float, default 5e-7
        Sparsity regularization parameter.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N.
    """
    X, r = _prepare_returns(dataset)
    alpha = compute_alpha(X, r)

    w0 = sparse_tracking_via_MM(alpha['X_tilde'], r, lambda_val=lambda_val, max_iter=10)
    w = w0 * alpha['last_alpha_lagged']

    w_sum = w.sum()
    if w_sum > 0:
        w /= w_sum
    return w

w = tracking_using_plain_returns_wfixed(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    })
sum(w), min(w)

w = tracking_using_plain_returns_wt(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    })
sum(w), min(w)

# Prep data
T, N = SP500_stocks_2015to2020.shape
T_trn = 2*252

# Run backtest
returns, weights = backtest(
    portfolio_funcs={
        "Tracking assuming fixed portfolio": tracking_using_plain_returns_wfixed,
        "Tracking assuming time-varying portfolio": tracking_using_plain_returns_wt
    },
    dataset={
        'prices': SP500_stocks_2015to2020,  #.loc["2018":]
        'index': SP500_index_2015to2020  # .loc["2018":]
    },
    lookback=T_trn,
    rebalance_every=6*21,  # Rebalance every 6 months (assuming 21 trading days per month)
    optimize_every=6*21,
)

def _align_with_benchmark(
        returns: pd.DataFrame,
        benchmark: pd.DataFrame,
) -> tuple[pd.DataFrame, pd.Series]:
    """Compute benchmark returns and align with portfolio returns."""
    benchmark_rets = benchmark.pct_change().dropna().squeeze()
    common_idx = returns.index.intersection(benchmark_rets.index)
    return returns.loc[common_idx], benchmark_rets.loc[common_idx]


def compute_tracking_error(
        returns: pd.DataFrame,
        benchmark: pd.DataFrame,
) -> pd.Series:
    """
    Compute mean squared tracking error for each portfolio.

    The tracking error is defined as
    :math:`(1/T) \\sum_t (r_{p,t} - r_{b,t})^2`, consistent with the
    least-squares objective used in ``sparse_tracking_via_MM``.

    Parameters
    ----------
    returns : pd.DataFrame
        T x K DataFrame of portfolio returns (one column per strategy).
    benchmark : pd.DataFrame
        T x 1 DataFrame of benchmark prices (converted to returns internally).

    Returns
    -------
    pd.Series
        Mean squared tracking error for each portfolio.
    """
    returns_aligned, benchmark_rets = _align_with_benchmark(returns, benchmark)
    excess = returns_aligned.subtract(benchmark_rets, axis=0)
    return (excess ** 2).mean()

def plot_tracking_error(
        returns: pd.DataFrame,
        benchmark: pd.DataFrame,
        window: int = 252,
) -> pd.DataFrame:
    """
    Plot rolling mean squared tracking error over time.

    Parameters
    ----------
    returns : pd.DataFrame
        T x K DataFrame of portfolio returns.
    benchmark : pd.DataFrame
        T x 1 DataFrame of benchmark prices.
    window : int, default 252
        Rolling window size (in trading days).

    Returns
    -------
    pd.DataFrame
        T x K DataFrame of squared tracking errors (unsmoothed).
    """
    returns_aligned, benchmark_rets = _align_with_benchmark(returns, benchmark)
    excess = returns_aligned.subtract(benchmark_rets, axis=0)
    tracking_errors = excess ** 2

    tracking_errors.rolling(window=window).mean().plot(figsize=(12, 6))
    plt.title(f'Tracking Error (Rolling {window}-day window)')
    plt.xlabel(None)
    plt.ylabel('Tracking Error')
    plt.legend(title='Portfolio')
    plt.grid(True)
    plt.tight_layout()
    plt.show()

    return tracking_errors

TE = plot_tracking_error(returns, benchmark=SP500_index_2015to2020, window=252)

def _compute_cardinality(
        weights: dict[str, pd.DataFrame],
        thres: float = 1e-5,
) -> pd.DataFrame:
    """
    Compute portfolio cardinality over time.

    Parameters
    ----------
    weights : dict[str, pd.DataFrame]
        Dictionary mapping strategy names to T x N weight DataFrames.
    thres : float, default 1e-5
        Weights below this threshold are treated as zero.

    Returns
    -------
    pd.DataFrame
        T x K DataFrame with the number of active assets per strategy.
    """
    return pd.DataFrame({
        name: (w > thres).sum(axis=1)
        for name, w in weights.items()
    })


def compute_cardinality(
        weights: dict[str, pd.DataFrame],
        thres: float = 1e-5,
) -> pd.Series:
    """
    Compute mean portfolio cardinality for each strategy.

    Parameters
    ----------
    weights : dict[str, pd.DataFrame]
        Dictionary mapping strategy names to T x N weight DataFrames.
    thres : float, default 1e-5
        Weights below this threshold are treated as zero.

    Returns
    -------
    pd.Series
        Mean cardinality per strategy.
    """
    return _compute_cardinality(weights, thres).mean()


def plot_cardinality(
        weights: dict[str, pd.DataFrame],
        thres: float = 1e-5,
) -> pd.DataFrame:
    """
    Plot portfolio cardinality over time.

    Parameters
    ----------
    weights : dict[str, pd.DataFrame]
        Dictionary mapping strategy names to T x N weight DataFrames.
    thres : float, default 1e-5
        Weights below this threshold are treated as zero.

    Returns
    -------
    pd.DataFrame
        T x K DataFrame with the number of active assets per strategy.
    """
    cardinality_df = _compute_cardinality(weights, thres)

    cardinality_df.plot(figsize=(12, 6))
    plt.title('Portfolio Cardinality Over Time')
    plt.xlabel(None)
    plt.ylabel('Number of Assets')
    plt.legend(title='Portfolio')
    plt.grid(True)
    plt.tight_layout()
    plt.show()

    return cardinality_df

cardinality = plot_cardinality(weights)

# Calculate column means for TE and cardinality
summary_df = pd.DataFrame({
    "TE": TE.mean(),
    "cardinality": cardinality.mean()
})

# Display
summary_df.style.format({
    'TE': '{:.2e}',  # Scientific notation with 6 decimal places
    'cardinality': '{:.2f}'  # Fixed-point with 2 decimal places
})

def tracking_using_cumreturns_wfixed(
        dataset: dict[str, pd.DataFrame],
        lambda_val: float = 20e-7,
) -> np.ndarray:
    """
    Index tracking with fixed weights on cumulative returns.

    Solves the sparse tracking problem on cumulative asset returns,
    yielding a weight vector that is held constant over time.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    lambda_val : float, default 20e-7
        Sparsity regularization parameter.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N.
    """
    X, r = _prepare_returns(dataset)
    return sparse_tracking_via_MM(X.cumsum(), r.cumsum(), lambda_val=lambda_val, max_iter=10)


def tracking_using_cumreturns_wt(
        dataset: dict[str, pd.DataFrame],
        lambda_val: float = 15e-7,
) -> np.ndarray:
    """
    Index tracking with time-varying weights on cumulative returns.

    Optimizes over cumulative alpha-transformed returns and rescales
    the resulting weights by the final lagged alpha.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    lambda_val : float, default 15e-7
        Sparsity regularization parameter.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N.
    """
    X, r = _prepare_returns(dataset)
    alpha = compute_alpha(X, r)

    w0 = sparse_tracking_via_MM(
        alpha['X_tilde'].cumsum(), r.cumsum(),
        lambda_val=lambda_val, max_iter=10,
    )

    w = w0 * alpha['last_alpha_lagged']
    w_sum = w.sum()
    if w_sum > 0:
        w /= w_sum
    return w

# Prep data
T, N = SP500_stocks_2015to2020.shape
T_trn = 2*252

# Run backtest
returns, weights = backtest(
    portfolio_funcs={
        "Tracking using plain returns": tracking_using_plain_returns_wfixed,
        "Tracking using cumulative returns": tracking_using_cumreturns_wfixed
    },
    dataset={
        'prices': SP500_stocks_2015to2020,  #.loc["2018":]
        'index': SP500_index_2015to2020  # .loc["2018":]
    },
    lookback=T_trn,
    rebalance_every=6*21,  # Rebalance every 6 months (assuming 21 trading days per month)
    optimize_every=6*21,
)

def plot_prices(
        returns: pd.DataFrame,
        benchmark: pd.DataFrame,
        title: str = 'Tracking of benchmark index',
) -> None:
    """
    Plot cumulative portfolio value against the benchmark.

    Portfolio returns are compounded into a NAV series and rescaled so
    that all curves start at the benchmark's initial price.

    Parameters
    ----------
    returns : pd.DataFrame
        T x K DataFrame of portfolio returns (one column per strategy).
    benchmark : pd.DataFrame
        T x 1 DataFrame of benchmark prices.
    title : str, default 'Tracking of benchmark index'
        Plot title.
    """
    NAV = (1 + returns).cumprod()
    benchmark_aligned = benchmark.loc[NAV.index]
    B0 = benchmark_aligned.iloc[0, 0]
    all_prices = pd.concat([benchmark_aligned, B0 * NAV], axis=1)

    all_prices.plot(figsize=(12, 6))
    plt.title(title)
    plt.xlabel(None)
    plt.ylabel('dollars')
    plt.legend(title='Portfolio')
    plt.grid(True)
    plt.tight_layout()
    plt.show()

plot_prices(returns, benchmark=SP500_index_2015to2020)

# Prep data
X = SP500_stocks_2015to2020.pct_change().dropna()
r = SP500_index_2015to2020.pct_change().dropna()

# Run index tracking MM algorithm
w, obj_values = sparse_tracking_via_MM(X, r, lambda_val=5e-7, max_iter=10, return_obj=True)

# Plot convergence
plt.figure(figsize=(10, 6))
plt.plot(range(len(obj_values)), obj_values, linewidth=1.2, color='blue')
plt.title('Convergence over iterations')
plt.xlabel('k')
plt.ylabel('Objective Value')
plt.xticks(range(len(obj_values)))
plt.grid(True)
plt.tight_layout()
plt.show()

def _solve_tracking_qp(
        X_mat: np.ndarray,
        r_vec: np.ndarray,
        reg: float = 1e-7,
) -> np.ndarray:
    """
    Solve the dense least-squares tracking QP.

    Parameters
    ----------
    X_mat : np.ndarray
        T x N matrix of asset returns.
    r_vec : np.ndarray
        T x 1 vector of benchmark returns.
    reg : float, default 1e-7
        Tikhonov regularization as a fraction of the mean diagonal of X'X.

    Returns
    -------
    w : np.ndarray
        Weight vector of length N.
    """
    T, N = X_mat.shape

    # # Solve the convex subproblem using CVXPY
    # N = X_mat.shape[1]
    # w = cp.Variable(N)
    # objective = cp.Minimize(cp.sum_squares(X_mat.values @ w - r_vec.values))
    # constraints = [cp.sum(w) == 1, w >= 0]

    # Solve the convex subproblem using quadprog
    P = (2 / T) * X_mat.T @ X_mat
    P = 0.5 * (P + P.T)
    P += reg * np.mean(np.diag(P)) * np.eye(N)
    q = (-(2 / T) * X_mat.T @ r_vec).ravel()

    w = solve_qp(P, q, -np.eye(N), np.zeros(N),
                 np.ones((1, N)), np.array([1.0]), solver="quadprog")
    return w


def full_tracking(
        dataset: dict[str, pd.DataFrame],
        reg: float = 1e-7,
) -> np.ndarray:
    """
    Full (dense) index tracking via least-squares QP.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    reg : float, default 1e-7
        Tikhonov regularization as a fraction of the mean diagonal of X'X.

    Returns
    -------
    w : np.ndarray
        Portfolio weight vector of length N.
    """
    X, r = _prepare_returns(dataset)
    return _solve_tracking_qp(X.values, r.values.reshape(-1, 1), reg=reg)

w = full_tracking(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    })
sum(w), min(w)

def two_stage_naive_tracking(
        dataset: dict[str, pd.DataFrame],
        K: int = 20,
        reg: float = 1e-7,
) -> np.ndarray:
    """
    Two-stage naive index tracking.

    Stage 1: solve the dense (full) least-squares tracking QP.
    Stage 2: keep only the K largest weights and renormalize.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    K : int, default 20
        Number of assets to retain in the sparse portfolio.
    reg : float, default 1e-7
        Tikhonov regularization as a fraction of the mean diagonal of X'X.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N with at most K non-zeros.
    """
    # Stage 1: dense portfolio
    w_dense = full_tracking(dataset, reg=reg)

    # Stage 2: keep K largest weights and renormalize
    N = len(w_dense)
    idx_keep = np.argsort(w_dense)[-K:]

    w = np.zeros(N)
    w[idx_keep] = w_dense[idx_keep]

    w_sum = w.sum()
    if w_sum > 0:
        w /= w_sum

    return w

w = two_stage_naive_tracking(
    dataset={
        'prices': SP500_stocks_2015to2020.loc['2018':],
        'index': SP500_index_2015to2020.loc['2018':],
    },
    K=10,
)

print(f"Sum: {w.sum():.6f}, Min: {w.min():.6f}, Cardinality: {np.sum(w > 1e-5)}")

def two_stage_refitted_tracking(
        dataset: dict[str, pd.DataFrame],
        K: int = 20,
        reg: float = 1e-7,
) -> np.ndarray:
    """
    Two-stage index tracking with refitted weights.

    Stage 1: solve the dense tracking QP.
    Stage 2: select the K largest weights, then re-solve the QP
    restricted to those K assets.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    K : int, default 20
        Number of assets to retain.
    reg : float, default 1e-7
        Tikhonov regularization as a fraction of the mean diagonal of X'X.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N with at most K non-zeros.
    """
    X, r = _prepare_returns(dataset)
    X_mat = X.values
    r_vec = r.values.reshape(-1, 1)
    N = X_mat.shape[1]

    # Stage 1: dense portfolio
    w_dense = _solve_tracking_qp(X_mat, r_vec, reg=reg)

    # Stage 2: refit on K largest
    idx_keep = np.argsort(w_dense)[-K:]
    w_refit = _solve_tracking_qp(X_mat[:, idx_keep], r_vec, reg=reg)

    w = np.zeros(N)
    w[idx_keep] = w_refit
    return w

w = two_stage_refitted_tracking(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    }, K=15)
sum(w), min(w), sum(w > 1e-5)

def iterative_reweighted_l1_norm_tracking(
        dataset: dict[str, pd.DataFrame],
        lambda_val: float = 5e-7,
) -> np.ndarray:
    """
    Sparse index tracking via iteratively reweighted L1-norm (MM algorithm).

    Equivalent to fixed-weight tracking on plain returns using the
    log-penalty MM solver in ``sparse_tracking_via_MM``.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    lambda_val : float, default 5e-7
        Sparsity regularization parameter.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N.
    """
    X, r = _prepare_returns(dataset)
    return sparse_tracking_via_MM(X, r, lambda_val=lambda_val, max_iter=10)

w = iterative_reweighted_l1_norm_tracking(dataset={
        'prices': SP500_stocks_2015to2020.loc["2018":],
        'index': SP500_index_2015to2020.loc["2018":]
    })
sum(w), min(w), sum(w > 1e-5)

# Prep data
T, N = SP500_stocks_2015to2020.shape
T_trn = 2*252

# Parameters for the desired sparsity
K = 20
lambda_val=5e-7

# Run backtest
returns, weights = backtest(
    portfolio_funcs={
        "Full tracking": full_tracking,
        "Two-step design with proportional weights": two_stage_naive_tracking,
        "Two-step design with refitted weights": two_stage_refitted_tracking,
        "Iterative reweighted l1-norm method": iterative_reweighted_l1_norm_tracking
    },
    dataset={
        'prices': SP500_stocks_2015to2020,  #.loc["2018":]
        'index': SP500_index_2015to2020  # .loc["2018":]
    },
    lookback=T_trn,
    rebalance_every=6*21,  # Rebalance every 6 months (assuming 21 trading days per month)
    optimize_every=6*21,
)

plot_prices(returns, benchmark=SP500_index_2015to2020)

methods = ["Two-step design with proportional weights",
           "Two-step design with refitted weights",
           "Iterative reweighted l1-norm method"]

TE = plot_tracking_error(returns[methods], benchmark=SP500_index_2015to2020, window=252)

cardinality = plot_cardinality({key: weights[key] for key in methods})

# Calculate column means for TE and cardinality
summary_df = pd.DataFrame({
    "TE": TE.mean(),
    "cardinality": cardinality.mean()
})

# Display
summary_df.style.format({
    'TE': '{:.2e}',  # Scientific notation with 6 decimal places
    'cardinality': '{:.2f}'  # Fixed-point with 2 decimal places
})

# Prep data
T, N = SP500_stocks_2015to2020.shape
T_trn = 2*252

# Sweeping loop
lmd_sweep = 1e-5 * np.array([1, 2, 5, 10, 20, 50, 100, 200, 1000]) / 300
df_list = []
for lmd in lmd_sweep:
    print(f"Backtesting for lambda: {lmd:.2e}")

    # Run backtest
    returns, weights = backtest(
        portfolio_funcs={
            "Assuming fixed portfolio": lambda dataset: tracking_using_plain_returns_wfixed(dataset, lambda_val=lmd),
            "Assuming time-varying portfolio": lambda dataset: tracking_using_plain_returns_wt(dataset, lambda_val=lmd)
        },
        dataset={
            'prices': SP500_stocks_2015to2020,
            'index': SP500_index_2015to2020
        },
        lookback=T_trn,
        rebalance_every=6*21,  # Rebalance every 6 months (assuming 21 trading days per month)
        optimize_every=6*21,
        verbose=False
    )

    # Calculate tracking error and cardinality
    TE = compute_tracking_error(returns, benchmark=SP500_index_2015to2020)
    cardinality = compute_cardinality(weights)
    df_list.append(pd.DataFrame({
        "method": TE.index,
        "TE": TE.values,
        "K": cardinality.values,
    }))

# Combine all rows into a single DataFrame
df = pd.concat(df_list, ignore_index=True)

# Plot tracking error versus sparsity level
sns.lineplot(data=df, x='K', y='TE', hue='method', style='method', linewidth=1.2)
plt.title('Tracking error versus sparsity level')
plt.xlabel('K (active assets)')
plt.ylabel('Tracking Error')
plt.legend(title=None)
plt.grid(True)
plt.tight_layout()
plt.show()

def _prepare_cumulative_returns(
        dataset: dict[str, pd.DataFrame],
) -> tuple[pd.DataFrame, pd.Series, pd.DataFrame, pd.Series]:
    """
    Compute plain and cumulative returns, aligned on common valid dates.

    Returns
    -------
    X, r, X_cum, r_cum
    """
    X, r = _prepare_returns(dataset)
    return X, r, X.cumsum(), r.cumsum()

def _tracking_residual(X_cum: pd.DataFrame, w: np.ndarray, r_cum: pd.Series) -> np.ndarray:
    """Compute tracking residual e_t = X_cum @ w - r_cum as a 1-D array."""
    return X_cum.values @ w - r_cum.values

def TE_index_tracking(
        dataset: dict[str, pd.DataFrame],
        lambda_val: float = 150 * 6e-7,
        u: float = 0.5,
        max_iter: int = 10,
) -> np.ndarray:
    """
    Sparse index tracking on cumulative returns (squared tracking error).

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    lambda_val : float, default 9e-5
        Sparsity regularization parameter.
    u : float, default 0.5
        Scale parameter of the log penalty.
    max_iter : int, default 10
        Maximum MM iterations.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N.
    """
    _, _, X_cum, r_cum = _prepare_cumulative_returns(dataset)
    return sparse_tracking_via_MM(X_cum, r_cum, lambda_val=lambda_val, u=u, max_iter=max_iter)

def DR_index_tracking(
        dataset: dict[str, pd.DataFrame],
        lambda_val: float = 150 * 6e-7,
        u: float = 0.5,
        max_iter_mm: int = 10,
        max_iter_outer: int = 5,
) -> np.ndarray:
    """
    Downside-risk index tracking on cumulative returns.

    Iteratively modifies the benchmark by adding the positive part of
    the tracking residual, so that only underperformance is penalized.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    lambda_val : float, default 9e-5
        Sparsity regularization parameter.
    u : float, default 0.5
        Scale parameter of the log penalty.
    max_iter_mm : int, default 10
        Maximum MM iterations per subproblem.
    max_iter_outer : int, default 5
        Number of outer reweighting iterations.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N.
    """
    _, _, X_cum, r_cum = _prepare_cumulative_returns(dataset)

    w = sparse_tracking_via_MM(X_cum, r_cum, lambda_val=lambda_val, u=u, max_iter=max_iter_mm)

    for _ in range(max_iter_outer):
        residual = _tracking_residual(X_cum, w, r_cum)  # (T,)
        r_cum_tilde = r_cum + np.maximum(0, residual)
        w = sparse_tracking_via_MM(X_cum, r_cum_tilde, lambda_val=lambda_val, u=u, max_iter=max_iter_mm)

    return w

def TE1_index_tracking(
        dataset: dict[str, pd.DataFrame],
        lambda_val: float = 150 * 6e-7,
        u: float = 0.5,
        max_iter_mm: int = 10,
        max_iter_outer: int = 5,
) -> np.ndarray:
    """
    L1-norm index tracking on cumulative returns via IRLS.

    Iteratively reweights observations by the inverse square root of
    the absolute tracking residual to approximate the L1-norm objective.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    lambda_val : float, default 9e-5
        Sparsity regularization parameter.
    u : float, default 0.5
        Scale parameter of the log penalty.
    max_iter_mm : int, default 10
        Maximum MM iterations per subproblem.
    max_iter_outer : int, default 5
        Number of outer IRLS iterations.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N.
    """
    _, _, X_cum, r_cum = _prepare_cumulative_returns(dataset)

    w = sparse_tracking_via_MM(X_cum, r_cum, lambda_val=100 * 6e-7, u=u, max_iter=max_iter_mm)

    for _ in range(max_iter_outer):
        residual = np.abs(_tracking_residual(X_cum, w, r_cum))  # (T,)
        alpha_sqrt = np.sqrt(1 / (2 * np.maximum(residual, 1e-12)))
        X_cum_w = X_cum.multiply(alpha_sqrt, axis=0)
        r_cum_w = r_cum * alpha_sqrt
        w = sparse_tracking_via_MM(X_cum_w, r_cum_w, lambda_val=lambda_val, u=u, max_iter=max_iter_mm)

    return w

def HubTE_index_tracking(
        dataset: dict[str, pd.DataFrame],
        lambda_val: float = 50 * 5e-8,
        u: float = 0.5,
        M: float = 1e-4,
        max_iter_mm: int = 10,
        max_iter_outer: int = 5,
) -> np.ndarray:
    """
    Huber-loss index tracking on cumulative returns via IRLS.

    Like ``TE1_index_tracking`` but uses a Huber-style weight that clips
    the reweighting at threshold ``M``, making it robust to large
    tracking deviations.

    Parameters
    ----------
    dataset : dict[str, pd.DataFrame]
        Must contain ``'prices'`` and ``'index'``.
    lambda_val : float, default 2.5e-6
        Sparsity regularization parameter.
    u : float, default 0.5
        Scale parameter of the log penalty.
    M : float, default 1e-4
        Huber threshold controlling the transition from L2 to L1.
    max_iter_mm : int, default 10
        Maximum MM iterations per subproblem.
    max_iter_outer : int, default 5
        Number of outer IRLS iterations.

    Returns
    -------
    w : np.ndarray
        Sparse portfolio weight vector of length N.
    """
    _, _, X_cum, r_cum = _prepare_cumulative_returns(dataset)

    w = sparse_tracking_via_MM(X_cum, r_cum, lambda_val=lambda_val, u=u, max_iter=max_iter_mm)

    for _ in range(max_iter_outer):
        residual = np.abs(_tracking_residual(X_cum, w, r_cum))  # (T,)
        alpha_sqrt = np.sqrt(np.minimum(1.0, M / np.maximum(residual, 1e-12)))
        X_cum_w = X_cum.multiply(alpha_sqrt, axis=0)
        r_cum_w = r_cum * alpha_sqrt
        w = sparse_tracking_via_MM(X_cum_w, r_cum_w, lambda_val=lambda_val, u=u, max_iter=max_iter_mm)

    return w

# Prep data
T, N = SP500_stocks_2015to2020.shape
T_trn = 2*252

# Run backtest
returns, weights = backtest(
    portfolio_funcs={
        "TE based design": TE_index_tracking,
        "Hub-TE based design": HubTE_index_tracking,
        "TE_1 based design": TE1_index_tracking,
        "DR based design": DR_index_tracking
    },
    dataset={
        'prices': SP500_stocks_2015to2020,
        'index': SP500_index_2015to2020
    },
    lookback=T_trn,
    rebalance_every=6*21,
    optimize_every=6*21,
    verbose=True
)

plot_prices(returns, benchmark=SP500_index_2015to2020)

#
# T-Rex index tracking in Pyhon under development.
#
# See R version in: https://portfoliooptimizationbook.com/R-code/R-index-tracking.html
#

import time
from datetime import timedelta
from skfolio import Population
from skfolio.optimization import ConvexOptimization, MeanRisk, ObjectiveFunction
from skfolio.model_selection import cross_val_predict, WalkForward
from skfolio.preprocessing import prices_to_returns
from skfolio.measures import RiskMeasure
# For cardinality constraints (mixed-integer programming):
# pip install "cvxpy[scip]"

class Portfolio_via_CVXPY(ConvexOptimization):
    """Portfolio based on a custom CVXPY function (supports benchmark y)."""

    def __init__(self, cvxpy_fun=None, **kwargs):
        super().__init__(**kwargs)
        self.cvxpy_fun = cvxpy_fun

    def fit(self, X, y=None):
        if self.cvxpy_fun is None:
            raise ValueError("cvxpy_fun must be provided at initialization")
        weights = self.cvxpy_fun(X, y)
        self.weights_ = np.asarray(weights)
        return self

def batch_cross_val_predict(optimizers, X, cv, **kwargs) -> list:
    """Run cross_val_predict on multiple optimizers with timing information."""
    bt_list = []
    n_splits = cv.get_n_splits(X)
    total_start_time = time.time()

    print(f"Starting batch backtesting with {len(optimizers)} portfolio designs...")
    print(f"Cross-validation: {n_splits} splits")
    print("=" * 70)

    for i, optim in enumerate(optimizers, 1):
        if hasattr(optim, "portfolio_params") and optim.portfolio_params is not None:
            name = optim.portfolio_params.get("name", f"Portfolio {i}")
        else:
            name = f"Portfolio {i} ({type(optim).__name__})"

        print(f"[{i}/{len(optimizers)}] Backtesting '{name}'...")
        start_time = time.time()

        try:
            bt = cross_val_predict(
                optim, X, cv=cv,
                portfolio_params=getattr(optim, "portfolio_params", {}),
                **kwargs,
            )
            elapsed_str = str(timedelta(seconds=int(time.time() - start_time)))
            bt_list.append(bt)
            print(f"  Completed '{name}' in {elapsed_str}")

        except Exception as e:
            elapsed_str = str(timedelta(seconds=int(time.time() - start_time)))
            print(f"  Failed '{name}' after {elapsed_str}: {str(e)}")
            raise

    total_elapsed_str = str(timedelta(seconds=int(time.time() - total_start_time)))
    print("=" * 70)
    print(f"Batch backtesting completed!")
    print(f"Total time: {total_elapsed_str}")
    print(f"Successfully processed {len(bt_list)} portfolios")

    return bt_list

# Convert prices to returns (both assets and benchmark)
X, y = prices_to_returns(SP500_stocks_2015to2020, SP500_index_2015to2020)
print(f"Asset returns: {X.shape[0]} observations × {X.shape[1]} assets")
print(f"Benchmark returns: {y.shape[0]} observations")

portf_min_var_TE = MeanRisk(
    objective_function=ObjectiveFunction.MINIMIZE_RISK,
    risk_measure=RiskMeasure.VARIANCE,
    max_tracking_error=0.003,
    portfolio_params=dict(name="Min-Variance with TE constraint"),
)

T_trn = 2 * 252
portf_min_var_TE.fit(X.iloc[:T_trn], y.iloc[:T_trn])
w = portf_min_var_TE.weights_
print(f"Sum of weights: {w.sum():.4f}")
print(f"Cardinality (non-zero weights): {np.sum(w > 1e-5)}")
print(f"Min weight: {w.min():.6f}")

def tracking_wfixed(X, y, lambda_val=5e-7):
    """Sparse index tracking assuming a fixed portfolio (adapter for skfolio)."""
    r = y.squeeze()
    return sparse_tracking_via_MM(X, r, lambda_val=lambda_val, max_iter=10)


def tracking_wt(X, y, lambda_val=5e-7):
    """Sparse index tracking assuming a time-varying portfolio (adapter for skfolio)."""
    r = y.squeeze()
    alpha = compute_alpha(X, r)
    w0 = sparse_tracking_via_MM(alpha["X_tilde"], r, lambda_val=lambda_val, max_iter=10)
    w = w0 * alpha["last_alpha_lagged"]
    w_sum = w.sum()
    if w_sum > 0:
        w = w / w_sum
    return w

portf_sparse_fixed = Portfolio_via_CVXPY(
    cvxpy_fun=tracking_wfixed,
    portfolio_params=dict(name="Sparse tracking (fixed w)"),
)

portf_sparse_tv = Portfolio_via_CVXPY(
    cvxpy_fun=tracking_wt,
    portfolio_params=dict(name="Sparse tracking (time-varying w)"),
)

portf_sparse_fixed.fit(X.iloc[:T_trn], y.iloc[:T_trn])
w_fixed = portf_sparse_fixed.weights_
print("=== Fixed portfolio ===")
print(f"  Sum of weights: {w_fixed.sum():.4f}")
print(f"  Cardinality:    {np.sum(w_fixed > 1e-5)}")

portf_sparse_tv.fit(X.iloc[:T_trn], y.iloc[:T_trn])
w_tv = portf_sparse_tv.weights_
print("\n=== Time-varying portfolio ===")
print(f"  Sum of weights: {w_tv.sum():.4f}")
print(f"  Cardinality:    {np.sum(w_tv > 1e-5)}")

all_formulations = [
    portf_min_var_TE,
    portf_sparse_fixed,
    portf_sparse_tv,
]

cv = WalkForward(train_size=T_trn, test_size=6 * 21, purged_size=1)

bt_list = batch_cross_val_predict(
    all_formulations,
    X=X,
    cv=cv,
    y=y,
)

bt_population = Population(bt_list)

bt_population.summary()

fig = bt_population.plot_cumulative_returns()

# Overlay benchmark cumulative returns
existing_x = fig.data[0].x
bench_oos = y.loc[bt_population[0].observations].squeeze()
bench_cum = (1 + bench_oos).cumprod() - 1

fig.add_scatter(
    x=existing_x,
    y=bench_cum.values,
    name="S&P 500",
    line=dict(color="black", dash="dash", width=1.5),
)
fig.update_layout(
    title="Cumulative Returns (out-of-sample)",
    yaxis_title="Growth of $1",
    legend=dict(x=0.02, y=0.98),
)
fig.show()

results = []
for bt in bt_population:
    oos_ret = bt.returns          # numpy array
    oos_obs = bt.observations     # date array
    bench_ret = y.loc[oos_obs].values
    excess = oos_ret - bench_ret
    te_mse = (excess ** 2).mean()
    te_rmse = np.sqrt(te_mse)
    results.append({
        "Portfolio": bt.name,
        "TE (MSE)": f"{te_mse:.2e}",
        "TE (daily RMSE)": f"{te_rmse:.4%}",
        "TE (annualized RMSE)": f"{te_rmse * np.sqrt(252):.2%}",
        "Cardinality": int(np.sum(bt[0].weights > 1e-5)),
    })

pd.DataFrame(results).set_index("Portfolio")

from skfolio.optimization import BenchmarkTracker

# --- Approach 1: Return-based TE constraint (min CVaR s.t. TE ≤ 0.3%) ---
portf_te_constraint = MeanRisk(
    objective_function=ObjectiveFunction.MINIMIZE_RISK,
    risk_measure=RiskMeasure.CVAR,
    max_tracking_error=0.003,
    portfolio_params=dict(name="Min CVaR (TE ≤ 0.3%)"),
)

# --- Approach 2: Weight-based target (track equal-weighted portfolio) ---
n_assets = X.shape[1]
target_w = np.ones(n_assets) / n_assets
portf_target_weights = MeanRisk(
    objective_function=ObjectiveFunction.MINIMIZE_RISK,
    risk_measure=RiskMeasure.STANDARD_DEVIATION,
    target_weights=target_w,
    portfolio_params=dict(name="Track EW target (weight-based)"),
)

# --- Approach 3: BenchmarkTracker (minimize TE on excess returns) ---
portf_benchmark_tracker = BenchmarkTracker(
    risk_measure=RiskMeasure.STANDARD_DEVIATION,
    portfolio_params=dict(name="BenchmarkTracker (dense)"),
)

# --- Approach 3b: BenchmarkTracker with cardinality for sparsity ---
portf_benchmark_tracker_sparse = BenchmarkTracker(
    risk_measure=RiskMeasure.STANDARD_DEVIATION,
    cardinality=20,          # at most 20 assets
    solver="SCIP",           # MIP solver needed for cardinality
    solver_params=dict(scip_params={"limits/time": 60}),  # 60s per solve
    portfolio_params=dict(name="BenchmarkTracker (K=20)"),
)

native_formulations = [
    portf_te_constraint,
    portf_target_weights,
    portf_benchmark_tracker,
    portf_benchmark_tracker_sparse,
]

cv = WalkForward(train_size=T_trn, test_size=6 * 21, purged_size=1)

bt_native = batch_cross_val_predict(
    native_formulations,
    X=X,
    cv=cv,
    y=y,
)

bt_native_pop = Population(bt_native)

bt_native_pop.summary()

fig = bt_native_pop.plot_cumulative_returns()

existing_x = fig.data[0].x
bench_oos = y.loc[bt_native_pop[0].observations].squeeze()
bench_cum = (1 + bench_oos).cumprod() - 1

fig.add_scatter(
    x=existing_x,
    y=bench_cum.values,
    name="S&P 500",
    line=dict(color="black", dash="dash", width=1.5),
)
fig.update_layout(
    title="Cumulative Returns – Native skfolio Tracking (out-of-sample)",
    yaxis_title="Growth of $1",
    legend=dict(x=0.02, y=0.98),
)
fig.show()

results = []
for bt in bt_native_pop:
    oos_ret = bt.returns
    oos_obs = bt.observations
    bench_ret = y.loc[oos_obs].values
    excess = oos_ret - bench_ret
    te_mse = (excess ** 2).mean()
    te_rmse = np.sqrt(te_mse)
    card = int(np.sum(bt[0].weights > 1e-5))
    results.append({
        "Portfolio": bt.name,
        "TE (MSE)": f"{te_mse:.2e}",
        "TE (daily RMSE)": f"{te_rmse:.4%}",
        "TE (annualized RMSE)": f"{te_rmse * np.sqrt(252):.2%}",
        "Cardinality": card,
    })

pd.DataFrame(results).set_index("Portfolio")