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

# Clustering and distances
from scipy.cluster.hierarchy import dendrogram
from scipy.spatial.distance import pdist, squareform

# Optimization (for the GMVP and MVP benchmarks)
import cvxpy as cp

# 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

# Dependence measures used as "sophisticated" graphs
from skfolio.distance import DistanceCorrelation

# Plotting
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_theme(style="darkgrid")

# Progress bars
from tqdm.auto import tqdm

# color palette (the same used in the R companion)
BLUE, RED, BLACK = "#2980b9", "#c0392b", "black"

def hclust(D, method="single"):
    """
    Agglomerative hierarchical clustering via the Lance-Williams recursion.

    Parameters
    ----------
    D : np.ndarray
        N x N symmetric matrix of distances (the diagonal is ignored).
    method : {"single", "complete", "average", "ward"}, default "single"
        Linkage method. "ward" applies Ward's update to the distances themselves,
        as R's hclust(method = "ward.D").

    Returns
    -------
    np.ndarray
        (N - 1) x 4 linkage matrix: children of each merge, merge height and
        number of leaves of the new cluster.
    """
    D = np.asarray(D, dtype=float)
    N = len(D)
    if N == 1:
        return np.zeros((0, 4))

    # ids N, N+1, ... denote the internal nodes of the tree
    d = np.full((2 * N - 1, 2 * N - 1), np.inf)
    d[:N, :N] = D
    np.fill_diagonal(d, np.inf)

    size = np.ones(2 * N - 1)
    active = np.zeros(2 * N - 1, dtype=bool)
    active[:N] = True
    Z = np.zeros((N - 1, 4))
    next_id = N

    for step in range(N - 1):
        # closest pair among the active clusters
        sub = np.where(active[:, None] & active[None, :], d, np.inf)
        i, j = divmod(int(np.argmin(sub)), len(d))
        if i > j:
            i, j = j, i
        Z[step] = [i, j, d[i, j], size[i] + size[j]]

        # Lance-Williams update of the distances to the new cluster
        ni, nj = size[i], size[j]
        if method == "single":
            new_d = np.minimum(d[i], d[j])
        elif method == "complete":
            new_d = np.maximum(d[i], d[j])
        elif method == "average":
            new_d = (ni * d[i] + nj * d[j]) / (ni + nj)
        elif method in ("ward", "ward.D"):
            nk = size
            new_d = ((ni + nk) * d[i] + (nj + nk) * d[j] - nk * d[i, j]) / (ni + nj + nk)
        else:
            raise ValueError(f"unknown linkage method: {method}")

        new_d[i] = np.inf
        new_d[j] = np.inf
        new_d[~active] = np.inf

        d[next_id] = new_d
        d[:, next_id] = new_d
        d[i, :] = d[j, :] = np.inf
        d[:, i] = d[:, j] = np.inf
        size[next_id] = ni + nj
        active[i] = active[j] = False
        active[next_id] = True
        next_id += 1

    return Z


def leaves_order(Z, N):
    """
    Leaf order of a dendrogram (R's hc$order).

    Parameters
    ----------
    Z : np.ndarray
        Linkage matrix.
    N : int
        Number of leaves.

    Returns
    -------
    list of int
        Order in which the leaves appear from left to right.
    """
    children = {k + N: (int(row[0]), int(row[1])) for k, row in enumerate(Z)}
    order = []

    def walk(node):
        if node < N:
            order.append(node)
        else:
            left, right = children[node]
            walk(left)
            walk(right)

    walk(2 * N - 2)
    return order


def cut_labels(Z, N, k):
    """
    Cluster labels when cutting a dendrogram into k clusters (R's cutree).

    Parameters
    ----------
    Z : np.ndarray
        Linkage matrix.
    N : int
        Number of leaves.
    k : int
        Number of clusters.

    Returns
    -------
    np.ndarray
        Length-N vector of labels 1, 2, ..., k.
    """
    parent = list(range(2 * N - 1))

    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    for step in range(N - k):
        a, b = int(Z[step, 0]), int(Z[step, 1])
        new_id = N + step
        parent[find(a)] = new_id
        parent[find(b)] = new_id

    labels = np.empty(N, dtype=int)
    mapping = {}
    for i in range(N):
        root = find(i)
        if root not in mapping:
            mapping[root] = len(mapping) + 1
        labels[i] = mapping[root]
    return labels

def corrcoef(Sigma):
    """Correlation matrix from a covariance matrix."""
    sd = np.sqrt(np.diag(Sigma))
    return Sigma / np.outer(sd, sd)


