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

# Financial data and performance metrics
import empyrical as emp

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

# Plotting
import matplotlib.pyplot as plt
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 quadprog import solve_qp          # QP solver used for the mean-variance designs
from scipy.stats import chi2, gaussian_kde, norm  # statistical distributions
from tqdm.auto import tqdm             # progress bars

def compute_returns(prices: pd.DataFrame) -> pd.DataFrame:
    """
    Compute arithmetic returns from a price DataFrame.

    Parameters
    ----------
    prices : pd.DataFrame
        T x N DataFrame of asset prices.

    Returns
    -------
    pd.DataFrame
        (T - 1) x N DataFrame of arithmetic returns.
    """
    return prices.pct_change().dropna()


def estimate_parameters(X) -> tuple[np.ndarray, np.ndarray]:
    """
    Sample estimates of the mean vector and the covariance matrix.

    Parameters
    ----------
    X : array-like
        T x N matrix of asset returns.

    Returns
    -------
    mu : np.ndarray
        Length-N sample mean vector.
    Sigma : np.ndarray
        N x N sample covariance matrix (with ``ddof=1``, as R's ``cov``).
    """
    X = np.asarray(X, dtype=float)
    return X.mean(axis=0), np.cov(X.T)

def design_MVP(
        mu: np.ndarray,
        Sigma: np.ndarray,
        lambda_val: float = 1.0,
        reg: float = 1e-8,
) -> np.ndarray:
    """
    Design a mean-variance portfolio (MVP).

    Solves

        maximize    mu^T w - (lambda/2) w^T Sigma w
        subject to  w >= 0,  1^T w = 1.

    Parameters
    ----------
    mu : np.ndarray
        Length-N vector of expected returns.
    Sigma : np.ndarray
        N x N covariance matrix.
    lambda_val : float, default 1.0
        Risk-aversion parameter.
    reg : float, default 1e-8
        Tikhonov regularization, as a fraction of the mean diagonal of ``Sigma``,
        added to the quadratic term to keep the QP numerically well posed.

    Returns
    -------
    w : np.ndarray
        Length-N vector of portfolio weights.
    """
    n = len(mu)

    # # Solve the convex problem directly with CVXPY
    # w = cp.Variable(n)
    # objective = cp.Maximize(mu @ w - (lambda_val / 2) * cp.quad_form(w, cp.psd_wrap(Sigma)))
    # prob = cp.Problem(objective, [w >= 0, cp.sum(w) == 1])
    # prob.solve()
    # return w.value

    # Equivalent QP solved with quadprog (much faster, which matters in the
    # Monte Carlo experiments):
    #   minimize 1/2 w^T G w - a^T w   s.t.  C^T w >= b  (first meq columns of C are equalities)
    G = lambda_val * Sigma + reg * np.mean(np.diag(Sigma)) * np.eye(n)
    G = np.array(G, dtype=float)  # quadprog needs writable C-contiguous arrays
    a = np.array(mu, dtype=float)
    C = np.array(np.hstack([np.ones((n, 1)), np.eye(n)]), dtype=float)  # 1^T w = 1 and w >= 0
    b = np.array(np.concatenate([[1.0], np.zeros(n)]), dtype=float)
    return solve_qp(G, a, C, b, meq=1)[0]


def MVP_naive(X, lambda_val: float = 1.0) -> np.ndarray:
    """
    Naive mean-variance portfolio.

    Designs the mean-variance portfolio from the sample estimates of the mean
    vector and the covariance matrix of the given returns, ignoring the
    estimation error.

    Parameters
    ----------
    X : array-like
        T x N matrix of asset returns.
    lambda_val : float, default 1.0
        Risk-aversion parameter.

    Returns
    -------
    np.ndarray
        Length-N vector of portfolio weights.
    """
    mu_hat, Sigma_hat = estimate_parameters(X)
    return design_MVP(mu_hat, Sigma_hat, lambda_val=lambda_val)

def monte_carlo_estimation_error(
        design_funs: dict[str, callable],
        mu: np.ndarray,
        Sigma: np.ndarray,
        T: int,
        num_realizations: int = 1000,
        seed: int = 42,
        verbose: bool = True,
) -> dict[str, np.ndarray]:
    """
    Monte Carlo study of the sensitivity to estimation errors.

    For each realization, a synthetic sample of ``T`` observations is drawn from
    N(mu, Sigma) and every design function is applied to the same sample, so that
    the differences between designs are not contaminated by different noise.

    Parameters
    ----------
    design_funs : dict[str, callable]
        Mapping from a method name to a function that takes a T x N matrix of
        (noisy) returns and returns a length-N weight vector.
    mu : np.ndarray
        True mean vector of the Gaussian model.
    Sigma : np.ndarray
        True covariance matrix of the Gaussian model.
    T : int
        Number of observations of each synthetic sample (the lookback length).
    num_realizations : int, default 1000
        Number of Monte Carlo realizations.
    seed : int, default 42
        Seed of the random number generator.
    verbose : bool, default True
        If True, show a progress bar.

    Returns
    -------
    dict[str, np.ndarray]
        Mapping from a method name to a ``num_realizations x N`` matrix of weights.
    """
    rng = np.random.default_rng(seed)
    Sigma = 0.5 * (Sigma + Sigma.T)  # enforce symmetry for multivariate_normal
    n = len(mu)
    ws = {name: np.zeros((num_realizations, n)) for name in design_funs}

    for b in tqdm(range(num_realizations), disable=not verbose):
        X_noisy = rng.multivariate_normal(mu, Sigma, size=T)
        for name, fun in design_funs.items():
            ws[name][b] = np.asarray(fun(X_noisy), dtype=float)

    return ws