def dependence_matrix(X, kind="pearson"):
    """
    Matrix of dependence between the columns of X.

    Parameters
    ----------
    X : array-like
        T x N matrix of returns.
    kind : {"pearson", "spearman", "kendall", "distance-corr"}
        Dependence measure.

    Returns
    -------
    np.ndarray
        N x N matrix with values in [-1, 1] (or [0, 1] for the distance correlation).
    """
    X = np.asarray(X, dtype=float)
    if kind == "pearson":
        return np.corrcoef(X, rowvar=False)
    if kind in ("spearman", "kendall"):
        return pd.DataFrame(X).corr(method=kind).values
    if kind == "distance-corr":
        return DistanceCorrelation().fit(X).codependence_
    raise ValueError(f"unknown dependence measure: {kind}")


def hierarchical_distance(Sigma=None, X=None, distance="corr-dist2"):
    """
    Distance matrix used by the hierarchical clustering.

    Parameters
    ----------
    Sigma : np.ndarray or None
        Covariance matrix (used by the correlation-based distances).
    X : array-like or None
        T x N matrix of returns (used by the other dependence measures).
    distance : {"corr-dist", "corr-dist2", "spearman", "kendall", "distance-corr"}
        Distance: the correlation distance, its distance-of-distance, or the
        same two transformations applied to another dependence matrix.

    Returns
    -------
    np.ndarray
        N x N distance matrix.
    """
    if distance in ("corr-dist", "corr-dist2"):
        C = corrcoef(Sigma)
    else:
        C = dependence_matrix(X, kind=distance)

    D = np.sqrt(np.clip(0.5 * (1 - C), 0, None))
    np.fill_diagonal(D, 0.0)
    if distance == "corr-dist":
        return D
    return squareform(pdist(D))

def design_H1overN(Sigma=None, X=None, splitting="dendrogram", num_clusters=None,
                   linkage="single", distance="corr-dist2"):
    """
    Hierarchical 1/N portfolio.

    Parameters
    ----------
    Sigma : np.ndarray or None
        Covariance matrix (estimated from X when None).
    X : array-like or None
        T x N matrix of returns.
    splitting : {"dendrogram", "bisection"}, default "dendrogram"
        How the clusters are split: following the dendrogram or bisecting the
        current largest cluster (after the dendrogram reordering).
    num_clusters : int or None, default None
        Number of clusters (None clusters all the way down).
    linkage : {"single", "complete", "average", "Ward"}, default "single"
        Linkage method of the clustering.
    distance : str, default "corr-dist2"
        Distance matrix (see ``hierarchical_distance``).

    Returns
    -------
    np.ndarray
        Length-N vector of weights.
    """
    if Sigma is None:
        Sigma = np.cov(np.asarray(X, dtype=float), rowvar=False)
    Sigma = np.atleast_2d(np.asarray(Sigma, dtype=float))
    N = len(Sigma)
    num_clusters = N if num_clusters is None else num_clusters

    Z = hclust(hierarchical_distance(Sigma, X, distance), _linkage(linkage))
    order = leaves_order(Z, N)
    labels = np.column_stack([cut_labels(Z, N, k) for k in range(1, N + 1)])

    w = np.ones(N)
    L = [order]
    for kk in range(2, num_clusters + 1):
        i_split, items_1, items_2 = _split_set(L, labels, kk, splitting)
        L.pop(i_split)
        L += [items_1, items_2]
        w[items_1] *= 0.5
        w[items_2] *= 0.5

    # in case the clustering is not all the way down
    for items in L:
        w[items] = w[items] / len(items)

    return w


def _linkage(linkage):
    """Internal name of the linkage (as R's "Ward" -> "ward.D")."""
    return "ward" if str(linkage).lower() in ("ward", "ward.d", "ward.d2") else linkage


def _split_set(L, labels, kk, splitting):
    """Split one of the current sets of assets, as in the chapter's algorithm."""
    if splitting == "bisection":
        # maximal subset
        i_split = int(np.argmax([len(items) for items in L]))
        items = L[i_split]
        half = len(items) // 2
        return i_split, items[:half], items[half:]

    if splitting == "dendrogram":
        # element of L with the largest number of different clusters at level kk
        i_split = int(np.argmax([len(np.unique(labels[items, kk - 1])) for items in L]))
        items = L[i_split]
        values = np.unique(labels[items, kk - 1])
        first = [i for i in items if labels[i, kk - 1] == values[0]]
        second = [i for i in items if labels[i, kk - 1] == values[1]]
        return i_split, first, second

    raise ValueError(f"unknown splitting method: {splitting}")


def design_IVarP(Sigma):
    """
    Inverse-variance portfolio.

    Parameters
    ----------
    Sigma : np.ndarray
        Covariance matrix.

    Returns
    -------
    np.ndarray
        Length-N vector of weights.
    """
    Sigma = np.atleast_2d(np.asarray(Sigma, dtype=float))
    w = 1 / np.diag(Sigma)
    return w / w.sum()


def cluster_var(Sigma, items):
    """
    Variance of the inverse-variance portfolio on a subset of assets.

    Parameters
    ----------
    Sigma : np.ndarray
        Covariance matrix.
    items : list of int
        Indices of the assets in the cluster.

    Returns
    -------
    float
        Variance of the cluster portfolio.
    """
    sub = Sigma[np.ix_(items, items)]
    w = design_IVarP(sub)
    return float(w @ sub @ w)


def design_HRP(Sigma=None, X=None, splitting="bisection", num_clusters=None,
               linkage="single", distance="corr-dist2", final="IVarP"):
    """
    Hierarchical risk parity (HRP) portfolio.

    Parameters
    ----------
    Sigma : np.ndarray or None
        Covariance matrix (estimated from X when None).
    X : array-like or None
        T x N matrix of returns.
    splitting : {"bisection", "dendrogram"}, default "bisection"
        How the clusters are split.
    num_clusters : int or None, default None
        Number of clusters (None clusters all the way down).
    linkage : {"single", "complete", "average", "Ward"}, default "single"
        Linkage method of the clustering.
    distance : str, default "corr-dist2"
        Distance matrix (see ``hierarchical_distance``).
    final : {"IVarP", "1/N"}, default "IVarP"
        Allocation inside the final clusters.

    Returns
    -------
    np.ndarray
        Length-N vector of weights.
    """
    if Sigma is None:
        Sigma = np.cov(np.asarray(X, dtype=float), rowvar=False)
    Sigma = np.atleast_2d(np.asarray(Sigma, dtype=float))
    N = len(Sigma)
    num_clusters = N if num_clusters is None else num_clusters

    Z = hclust(hierarchical_distance(Sigma, X, distance), _linkage(linkage))
    order = leaves_order(Z, N)
    labels = np.column_stack([cut_labels(Z, N, k) for k in range(1, N + 1)])

    w = np.ones(N)
    L = [order]
    for kk in range(2, num_clusters + 1):
        i_split, items_1, items_2 = _split_set(L, labels, kk, splitting)
        L.pop(i_split)
        L += [items_1, items_2]

        var_1 = cluster_var(Sigma, items_1)
        var_2 = cluster_var(Sigma, items_2)
        alpha = 1 - var_1 / (var_1 + var_2)
        w[items_1] *= alpha
        w[items_2] *= 1 - alpha

    # in case the clustering is not all the way down
    for items in L:
        if len(items) > 1:
            if final == "1/N":
                w[items] = w[items] / len(items)
            elif final == "IVarP":
                w[items] = w[items[0]] * design_IVarP(Sigma[np.ix_(items, items)])
            else:
                raise ValueError(f"unknown final allocation: {final}")

    return w

def design_GMVP(Sigma):
    """
    Global minimum variance portfolio.

    Solves  minimize w' Sigma w  subject to  w >= 0, 1' w = 1.

    Parameters
    ----------
    Sigma : np.ndarray
        Covariance matrix.

    Returns
    -------
    np.ndarray
        Length-N vector of weights.
    """
    n = len(Sigma)
    w = cp.Variable(n)
    prob = cp.Problem(cp.Minimize(cp.quad_form(w, cp.psd_wrap(Sigma))), [w >= 0, cp.sum(w) == 1])
    prob.solve()
    return np.clip(w.value, 0, None)


def design_MVP(mu, Sigma, lambda_val=1.0):
    """
    Mean-variance portfolio.

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

    Parameters
    ----------
    mu : np.ndarray
        Expected returns.
    Sigma : np.ndarray
        Covariance matrix.
    lambda_val : float, default 1.0
        Risk-aversion parameter.

    Returns
    -------
    np.ndarray
        Length-N vector of weights.
    """
    n = len(Sigma)
    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 np.clip(w.value, 0, None)