def mean_variance_value(w: np.ndarray, mu: np.ndarray, Sigma: np.ndarray) -> float:
    """Mean-variance objective value mu^T w - (1/2) w^T Sigma w."""
    return float(w @ mu - 0.5 * w @ Sigma @ w)


def sharpe_ratio(w: np.ndarray, mu: np.ndarray, Sigma: np.ndarray) -> float:
    """Sharpe ratio (no risk-free rate) of a portfolio under the model parameters."""
    return float((w @ mu) / np.sqrt(w @ Sigma @ w))


def evaluate_portfolios(
        ws: dict[str, np.ndarray],
        mu: np.ndarray,
        Sigma: np.ndarray,
) -> pd.DataFrame:
    """
    Evaluate a collection of portfolio realizations.

    Parameters
    ----------
    ws : dict[str, np.ndarray]
        Mapping from a method name to a ``num_realizations x N`` matrix of weights.
    mu : np.ndarray
        True mean vector used to evaluate the portfolios.
    Sigma : np.ndarray
        True covariance matrix used to evaluate the portfolios.

    Returns
    -------
    pd.DataFrame
        Tidy (long-format) DataFrame with one row per (method, realization) and
        columns ``method``, ``realization``, ``mean-variance`` and ``Sharpe ratio``.
    """
    rows = []
    for name, W in ws.items():
        rows.append(pd.DataFrame({
            "method": name,
            "realization": np.arange(1, len(W) + 1),
            "mean-variance": [mean_variance_value(w, mu, Sigma) for w in W],
            "Sharpe ratio": [sharpe_ratio(w, mu, Sigma) for w in W],
        }))

    df = pd.concat(rows, ignore_index=True)
    df["method"] = pd.Categorical(df["method"], categories=list(ws), ordered=True)
    return df

def plot_weights(
        ws: np.ndarray,
        assets: list[str],
        title: str = "Portfolio for different estimation errors",
        ax: plt.Axes | None = None,
) -> None:
    """
    Grouped bar chart of the portfolio weights across realizations.

    Parameters
    ----------
    ws : np.ndarray
        ``num_realizations x N`` matrix of weights, one row per realization.
    assets : list[str]
        Length-N list of asset names.
    title : str, default 'Portfolio for different estimation errors'
        Title of the plot.
    ax : plt.Axes or None, default None
        Axes where to draw the plot. If None, a new figure is created.
    """
    df = pd.DataFrame(
        np.asarray(ws).T,
        index=assets,
        columns=[f"#{i + 1}" for i in range(len(ws))],
    )

    own_figure = ax is None
    if own_figure:
        _, ax = plt.subplots(figsize=(12, 4))

    df.plot(kind="bar", ax=ax, width=0.85, edgecolor="black", linewidth=0.4)
    ax.set_title(title, fontsize=11)
    ax.set_xlabel(None)
    ax.set_ylabel("weight")
    ax.tick_params(axis="x", rotation=0, labelsize=8)
    ax.legend(title="realization", ncol=len(df.columns), fontsize=8, title_fontsize=8)

    if own_figure:
        plt.tight_layout()
        plt.show()


def plot_distribution(
        df: pd.DataFrame,
        x: str,
        title: str,
        xlim: tuple[float | None, float | None] | None = None,
        reference: float | None = None,
        reference_label: str = "clairvoyant portfolio",
        bw_method: float | str | None = None,
) -> None:
    """
    Ridgeline plot of the empirical distribution of a performance metric.

    Each method is drawn on its own baseline (as in the book's
    ``geom_density_ridges`` figures), with the densities shifted vertically so
    that they overlap partially.

    Parameters
    ----------
    df : pd.DataFrame
        Tidy DataFrame produced by ``evaluate_portfolios``.
    x : str
        Name of the column containing the metric to plot.
    title : str
        Title of the plot.
    xlim : tuple or None, default None
        Limits of the x-axis (``None`` leaves the corresponding limit free).
    reference : float or None, default None
        Value to mark with a vertical line (e.g., the clairvoyant portfolio).
    reference_label : str, default 'clairvoyant portfolio'
        Label of the reference line.
    bw_method : float, str or None, default None
        Bandwidth of the Gaussian kernel density estimate.
    """
    methods = list(df["method"].cat.categories)
    n_methods = len(methods)
    colors = sns.color_palette("deep", n_methods)

    # Grid where the densities are evaluated
    values = df[x].to_numpy()
    grid = np.linspace(np.nanmin(values), np.nanmax(values), 500)
    densities = {
        name: gaussian_kde(df.loc[df["method"] == name, x].to_numpy(), bw_method=bw_method)(grid)
        for name in methods
    }
    max_density = max(d.max() for d in densities.values())

    # Vertical offsets: the first method is drawn on top
    offsets = np.arange(n_methods)[::-1]

    _, ax = plt.subplots(figsize=(9, 1.0 + 0.85 * n_methods))
    for k, name in enumerate(methods):
        y = offsets[k] + 0.9 * densities[name] / max_density
        ax.fill_between(grid, offsets[k], y, alpha=0.6, color=colors[k],
                        edgecolor="black", linewidth=1.0)

    if reference is not None:
        ax.axvline(reference, color="black", linestyle="--", linewidth=1.2,
                   label=reference_label)
        ax.legend(loc="upper right", fontsize=8)

    ax.set_yticks(offsets)
    ax.set_yticklabels(methods, fontsize=9)
    ax.set_ylim(-0.2, n_methods)
    if xlim is not None:
        ax.set_xlim(xlim[0], xlim[1])

    ax.set_title(title)
    ax.set_xlabel(x)
    ax.set_ylabel(None)
    plt.tight_layout()
    plt.show()