def backtest(portfolio_funcs, prices, lookback, optimize_every=21, rebalance_every=1,
             cost=0.0):
    """
    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``
        window of prices and returns a length-N vector of weights.
    prices : pd.DataFrame
        T x N DataFrame of prices.
    lookback : int
        Number of past prices 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.
    cost : float, default 0.0
        Transaction cost per unit of traded amount (e.g., 60e-4 for 60 bps).

    Returns
    -------
    returns : pd.DataFrame
        (T - lookback - 1) x K DataFrame of daily out-of-sample returns.
    weights : dict[str, pd.DataFrame]
        Mapping from a strategy name to a (T - lookback - 1) x N DataFrame of
        the weights designed at each re-optimization date (NaN elsewhere).
    """
    if optimize_every % rebalance_every != 0:
        raise ValueError("optimize_every must be a multiple of rebalance_every")

    P = prices.to_numpy(dtype=float)
    X = prices.pct_change().fillna(0).to_numpy()
    T, N = P.shape

    returns = {}
    weights = {}
    for name, fun in portfolio_funcs.items():
        w_bop = np.full((T + 1, N), np.nan)      # weights at the beginning of the period
        w_eop = np.zeros((T + 1, N))             # weights at the end of the period
        NAV = np.full(T + 1, np.nan)
        w_designed = np.full((T + 1, N), np.nan)

        w_eop[lookback] = 0.0
        NAV[lookback] = 1.0
        target = np.zeros(N)

        for t in range(lookback, T):
            if t > lookback:
                if np.isnan(w_bop[t]).any():     # no rebalancing: the weights drift
                    w_bop[t] = w_eop[t - 1]
                tc = cost * np.abs(w_bop[t] - w_eop[t - 1]).sum()
                cash = 1 - w_bop[t].sum() - tc
                w_eop_ = (1 + X[t]) * w_bop[t]
                NAV[t] = (w_eop_.sum() + cash) * NAV[t - 1]
                w_eop[t] = w_eop_ / (w_eop_.sum() + cash)

            if (t - lookback) % rebalance_every == 0 and t + 1 < T:
                if (t - lookback) % optimize_every == 0:
                    window = prices.iloc[t - lookback + 1:t + 1]
                    target = np.asarray(fun(window), dtype=float)
                    w_designed[t] = target
                w_bop[t + 1] = target

        rets = NAV[lookback + 1:T] / NAV[lookback:T - 1] - 1
        returns[name] = pd.Series(rets, index=prices.index[lookback + 1:T])
        weights[name] = pd.DataFrame(w_designed[lookback:T - 1], index=prices.index[lookback:T - 1],
                                     columns=prices.columns)

    return pd.DataFrame(returns), weights


def performance_metrics(returns):
    """
    Performance statistics of a daily return series.

    Parameters
    ----------
    returns : pd.Series
        Daily returns.

    Returns
    -------
    pd.Series
        Annualized return, annualized volatility, Sharpe ratio, Sortino ratio
        and maximum drawdown.
    """
    return pd.Series({
        "annual return": emp.annual_return(returns),
        "annual volatility": emp.annual_volatility(returns),
        "Sharpe ratio": emp.sharpe_ratio(returns),
        "Sortino ratio": emp.sortino_ratio(returns),
        "max drawdown": emp.max_drawdown(returns),
    })

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

    Parameters
    ----------
    returns : pd.DataFrame
        T x K DataFrame of daily 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 plot_weights(weights, assets):
    """
    Grouped bar chart of the portfolio weights.

    Parameters
    ----------
    weights : dict[str, np.ndarray]
        Mapping from a strategy name to a length-N vector of weights.
    assets : list of str
        Names of the assets.
    """
    names = list(weights)
    N = len(assets)
    x = np.arange(N)
    width = 0.8 / len(names)

    _, ax = plt.subplots(figsize=(13, 5))
    for k, name in enumerate(names):
        ax.bar(x + (k - (len(names) - 1) / 2) * width, np.asarray(weights[name]),
               width=width, label=name, edgecolor="black", linewidth=0.4)
    ax.set_xticks(x)
    ax.set_xticklabels(assets, rotation=45, ha="right")
    ax.set_ylabel("weight")
    ax.set_title("Portfolio weights")
    ax.legend(fontsize=8)
    plt.tight_layout()
    plt.show()