def backtest(
        portfolio_funcs: dict[str, callable],
        prices: pd.DataFrame,
        lookback: int,
        optimize_every: int = 21,
        rebalance_every: int = 1,
        verbose: bool = True,
) -> tuple[pd.DataFrame, dict[str, pd.DataFrame]]:
    """
    Walk-forward backtest of several portfolio designs.

    Parameters
    ----------
    portfolio_funcs : dict[str, callable]
        Mapping from a strategy name to a function that takes a
        ``lookback x N`` matrix of returns and returns a length-N weight vector.
    prices : pd.DataFrame
        T x N DataFrame of asset prices.
    lookback : int
        Number of past returns used to design the portfolios.
    optimize_every : int, default 21
        Number of periods between re-optimizations.
    rebalance_every : int, default 1
        Number of periods between rebalancings to the latest target weights.
    verbose : bool, default True
        If True, print the re-optimization dates.

    Returns
    -------
    returns : pd.DataFrame
        (T - lookback) x K DataFrame of out-of-sample portfolio returns.
    weights : dict[str, pd.DataFrame]
        Mapping from a strategy name to a (T - lookback) x N DataFrame of weights.
    """
    if optimize_every % rebalance_every != 0:
        raise ValueError("optimize_every must be a multiple of rebalance_every")

    X = compute_returns(prices)
    rets = X.values
    T, N = X.shape
    portf_rets = {name: np.full(T, np.nan) for name in portfolio_funcs}
    all_w = {name: np.full((T, N), np.nan) for name in portfolio_funcs}

    for name, fun in portfolio_funcs.items():
        current_w = np.zeros(N)
        target_w = np.zeros(N)

        for t in range(lookback, T):
            if (t - lookback) % rebalance_every == 0:
                if (t - lookback) % optimize_every == 0:
                    target_w = np.asarray(fun(X.iloc[t - lookback:t]), dtype=float)
                    if verbose:
                        print(f"  Re-optimizing '{name}' at {X.index[t].date()}")
                current_w = target_w.copy()

            all_w[name][t] = current_w
            portf_rets[name][t] = rets[t] @ current_w

            # drift the weights with the realized returns
            current_w = current_w * (1 + rets[t])
            current_w = current_w / current_w.sum()

    returns = pd.DataFrame(portf_rets, index=X.index).iloc[lookback:]
    weights = {
        name: pd.DataFrame(w, index=X.index, columns=prices.columns).iloc[lookback:]
        for name, w in all_w.items()
    }
    return returns, weights


def plot_backtest(returns: pd.DataFrame, title: str | None = None) -> None:
    """
    Cumulative P&L (top panel) and drawdown (bottom panel) of each portfolio.

    Parameters
    ----------
    returns : pd.DataFrame
        T x K DataFrame of portfolio returns.
    title : str or None, default None
        Optional suptitle of the figure.
    """
    NAV = (1 + returns).cumprod()
    drawdown = NAV / NAV.cummax() - 1

    _, axes = plt.subplots(2, 1, figsize=(12, 7), sharex=True)
    NAV.plot(ax=axes[0], linewidth=1.1)
    axes[0].set_title("Cumulative P&L")
    axes[0].set_ylabel("wealth")
    axes[0].set_xlabel(None)
    axes[0].legend(title=None, fontsize=8)

    drawdown.plot(ax=axes[1], linewidth=1.1)
    axes[1].set_title("Drawdown")
    axes[1].set_ylabel("drawdown")
    axes[1].set_xlabel(None)
    axes[1].legend(title=None, fontsize=8)

    if title is not None:
        plt.suptitle(title)

    plt.tight_layout()
    plt.show()


def performance_metrics(returns: pd.DataFrame) -> pd.DataFrame:
    """
    Summary performance statistics of each portfolio.

    Parameters
    ----------
    returns : pd.DataFrame
        T x K DataFrame of daily portfolio returns.

    Returns
    -------
    pd.DataFrame
        K x 4 DataFrame with the annualized return, annualized volatility,
        Sharpe ratio and maximum drawdown of each portfolio.
    """
    return pd.DataFrame({
        "annual return": returns.apply(emp.annual_return),
        "annual volatility": returns.apply(emp.annual_volatility),
        "Sharpe ratio": returns.apply(emp.sharpe_ratio),
        "max drawdown": returns.apply(emp.max_drawdown),
    })

def financial_data_resample(
        prices: pd.DataFrame,
        num_datasets: int = 100,
        N_sample: int = 50,
        T_sample: int = 252,
        seed: int = 42,
) -> list[pd.DataFrame]:
    """
    Generate resampled price datasets.

    Each dataset is built from a random subset of ``N_sample`` assets and a random
    contiguous time window of ``T_sample`` observations.

    Parameters
    ----------
    prices : pd.DataFrame
        T x N DataFrame of asset prices.
    num_datasets : int, default 100
        Number of resampled datasets.
    N_sample : int, default 50
        Number of assets of each dataset.
    T_sample : int, default 252
        Number of observations of each dataset.
    seed : int, default 42
        Seed of the random number generator.

    Returns
    -------
    list[pd.DataFrame]
        List of resampled price DataFrames.
    """
    rng = np.random.default_rng(seed)
    datasets = []

    for _ in range(num_datasets):
        cols = np.sort(rng.choice(prices.shape[1], size=N_sample, replace=False))
        start = rng.integers(0, prices.shape[0] - T_sample)
        datasets.append(prices.iloc[start:start + T_sample, cols])

    return datasets


def multiple_backtests(
        portfolio_funcs: dict[str, callable],
        datasets: list[pd.DataFrame],
        lookback: int,
        optimize_every: int = 21,
        rebalance_every: int = 1,
        verbose: bool = True,
) -> pd.DataFrame:
    """
    Run the same walk-forward backtest on each resampled dataset.

    Parameters
    ----------
    portfolio_funcs : dict[str, callable]
        Mapping from a strategy name to a portfolio design function.
    datasets : list[pd.DataFrame]
        List of price DataFrames.
    lookback : int
        Number of past returns used to design the portfolios.
    optimize_every : int, default 21
        Number of periods between re-optimizations.
    rebalance_every : int, default 1
        Number of periods between rebalancings.
    verbose : bool, default True
        If True, show a progress bar.

    Returns
    -------
    pd.DataFrame
        Tidy DataFrame with one row per (dataset, method) and columns
        ``dataset``, ``method``, ``Sharpe ratio`` and ``max drawdown``.
    """
    rows = []

    for i, prices in enumerate(tqdm(datasets, disable=not verbose)):
        returns, _ = backtest(portfolio_funcs, prices, lookback=lookback,
                              optimize_every=optimize_every,
                              rebalance_every=rebalance_every, verbose=False)
        for name in returns.columns:
            rows.append({
                "dataset": i + 1,
                "method": name,
                "Sharpe ratio": emp.sharpe_ratio(returns[name]),
                "max drawdown": emp.max_drawdown(returns[name]),
            })

    perf = pd.DataFrame(rows)
    perf["method"] = pd.Categorical(perf["method"], categories=list(portfolio_funcs), ordered=True)
    return perf


def plot_performance_boxplots(
        perf: pd.DataFrame,
        metrics: tuple[str, ...] = ("Sharpe ratio", "max drawdown"),
) -> None:
    """
    Box plots of the out-of-sample performance metrics across resampled datasets.

    Parameters
    ----------
    perf : pd.DataFrame
        Tidy DataFrame produced by ``multiple_backtests``.
    metrics : tuple[str, ...]
        Names of the performance metrics to plot.
    """
    _, axes = plt.subplots(1, len(metrics), figsize=(11, 3.2))

    for ax, metric in zip(np.atleast_1d(axes), metrics):
        sns.boxplot(data=perf, x="method", y=metric, ax=ax, width=0.5, color="#4C72B0")
        ax.set_title(metric)
        ax.set_xlabel(None)
        ax.tick_params(axis="x", rotation=15, labelsize=8)

    plt.tight_layout()
    plt.show()

# Ten liquid stocks from 2019 onwards
tickers = ["AAPL", "AMZN", "AMD", "GM", "GOOGL", "MGM", "MSFT", "QCOM", "TSCO", "UPS"]
stock_prices = SP500_stocks_2015to2020[tickers].loc["2019":]
X = compute_returns(stock_prices)
mu, Sigma = estimate_parameters(X)

# Six realizations of the estimation error
ws = monte_carlo_estimation_error(
    design_funs={"Naive MVP": MVP_naive},
    mu=mu, Sigma=Sigma, T=len(X),
    num_realizations=6, seed=42, verbose=False,
)["Naive MVP"]

plot_weights(ws, assets=list(stock_prices.columns))

def design_MVP_robust_mu_box(
        mu_hat: np.ndarray,
        delta: np.ndarray,
        Sigma_hat: np.ndarray,
        lambda_val: float = 1.0,
) -> np.ndarray:
    """
    Robust mean-variance portfolio under a box uncertainty set for the mean.

    Solves

        maximize    w^T mu_hat - |w|^T delta - (lambda/2) w^T Sigma_hat w
        subject to  w >= 0,  1^T w = 1.

    Parameters
    ----------
    mu_hat : np.ndarray
        Length-N vector of estimated expected returns.
    delta : np.ndarray
        Length-N vector with the half-widths of the box.
    Sigma_hat : np.ndarray
        N x N estimated covariance matrix.
    lambda_val : float, default 1.0
        Risk-aversion parameter.

    Returns
    -------
    w : np.ndarray
        Length-N vector of portfolio weights.
    """
    n = len(mu_hat)
    w = cp.Variable(n)
    objective = cp.Maximize(
        mu_hat @ w
        - cp.sum(cp.multiply(cp.abs(w), delta))
        - (lambda_val / 2) * cp.quad_form(w, cp.psd_wrap(Sigma_hat))
    )
    prob = cp.Problem(objective, [w >= 0, cp.sum(w) == 1])
    prob.solve()
    return w.value