def resample_datasets(prices, num_datasets, N_sample, T_sample, seed=42):
    """
    Random resamples of a price panel.

    Parameters
    ----------
    prices : pd.DataFrame
        T x N DataFrame of prices.
    num_datasets : int
        Number of resampled datasets.
    N_sample : int
        Number of stocks in each resample.
    T_sample : int
        Number of days in each resample.
    seed : int, default 42
        Seed of the random number generator.

    Returns
    -------
    list of pd.DataFrame
        The resampled price panels.
    """
    rng = np.random.default_rng(seed)
    datasets = []
    for _ in range(num_datasets):
        columns = rng.choice(prices.columns, N_sample, replace=False)
        start = rng.integers(0, len(prices) - T_sample + 1)
        datasets.append(prices.loc[:, columns].iloc[start:start + T_sample])
    return datasets


def multiple_backtests(portfolio_funcs, datasets, lookback=252, optimize_every=20,
                       rebalance_every=5, cost=60e-4, verbose=True):
    """
    Walk-forward backtest of several designs on several resampled datasets.

    Parameters
    ----------
    portfolio_funcs : dict[str, callable]
        Mapping from a strategy name to a design function.
    datasets : list of pd.DataFrame
        Resampled price panels.
    lookback : int, default 252
        Lookback window of the design functions.
    optimize_every : int, default 20
        Number of periods between re-optimizations.
    rebalance_every : int, default 5
        Number of periods between rebalancings.
    cost : float, default 60e-4
        Transaction cost per unit of traded amount.
    verbose : bool, default True
        If True, show a progress bar.

    Returns
    -------
    pd.DataFrame
        Long-format DataFrame with one row per (portfolio, dataset) and the
        performance measures as columns.
    """
    rows = []
    for b, prices in enumerate(tqdm(datasets, disable=not verbose)):
        returns, _ = backtest(portfolio_funcs, prices, lookback=lookback,
                              optimize_every=optimize_every,
                              rebalance_every=rebalance_every, cost=cost)
        for name in portfolio_funcs:
            row = performance_metrics(returns[name])
            row["portfolio"] = name
            row["dataset"] = b
            rows.append(row)
    return pd.DataFrame(rows)


def plot_metric_boxplots(results, metrics=("Sharpe ratio", "max drawdown"), figsize=(12, 6)):
    """
    Horizontal boxplots of the chosen performance measures across datasets.

    Parameters
    ----------
    results : pd.DataFrame
        Long-format DataFrame from ``multiple_backtests``.
    metrics : tuple of str
        Measures to plot.
    figsize : tuple, default (12, 6)
        Size of the figure.
    """
    _, axes = plt.subplots(len(metrics), 1, figsize=figsize)
    for ax, metric in zip(np.atleast_1d(axes), metrics):
        sns.boxplot(data=results, y="portfolio", x=metric, ax=ax, color=BLUE, width=0.6)
        ax.set_title(metric)
        ax.set_ylabel(None)
        ax.set_xlabel(None)
    plt.tight_layout()
    plt.show()

tickers = ["MGM", "WYNN", "LVS", "CMS", "DUK", "BAC", "BK", "USB",
           "KMB", "PG", "FB", "AMZN", "MSFT", "GOOGL"]

rng = np.random.default_rng(42)
stock_prices = SP500_stocks_2015to2020[tickers].tail(200)
stock_prices = stock_prices.iloc[:, rng.permutation(len(tickers))]
X = np.log(stock_prices).diff().dropna()

C = X.corr().values
Sigma = np.cov(X.values, rowvar=False)
D = hierarchical_distance(Sigma, distance="corr-dist")

fig, axes = plt.subplots(2, 2, figsize=(12, 8))
for ax, method, title in zip(axes.ravel(),
                             ["single", "complete", "average", "ward"],
                             ["Single linkage", "Complete linkage",
                              "Average linkage", "Ward's method"]):
    Z = hclust(D, method)
    dendrogram(Z, ax=ax, labels=list(stock_prices.columns), leaf_rotation=90,
               leaf_font_size=8, color_threshold=0, above_threshold_color=BLACK)
    # cut into four clusters
    height = (Z[-4, 2] + Z[-3, 2]) / 2
    ax.axhline(height, color=BLUE, linestyle="--")
    ax.set_title(title)
    ax.set_yticks([])
    ax.set_ylabel(None)
    ax.set_facecolor("white")
    ax.grid(False)
plt.tight_layout()
plt.show()