def design_MVP_robust_mu_ellipsoid(
        mu_hat: np.ndarray,
        Sigma_hat: np.ndarray,
        kappa_mu: float = 0.1,
        S_mu: np.ndarray | None = None,
        lambda_val: float = 1.0,
) -> np.ndarray:
    """
    Robust mean-variance portfolio under an ellipsoidal uncertainty set for the mean.

    Solves

        maximize    w^T mu_hat - kappa ||S^{1/2} w||_2 - (lambda/2) w^T Sigma_hat w
        subject to  w >= 0,  1^T w = 1,

    where S is the shape matrix of the ellipsoid (by default S = Sigma_hat).

    Parameters
    ----------
    mu_hat : np.ndarray
        Length-N vector of estimated expected returns.
    Sigma_hat : np.ndarray
        N x N estimated covariance matrix (used for the risk term).
    kappa_mu : float, default 0.1
        Size of the ellipsoid (radius of the uncertainty set).
    S_mu : np.ndarray or None, default None
        N x N shape matrix of the ellipsoid. If None, ``Sigma_hat`` is used.
    lambda_val : float, default 1.0
        Risk-aversion parameter.

    Returns
    -------
    w : np.ndarray
        Length-N vector of portfolio weights.
    """
    if S_mu is None:
        S_mu = Sigma_hat

    n = len(mu_hat)
    S12 = np.linalg.cholesky(S_mu)  # lower triangular with S12 @ S12.T = S_mu
    w = cp.Variable(n)
    objective = cp.Maximize(
        mu_hat @ w
        - kappa_mu * cp.norm(S12.T @ w, 2)
        - (lambda_val / 2) * cp.quad_form(w, cp.psd_wrap(Sigma_hat))
    )
    prob = cp.Problem(objective, [w >= 0, cp.sum(w) == 1])
    prob.solve()
    return w.value


def MVP_robust_mu_box(X, kappa_mu: float = 0.2, lambda_val: float = 1.0) -> np.ndarray:
    """
    Robust mean-variance portfolio with a box uncertainty set for the mean.

    The half-widths of the box are set to ``kappa_mu`` times the standard
    deviations of the assets.

    Parameters
    ----------
    X : array-like
        T x N matrix of asset returns.
    kappa_mu : float, default 0.2
        Size of the box uncertainty set.
    lambda_val : float, default 1.0
        Risk-aversion parameter.

    Returns
    -------
    np.ndarray
        Length-N vector of portfolio weights.
    """
    mu_hat, Sigma_hat = estimate_parameters(X)
    delta = kappa_mu * np.sqrt(np.diag(Sigma_hat))
    return design_MVP_robust_mu_box(mu_hat, delta, Sigma_hat, lambda_val=lambda_val)


def MVP_robust_mu_ellipsoid(X, kappa_mu: float = 0.2, lambda_val: float = 1.0) -> np.ndarray:
    """
    Robust mean-variance portfolio with an ellipsoidal uncertainty set for the mean.

    Parameters
    ----------
    X : array-like
        T x N matrix of asset returns.
    kappa_mu : float, default 0.2
        Size of the ellipsoidal uncertainty set.
    lambda_val : float, default 1.0
        Risk-aversion parameter.

    Returns
    -------
    np.ndarray
        Length-N vector of portfolio weights.
    """
    mu_hat, Sigma_hat = estimate_parameters(X)
    return design_MVP_robust_mu_ellipsoid(mu_hat, Sigma_hat, kappa_mu=kappa_mu,
                                          lambda_val=lambda_val)

# Same ten stocks as in the introduction
stock_prices = SP500_stocks_2015to2020[tickers].loc["2019":]
X = compute_returns(stock_prices)
mu, Sigma = estimate_parameters(X)

ws = monte_carlo_estimation_error(
    design_funs={
        "Naive MVP": MVP_naive,
        "Robust MVP (box)": lambda X: MVP_robust_mu_box(X, kappa_mu=0.5),
        "Robust MVP (ellipsoidal)": lambda X: MVP_robust_mu_ellipsoid(X, kappa_mu=0.5),
    },
    mu=mu, Sigma=Sigma, T=len(X),
    num_realizations=6, seed=42, verbose=False,
)

_, axes = plt.subplots(1, 3, figsize=(16, 4))
for ax, (name, W) in zip(axes, ws.items()):
    plot_weights(W, assets=list(stock_prices.columns), title=name, ax=ax)
plt.tight_layout()
plt.show()

# Parameters
num_realizations = 1000
N = 50
lookback_months = 12
T_lookback = lookback_months * 21

# Random subset of 50 stocks from 2019 onwards
rng = np.random.default_rng(42)
cols = np.sort(rng.choice(SP500_stocks_2015to2020.shape[1], size=N, replace=False))
stock_prices = SP500_stocks_2015to2020.iloc[:, cols].loc["2019":]

# "True" parameters of the Gaussian model, estimated on the first training window
X_train = compute_returns(stock_prices).iloc[:T_lookback]
mu_true, Sigma_true = estimate_parameters(X_train)
w_clairvoyant = design_MVP(mu_true, Sigma_true, lambda_val=1)

# Monte Carlo over estimation-error realizations
design_funs = {
    "Naive MVP": MVP_naive,
    "Robust MVP (box uncertainty set)": MVP_robust_mu_box,
    "Robust MVP (ellipsoidal uncertainty set)": MVP_robust_mu_ellipsoid,
}

ws = monte_carlo_estimation_error(design_funs, mu_true, Sigma_true, T=len(X_train),
                                  num_realizations=num_realizations, seed=42, verbose=False)
df = evaluate_portfolios(ws, mu_true, Sigma_true)

plot_distribution(
    df, x="mean-variance",
    title="Empirical distribution of the mean-variance value",
    xlim=(-0.001, None),
    reference=mean_variance_value(w_clairvoyant, mu_true, Sigma_true),
)

plot_distribution(
    df, x="Sharpe ratio",
    title="Empirical distribution of the Sharpe ratio",
    xlim=(-0.05, None),
    reference=sharpe_ratio(w_clairvoyant, mu_true, Sigma_true),
)

N = 50
lookback_months = 6
rng = np.random.default_rng(42)
cols = np.sort(rng.choice(SP500_stocks_2015to2020.shape[1], size=N, replace=False))
stock_prices = SP500_stocks_2015to2020.iloc[:, cols].loc["2017":]

portfolio_funcs = {
    "Naive MVP": MVP_naive,
    "Robust MVP (box uncertainty set)": MVP_robust_mu_box,
    "Robust MVP (ellipsoidal uncertainty set)": MVP_robust_mu_ellipsoid,
}

returns, weights = backtest(
    portfolio_funcs,
    prices=stock_prices,
    lookback=lookback_months * 21,
    optimize_every=21,
    rebalance_every=1,
    verbose=False,
)

plot_backtest(returns)

performance_metrics(returns).style.format({
    "annual return": "{:.2%}",
    "annual volatility": "{:.2%}",
    "Sharpe ratio": "{:.2f}",
    "max drawdown": "{:.2%}",
})

datasets = financial_data_resample(
    SP500_stocks_2015to2020,
    num_datasets=100, N_sample=50, T_sample=12 * 21, seed=42,
)

perf = multiple_backtests(
    portfolio_funcs,
    datasets=datasets,
    lookback=6 * 21,
    optimize_every=21,
    rebalance_every=1,
    verbose=False,
)

plot_performance_boxplots(perf)

B = 1000  # number of bootstraps
n = 100   # number of observations for the estimation
mu_true, sd_true = 0.0, 1.0
rng = np.random.default_rng(42)

# Nominal sample
x_nominal = rng.normal(mu_true, sd_true, size=n)
mu_nominal = x_nominal.mean()

# Nonparametric bootstrap
bias = np.array([
    rng.choice(x_nominal, size=n, replace=True).mean() - mu_nominal
    for _ in range(B)
])

# Plot
_, ax = plt.subplots(figsize=(9, 4))
ax.hist(bias, bins=40, density=True, alpha=0.7, color="#4C72B0",
        edgecolor="0.3", label="nonparametric bootstrap")
grid = np.linspace(-0.4, 0.4, 500)
ax.plot(grid, norm.pdf(grid, loc=mu_true, scale=sd_true / np.sqrt(n)),
        color="black", linewidth=1.2, label="true distribution")
ax.set_xlim(-0.4, 0.4)
ax.set_yticks([])
ax.set_xlabel("bias")
ax.set_title("Empirical distribution of the mean estimation bias")
ax.legend(title=None, loc="upper right")
plt.tight_layout()
plt.show()

def bootstrap_portfolio(
        portfolio_fun: callable,
        X,
        num_bootstraps: int = 20,
        SSR_exp: float = 0.7,
        type_bootstrap: str = "vanilla",
        seed: int | None = None,
) -> np.ndarray:
    """
    Aggregated portfolio by resampling (bagging).

    Each of the ``num_bootstraps`` resampled designs is obtained by resampling the
    data and designing a portfolio, and the final portfolio is the average of the
    resampled designs.

    Parameters
    ----------
    portfolio_fun : callable
        Function that takes a returns matrix and returns a weight vector.
    X : array-like
        T x N matrix of asset returns.
    num_bootstraps : int, default 20
        Number of resampled designs B.
    SSR_exp : float, default 0.7
        Exponent that controls the number of assets in subset resampling:
        ceil(N**SSR_exp) assets are sampled at each design.
    type_bootstrap : str, default 'vanilla'
        Resampling scheme, one of 'vanilla' (resample the observations),
        'subset resampling' (resample the assets), or
        'vanilla bootstrap + subset resampling' (resample both).
    seed : int or None, default None
        Seed of the random number generator.

    Returns
    -------
    np.ndarray
        Aggregated length-N vector of portfolio weights.
    """
    valid_types = ("vanilla", "subset resampling", "vanilla bootstrap + subset resampling")
    if type_bootstrap not in valid_types:
        raise ValueError(f"type_bootstrap must be one of {valid_types}")

    X = np.asarray(X, dtype=float)
    T, N = X.shape
    rng = np.random.default_rng(seed)
    n_subset = int(np.ceil(N ** SSR_exp))
    w_sum = np.zeros(N)

    for _ in range(num_bootstraps):
        # resample the assets (columns) and/or the observations (rows)
        if "subset resampling" in type_bootstrap:
            cols = np.sort(rng.choice(N, size=n_subset, replace=False))
        else:
            cols = np.arange(N)
        if type_bootstrap.startswith("vanilla"):
            rows = rng.integers(0, T, size=T)
        else:
            rows = np.arange(T)

        w_sample = np.asarray(portfolio_fun(X[rows][:, cols]), dtype=float)
        w_sum[cols] += w_sample  # assets that are not sampled get zero weight

    return w_sum / num_bootstraps