def reorder_cormat(cormat, method="single"):
    """
    Reorder a correlation matrix according to the leaves of the dendrogram.

    Parameters
    ----------
    cormat : np.ndarray
        Correlation matrix.
    method : str, default "single"
        Linkage method.

    Returns
    -------
    cormat_ordered : np.ndarray
        Reordered correlation matrix.
    order : list of int
        Order of the assets (indices of the original matrix).
    """
    D = np.sqrt(np.clip(0.5 * (1 - cormat), 0, None))
    Z = hclust(squareform(pdist(D)), method)
    order = leaves_order(Z, len(cormat))
    return cormat[np.ix_(order, order)], order


cormat = np.corrcoef(X.values, rowvar=False)
cormat_names = list(stock_prices.columns)
cormat_ordered, order = reorder_cormat(cormat)
names_ordered = [cormat_names[i] for i in order]

_, axes = plt.subplots(1, 2, figsize=(13, 5))
for ax, mat, names, title in [(axes[0], cormat, cormat_names, "Original correlation matrix"),
                              (axes[1], cormat_ordered, names_ordered,
                               "Quasi-diagonal correlation matrix")]:
    im = ax.imshow(mat, cmap="RdBu_r", vmin=-1, vmax=1)
    ax.set_xticks(range(len(names)))
    ax.set_xticklabels(names, rotation=90)
    ax.set_yticks(range(len(names)))
    ax.set_yticklabels(names)
    ax.set_title(title)
    ax.set_facecolor("white")
    ax.grid(False)
fig.colorbar(im, ax=axes, shrink=0.7)
plt.show()

def returns_from_window(prices):
    """Log-returns of a window of prices."""
    return np.log(prices).diff().dropna().values


def EWP(prices):
    """Equally weighted (1/N) portfolio."""
    N = prices.shape[1]
    return np.full(N, 1.0 / N)


def GMVP(prices):
    """Global minimum variance portfolio."""
    X_win = returns_from_window(prices)
    return design_GMVP(np.cov(X_win, rowvar=False))


def MVP(prices):
    """Mean-variance portfolio."""
    X_win = returns_from_window(prices)
    return design_MVP(X_win.mean(axis=0), np.cov(X_win, rowvar=False))


def H1overN_single(prices):
    return design_H1overN(X=returns_from_window(prices), linkage="single", distance="corr-dist")


def H1overN_complete(prices):
    return design_H1overN(X=returns_from_window(prices), linkage="complete", distance="corr-dist")


def H1overN_average(prices):
    return design_H1overN(X=returns_from_window(prices), linkage="average", distance="corr-dist")


def H1overN_Ward(prices):
    return design_H1overN(X=returns_from_window(prices), linkage="Ward", distance="corr-dist")

T = len(stock_prices)
T_trn = round(0.70 * T)

bt_linkage = {
    "1/N": EWP,
    "Hierarchical 1/N (single linkage)": H1overN_single,
    "Hierarchical 1/N (complete linkage)": H1overN_complete,
    "Hierarchical 1/N (average linkage)": H1overN_average,
    "Hierarchical 1/N (Ward's method)": H1overN_Ward,
}
returns_linkage, weights_linkage = backtest(bt_linkage, stock_prices, lookback=T_trn,
                                            optimize_every=30, rebalance_every=1)

plot_weights({name: weights_linkage[name].iloc[0].values for name in bt_linkage},
             list(stock_prices.columns))

plot_backtest(returns_linkage)

def H1overN_Ward_dist2(prices):
    return design_H1overN(X=returns_from_window(prices), linkage="Ward", distance="corr-dist2")


def H1overN_Ward_spearman(prices):
    return design_H1overN(X=returns_from_window(prices), linkage="Ward", distance="spearman")


def H1overN_Ward_kendall(prices):
    return design_H1overN(X=returns_from_window(prices), linkage="Ward", distance="kendall")


def H1overN_Ward_distcorr(prices):
    return design_H1overN(X=returns_from_window(prices), linkage="Ward", distance="distance-corr")

bt_distance = {
    "1/N": EWP,
    "Hierarchical 1/N (corr-distance)": H1overN_Ward,
    "Hierarchical 1/N (corr-distance-of-distance)": H1overN_Ward_dist2,
    "Hierarchical 1/N (Spearman)": H1overN_Ward_spearman,
    "Hierarchical 1/N (Kendall)": H1overN_Ward_kendall,
    "Hierarchical 1/N (distance correlation)": H1overN_Ward_distcorr,
}
returns_distance, weights_distance = backtest(bt_distance, stock_prices, lookback=T_trn,
                                              optimize_every=30, rebalance_every=1)

plot_weights({name: weights_distance[name].iloc[0].values for name in bt_distance},
             list(stock_prices.columns))

plot_backtest(returns_distance)