def MVP_bagged(X, num_bootstraps: int = 20, lambda_val: float = 1.0, seed: int = 42) -> np.ndarray:
    """
    Bagged mean-variance portfolio (bootstrap aggregation over the observations).

    Parameters
    ----------
    X : array-like
        T x N matrix of asset returns.
    num_bootstraps : int, default 20
        Number of bootstrap designs B.
    lambda_val : float, default 1.0
        Risk-aversion parameter.
    seed : int, default 42
        Seed of the random number generator.

    Returns
    -------
    np.ndarray
        Length-N vector of portfolio weights.
    """
    return bootstrap_portfolio(
        portfolio_fun=lambda X_sample: MVP_naive(X_sample, lambda_val=lambda_val),
        X=X,
        num_bootstraps=num_bootstraps,
        type_bootstrap="vanilla",
        seed=seed,
    )


def MVP_SSR(X, num_bootstraps: int = 20, lambda_val: float = 1.0, seed: int = 42) -> np.ndarray:
    """
    Subset-resampled (SSR) mean-variance portfolio (resampling over the assets).

    Parameters
    ----------
    X : array-like
        T x N matrix of asset returns.
    num_bootstraps : int, default 20
        Number of resampled designs B.
    lambda_val : float, default 1.0
        Risk-aversion parameter.
    seed : int, default 42
        Seed of the random number generator.

    Returns
    -------
    np.ndarray
        Length-N vector of portfolio weights.
    """
    return bootstrap_portfolio(
        portfolio_fun=lambda X_sample: MVP_naive(X_sample, lambda_val=lambda_val),
        X=X,
        num_bootstraps=num_bootstraps,
        type_bootstrap="subset resampling",
        seed=seed,
    )


def MVP_SSR_bagged(X, num_bootstraps: int = 20, lambda_val: float = 1.0, seed: int = 42) -> np.ndarray:
    """
    Bagged subset-resampled mean-variance portfolio (resampling over observations and assets).

    Parameters
    ----------
    X : array-like
        T x N matrix of asset returns.
    num_bootstraps : int, default 20
        Number of resampled designs B.
    lambda_val : float, default 1.0
        Risk-aversion parameter.
    seed : int, default 42
        Seed of the random number generator.

    Returns
    -------
    np.ndarray
        Length-N vector of portfolio weights.
    """
    return bootstrap_portfolio(
        portfolio_fun=lambda X_sample: MVP_naive(X_sample, lambda_val=lambda_val),
        X=X,
        num_bootstraps=num_bootstraps,
        type_bootstrap="vanilla bootstrap + subset resampling",
        seed=seed,
    )

# Same ten stocks as in the introduction
stock_prices = SP500_stocks_2015to2020[tickers].loc["2019":]
X = compute_returns(stock_prices)
mu, Sigma = estimate_parameters(X)

ws = monte_carlo_estimation_error(
    design_funs={
        "Naive MVP": MVP_naive,
        "Bagged MVP": MVP_bagged,
        "Subset resampled MVP": MVP_SSR,
        "Subset bagged MVP": MVP_SSR_bagged,
    },
    mu=mu, Sigma=Sigma, T=len(X),
    num_realizations=6, seed=42, verbose=False,
)

_, axes = plt.subplots(2, 2, figsize=(15, 8))
for ax, (name, W) in zip(axes.ravel(), ws.items()):
    plot_weights(W, assets=list(stock_prices.columns), title=name, ax=ax)
plt.tight_layout()
plt.show()

# Parameters
num_realizations = 1000
N = 50
lookback_months = 12
T_lookback = lookback_months * 21

# Random subset of 50 stocks from 2019 onwards
rng = np.random.default_rng(42)
cols = np.sort(rng.choice(SP500_stocks_2015to2020.shape[1], size=N, replace=False))
stock_prices = SP500_stocks_2015to2020.iloc[:, cols].loc["2019":]

# "True" parameters of the Gaussian model, estimated on the first training window
X_train = compute_returns(stock_prices).iloc[:T_lookback]
mu_true, Sigma_true = estimate_parameters(X_train)
w_clairvoyant = design_MVP(mu_true, Sigma_true, lambda_val=1)

# Monte Carlo over estimation-error realizations
design_funs = {
    "Naive MVP": MVP_naive,
    "Bagged MVP": MVP_bagged,
    "Subset resampled MVP": MVP_SSR,
    "Subset bagged MVP": MVP_SSR_bagged,
}

ws = monte_carlo_estimation_error(design_funs, mu_true, Sigma_true, T=len(X_train),
                                  num_realizations=num_realizations, seed=42, verbose=False)
df = evaluate_portfolios(ws, mu_true, Sigma_true)

plot_distribution(
    df, x="mean-variance",
    title="Empirical distribution of the mean-variance value",
    xlim=(-0.001, None),
    reference=mean_variance_value(w_clairvoyant, mu_true, Sigma_true),
)

plot_distribution(
    df, x="Sharpe ratio",
    title="Empirical distribution of the Sharpe ratio",
    xlim=(-0.05, None),
    reference=sharpe_ratio(w_clairvoyant, mu_true, Sigma_true),
)

N = 50
lookback_months = 6
rng = np.random.default_rng(42)
cols = np.sort(rng.choice(SP500_stocks_2015to2020.shape[1], size=N, replace=False))
stock_prices = SP500_stocks_2015to2020.iloc[:, cols].loc["2017":]

portfolio_funcs = {
    "Naive MVP": MVP_naive,
    "Bagged MVP": MVP_bagged,
    "Subset resampled MVP": MVP_SSR,
    "Subset bagged MVP": MVP_SSR_bagged,
}