bt_final = {
    "1/N": EWP,
    "GMVP": GMVP,
    "MVP": MVP,
    "Hierarchical 1/N": H1overN_Ward_dist2,
}
returns_final, weights_final = backtest(bt_final, stock_prices, lookback=T_trn,
                                        optimize_every=30, rebalance_every=1)

plot_weights({name: weights_final[name].iloc[0].values for name in bt_final},
             list(stock_prices.columns))

plot_backtest(returns_final)

def IVarP(prices):
    """Inverse-variance portfolio."""
    X_win = returns_from_window(prices)
    return design_IVarP(np.cov(X_win, rowvar=False))


def HRP_bisection(prices):
    return design_HRP(X=returns_from_window(prices), splitting="bisection", linkage="single",
                      distance="corr-dist2")


def HRP_dendrogram(prices):
    return design_HRP(X=returns_from_window(prices), splitting="dendrogram", linkage="single",
                      distance="corr-dist2")


def HRP_dendrogram_spearman(prices):
    return design_HRP(X=returns_from_window(prices), splitting="dendrogram", linkage="single",
                      distance="spearman")


def HRP_dendrogram_distcorr(prices):
    return design_HRP(X=returns_from_window(prices), splitting="dendrogram", linkage="single",
                      distance="distance-corr")

bt_hrp = {
    "GMVP": GMVP,
    "IVarP": IVarP,
    "HRP-bisection": HRP_bisection,
    "HRP-dendrogram": HRP_dendrogram,
    "HRP-dendrogram (Spearman)": HRP_dendrogram_spearman,
    "HRP-dendrogram (distance correlation)": HRP_dendrogram_distcorr,
}
returns_hrp, weights_hrp = backtest(bt_hrp, stock_prices, lookback=T_trn,
                                    optimize_every=30, rebalance_every=1)

plot_weights({name: weights_hrp[name].iloc[0].values for name in bt_hrp},
             list(stock_prices.columns))

plot_backtest(returns_hrp)

def H1overN_Ward_dist2_dendrogram(prices):
    return design_H1overN(X=returns_from_window(prices), splitting="dendrogram", linkage="Ward",
                          distance="corr-dist2")


def HERC_1overN_dendrogram(prices):
    return design_H1overN(X=returns_from_window(prices), splitting="dendrogram", num_clusters=4,
                          linkage="Ward", distance="corr-dist2")


def HERC_1overN_bisection(prices):
    return design_H1overN(X=returns_from_window(prices), splitting="bisection", num_clusters=4,
                          linkage="Ward", distance="corr-dist2")


def HERC_IVarP_dendrogram(prices):
    return design_HRP(X=returns_from_window(prices), splitting="dendrogram", num_clusters=4,
                      linkage="Ward", distance="corr-dist2", final="1/N")


def HERC_IVarP_bisection(prices):
    return design_HRP(X=returns_from_window(prices), splitting="bisection", num_clusters=4,
                      linkage="Ward", distance="corr-dist2", final="1/N")

bt_herc = {
    "1/N": EWP,
    "Hierarchical 1/N": H1overN_Ward_dist2_dendrogram,
    "HERC-dendrogram 1/N": HERC_1overN_dendrogram,
    "HERC-bisection 1/N": HERC_1overN_bisection,
    "IVarP": IVarP,
    "HRP-bisection": HRP_bisection,
    "HERC-dendrogram IVarP": HERC_IVarP_dendrogram,
    "HERC-bisection IVarP": HERC_IVarP_bisection,
}
returns_herc, weights_herc = backtest(bt_herc, stock_prices, lookback=T_trn,
                                      optimize_every=30, rebalance_every=1)

plot_weights({name: weights_herc[name].iloc[0].values for name in bt_herc},
             list(stock_prices.columns))

plot_backtest(returns_herc)

stock_panel = SP500_stocks_2015to2020

datasets = resample_datasets(stock_panel, num_datasets=100, N_sample=50, T_sample=252 * 2, seed=42)
print(f"{len(datasets)} datasets of {datasets[0].shape[0]} days and {datasets[0].shape[1]} stocks")

def H1overN_bisection(prices):
    return design_H1overN(X=returns_from_window(prices), splitting="bisection", linkage="Ward",
                          distance="corr-dist2")


def H1overN_dendrogram(prices):
    return design_H1overN(X=returns_from_window(prices), splitting="dendrogram", linkage="Ward",
                          distance="corr-dist2")