returns, weights = backtest(
    portfolio_funcs,
    prices=stock_prices,
    lookback=lookback_months * 21,
    optimize_every=21,
    rebalance_every=1,
    verbose=False,
)

plot_backtest(returns)

performance_metrics(returns).style.format({
    "annual return": "{:.2%}",
    "annual volatility": "{:.2%}",
    "Sharpe ratio": "{:.2f}",
    "max drawdown": "{:.2%}",
})

datasets = financial_data_resample(
    SP500_stocks_2015to2020,
    num_datasets=100, N_sample=50, T_sample=12 * 21, seed=42,
)

perf = multiple_backtests(
    portfolio_funcs,
    datasets=datasets,
    lookback=6 * 21,
    optimize_every=21,
    rebalance_every=1,
    verbose=False,
)

plot_performance_boxplots(perf)

from skfolio.optimization import DistributionallyRobustCVaR, MeanRisk, ObjectiveFunction
from skfolio.measures import RiskMeasure
from skfolio.uncertainty_set import BootstrapMuUncertaintySet, EmpiricalMuUncertaintySet
from skfolio.model_selection import WalkForward, cross_val_predict
from skfolio.population import Population
from skfolio.preprocessing import prices_to_returns

# Same 50 stocks from 2017 onwards as in the backtests
rng = np.random.default_rng(42)
cols = np.sort(rng.choice(SP500_stocks_2015to2020.shape[1], size=50, replace=False))
stock_prices = SP500_stocks_2015to2020.iloc[:, cols].loc["2017":]
X_sk = prices_to_returns(stock_prices)

# Training window and statistically calibrated radius
T_trn = 252
X_first = X_sk.iloc[:T_trn]
mu_hat, Sigma_hat = estimate_parameters(X_first)
n_assets = X_sk.shape[1]
kappa_sk = float(np.sqrt(chi2.ppf(0.95, n_assets)))

w_ours = design_MVP_robust_mu_ellipsoid(mu_hat, Sigma_hat, kappa_mu=kappa_sk,
                                        S_mu=Sigma_hat / T_trn, lambda_val=1)

mr_robust_full = MeanRisk(
    objective_function=ObjectiveFunction.MAXIMIZE_UTILITY,
    risk_measure=RiskMeasure.VARIANCE,
    risk_aversion=0.5,
    mu_uncertainty_set_estimator=EmpiricalMuUncertaintySet(confidence_level=0.95, diagonal=False),
).fit(X_first)

print(f"Kappa (chi-squared): {kappa_sk:.3f}")
print(f"Max |difference| between our design and skfolio's: "
      f"{np.abs(mr_robust_full.weights_ - w_ours).max():.2e}")

optimizers = {
    "Naive MVP": MeanRisk(
        objective_function=ObjectiveFunction.MAXIMIZE_UTILITY,
        risk_measure=RiskMeasure.VARIANCE,
        risk_aversion=0.5,
    ),
    "Robust MVP (ellipsoidal, diagonal)": MeanRisk(
        objective_function=ObjectiveFunction.MAXIMIZE_UTILITY,
        risk_measure=RiskMeasure.VARIANCE,
        risk_aversion=0.5,
        mu_uncertainty_set_estimator=EmpiricalMuUncertaintySet(confidence_level=0.95, diagonal=True),
    ),
    "Robust MVP (ellipsoidal, full)": MeanRisk(
        objective_function=ObjectiveFunction.MAXIMIZE_UTILITY,
        risk_measure=RiskMeasure.VARIANCE,
        risk_aversion=0.5,
        mu_uncertainty_set_estimator=EmpiricalMuUncertaintySet(confidence_level=0.95, diagonal=False),
    ),
    "Robust MVP (bootstrap ellipsoidal)": MeanRisk(
        objective_function=ObjectiveFunction.MAXIMIZE_UTILITY,
        risk_measure=RiskMeasure.VARIANCE,
        risk_aversion=0.5,
        mu_uncertainty_set_estimator=BootstrapMuUncertaintySet(
            confidence_level=0.95, diagonal=False, n_bootstrap_samples=500, seed=42),
    ),
    "DRO CVaR (Wasserstein ball)": DistributionallyRobustCVaR(
        risk_aversion=1.0, cvar_beta=0.95, wasserstein_ball_radius=0.02,
    ),
}

cv = WalkForward(train_size=T_trn, test_size=6 * 21, purged_size=1)
print(f"Number of walk-forward folds: {cv.get_n_splits(X_sk)}")

population = Population([
    cross_val_predict(optimizer, X_sk, cv=cv, portfolio_params=dict(name=name))
    for name, optimizer in optimizers.items()
])

summary = population.summary()
summary.loc[[
    "Annualized Mean",
    "Annualized Standard Deviation",
    "Annualized Sharpe Ratio",
    "MAX Drawdown",
    "Worst Realization",
    "CVaR at 95%",
]]

cardinality = pd.Series(
    {
        ptf.name: float(np.mean([(sub.weights > 1e-3).sum() for sub in ptf]))
        for ptf in population
    },
    name="average number of active assets",
)
cardinality.round(1)

fig = population.plot_cumulative_returns()
fig.update_layout(
    title="Cumulative returns (out-of-sample)",
    yaxis_title="Cumulative return",
    legend=dict(x=0.02, y=0.98),
)
fig.show()

fig = population.plot_drawdowns()
fig.update_layout(title="Drawdowns (out-of-sample)")
fig.show()