def HRP_bisection_single(prices):
    return design_HRP(X=returns_from_window(prices), splitting="bisection", linkage="single",
                      distance="corr-dist2")


def HRP_dendrogram_single(prices):
    return design_HRP(X=returns_from_window(prices), splitting="dendrogram", linkage="single",
                      distance="corr-dist2")


def HERC_1overN_bisection_10(prices):
    return design_H1overN(X=returns_from_window(prices), splitting="bisection", num_clusters=10,
                          linkage="Ward", distance="corr-dist2")


def HERC_1overN_dendrogram_10(prices):
    return design_H1overN(X=returns_from_window(prices), splitting="dendrogram", num_clusters=10,
                          linkage="Ward", distance="corr-dist2")


def HERC_IVarP_bisection_10(prices):
    return design_HRP(X=returns_from_window(prices), splitting="bisection", num_clusters=10,
                      linkage="Ward", distance="corr-dist2", final="1/N")


def HERC_IVarP_dendrogram_10(prices):
    return design_HRP(X=returns_from_window(prices), splitting="dendrogram", num_clusters=10,
                      linkage="Ward", distance="corr-dist2", final="1/N")


bt_splitting = {
    "1/N": EWP,
    "Hierarchical 1/N - bisection": H1overN_bisection,
    "Hierarchical 1/N - dendrogram": H1overN_dendrogram,
    "HERC 1/N - bisection": HERC_1overN_bisection_10,
    "HERC 1/N - dendrogram": HERC_1overN_dendrogram_10,
    "IVarP": IVarP,
    "HRP - bisection": HRP_bisection_single,
    "HRP - dendrogram": HRP_dendrogram_single,
    "HERC IVarP - bisection": HERC_IVarP_bisection_10,
    "HERC IVarP - dendrogram": HERC_IVarP_dendrogram_10,
}

results_splitting = multiple_backtests(bt_splitting, datasets, lookback=252, optimize_every=20,
                                       rebalance_every=5, cost=60e-4, verbose=False)

plot_metric_boxplots(results_splitting, figsize=(12, 7))

def H1overN_spearman(prices):
    return design_H1overN(X=returns_from_window(prices), splitting="bisection", linkage="Ward",
                          distance="spearman")


def H1overN_kendall(prices):
    return design_H1overN(X=returns_from_window(prices), splitting="bisection", linkage="Ward",
                          distance="kendall")


def HRP_spearman(prices):
    return design_HRP(X=returns_from_window(prices), splitting="bisection", linkage="single",
                      distance="spearman")


def HRP_kendall(prices):
    return design_HRP(X=returns_from_window(prices), splitting="bisection", linkage="single",
                      distance="kendall")


bt_graph = {
    "1/N": EWP,
    "Hierarchical 1/N (Pearson)": H1overN_bisection,
    "Hierarchical 1/N (Spearman)": H1overN_spearman,
    "Hierarchical 1/N (Kendall)": H1overN_kendall,
    "IVarP": IVarP,
    "HRP (Pearson)": HRP_bisection_single,
    "HRP (Spearman)": HRP_spearman,
    "HRP (Kendall)": HRP_kendall,
}

datasets_graph = resample_datasets(stock_panel, num_datasets=50, N_sample=50, T_sample=252 * 2,
                                   seed=42)

results_graph = multiple_backtests(bt_graph, datasets_graph, lookback=252, optimize_every=20,
                                   rebalance_every=5, cost=60e-4, verbose=False)

plot_metric_boxplots(results_graph, figsize=(12, 6))

bt_comparison = {
    "1/N": EWP,
    "Hierarchical 1/N": H1overN_bisection,
    "HERC 1/N": HERC_1overN_bisection_10,
    "IVarP": IVarP,
    "HRP": HRP_bisection_single,
    "HERC IVarP": HERC_IVarP_bisection_10,
}

results_comparison = multiple_backtests(bt_comparison, datasets, lookback=252, optimize_every=20,
                                        rebalance_every=5, cost=60e-4, verbose=False)

measures = ["Sharpe ratio", "annual return", "annual volatility", "Sortino ratio", "max drawdown"]
summary = results_comparison.groupby("portfolio")[measures].median().loc[list(bt_comparison)]
summary.round(3)

summary[["max drawdown", "annual volatility"]].plot(kind="bar", figsize=(12, 4), rot=20)
plt.ylabel(None)
plt.legend(title="Portfolio")
plt.title("Median performance across datasets")
plt.tight_layout()
plt.show()

plot_metric_boxplots(results_comparison, figsize=(12, 6))