!pip install numpy pandas matplotlib seaborn scipy yfinance --quiet

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.cluster.hierarchy import linkage, dendrogram, leaves_list
from scipy.spatial.distance import squareform
import warnings

warnings.filterwarnings('ignore')
%matplotlib inline
plt.rcParams['figure.figsize'] = (14, 5)
plt.rcParams['axes.spines.top'] = False
plt.rcParams['axes.spines.right'] = False
print('Imports ready.')

# ── Data Source Toggle ─────────────────────────────
USE_LIVE_DATA = False
TICKERS       = ['BTC-USD','ETH-USD','SOL-USD','BNB-USD','AVAX-USD','MATIC-USD','DOT-USD','LINK-USD']
START_DATE    = '2022-01-01'
END_DATE      = '2024-12-31'
# ──────────────────────────────────────────────────
ASSETS = ['BTC','ETH','SOL','BNB','AVAX','MATIC','DOT','LINK']
print('Config ready.')

if USE_LIVE_DATA:
    import yfinance as yf
    prices  = yf.download(TICKERS, start=START_DATE, end=END_DATE)['Close']
    prices.columns = ASSETS
    prices.dropna(inplace=True)
    returns = prices.pct_change().dropna()
    print(f'Live data: {len(returns)} days')
else:
    rng = np.random.default_rng(42)
    n   = 730
    n_a = len(ASSETS)
    # Create two clusters: large-caps [BTC,ETH,BNB] and alt-caps [SOL,AVAX,MATIC,DOT,LINK]
    corr = np.full((n_a, n_a), 0.50)
    np.fill_diagonal(corr, 1.0)
    corr[0,1]=corr[1,0]=0.85; corr[0,3]=corr[3,0]=0.75; corr[1,3]=corr[3,1]=0.72  # large-cap cluster
    corr[2,4]=corr[4,2]=0.80; corr[2,5]=corr[5,2]=0.78; corr[4,5]=corr[5,4]=0.82  # alt-cap cluster
    vols = np.array([0.65,0.75,1.20,0.70,1.10,1.40,1.00,1.20]) / np.sqrt(252)
    cov  = np.outer(vols,vols)*corr
    L    = np.linalg.cholesky(cov)
    mu   = np.array([0.50,0.45,0.80,0.35,0.70,0.90,0.55,0.75])/252
    data = rng.standard_normal((n,n_a)) @ L.T + mu
    idx  = pd.date_range('2022-01-01', periods=n, freq='B')
    returns = pd.DataFrame(data, columns=ASSETS, index=idx)
    print(f'Synthetic data: {len(returns)} days')

print(f'Assets: {len(ASSETS)}')

def hrp_weights(returns: pd.DataFrame) -> pd.Series:
    """
    Compute Hierarchical Risk Parity weights.

    Parameters
    ----------
    returns : pd.DataFrame  Asset return series.

    Returns
    -------
    pd.Series  Portfolio weights indexed by asset name.
    """
    # Step 1: Correlation distance matrix
    corr = returns.corr()
    dist = np.sqrt((1 - corr) / 2)
    dist_condensed = squareform(dist.values, checks=False)

    # Step 2: Hierarchical clustering
    link = linkage(dist_condensed, method='single')

    # Step 3: Quasi-diagonalisation — reorder assets
    sort_ix = leaves_list(link)
    sorted_assets = [returns.columns[i] for i in sort_ix]
    cov_sorted = returns[sorted_assets].cov()

    # Step 4: Recursive bisection
    weights = pd.Series(1.0, index=sorted_assets)
    clusters = [sorted_assets]
    while clusters:
        clusters = [c[j:k] for c in clusters
                    for j, k in ((0, len(c) // 2), (len(c) // 2, len(c)))
                    if len(c) > 1]
        for c in [c for c in clusters if len(c) > 1]:
            pass
        # Inverse variance allocation within each sub-cluster
        all_items = [item for cluster in clusters for item in cluster]
        for i in range(0, len(all_items) - 1, 2):
            left  = [all_items[i]] if i < len(all_items) else []
            right = [all_items[i+1]] if i+1 < len(all_items) else []
            if not left or not right:
                continue
            var_l = cov_sorted.loc[left, left].values.sum()
            var_r = cov_sorted.loc[right, right].values.sum()
            alpha = 1 - var_l / (var_l + var_r + 1e-9)
            weights[left]  *= alpha
            weights[right] *= (1 - alpha)
        break

    # Fallback: simple inverse variance
    cov_all = returns.cov()
    inv_var = 1 / np.diag(cov_all.values)
    w = pd.Series(inv_var / inv_var.sum(), index=returns.columns)
    return w


# Proper HRP implementation
corr = returns.corr()
cov  = returns.cov()
dist = np.sqrt((1 - corr) / 2)

link      = linkage(squareform(dist.values, checks=False), method='ward')
sort_ix   = leaves_list(link)
sorted_assets = [ASSETS[i] for i in sort_ix]

# Recursive bisection on sorted assets
def get_ivp(cov, assets):
    iv = 1 / np.diag(cov.loc[assets, assets].values)
    return iv / iv.sum()

def recursive_bisection(cov, sort_idx):
    w = pd.Series(1.0, index=sort_idx)
    cluster_items = [sort_idx]
    while len(cluster_items) > 0:
        cluster_items = [i[j:k] for i in cluster_items
                          for j, k in ((0, len(i)//2),(len(i)//2, len(i)))
                          if len(i) > 1]
        for i in range(0, len(cluster_items), 2):
            if i+1 >= len(cluster_items): break
            c_left  = cluster_items[i]
            c_right = cluster_items[i+1]
            var_l = get_ivp(cov, c_left)
            var_r = get_ivp(cov, c_right)
            cov_l = (get_ivp(cov, c_left) * np.diag(cov.loc[c_left, c_left])).sum()
            cov_r = (get_ivp(cov, c_right) * np.diag(cov.loc[c_right, c_right])).sum()
            alpha = 1 - cov_l / (cov_l + cov_r + 1e-9)
            w[c_left]  *= alpha
            w[c_right] *= 1 - alpha
    return w

w_hrp = recursive_bisection(cov, sorted_assets)
w_hrp = w_hrp.reindex(ASSETS).fillna(0)
w_hrp = w_hrp / w_hrp.sum()

print('HRP Weights:')
print((w_hrp * 100).round(2).to_string())

fig, axes = plt.subplots(1, 2, figsize=(14, 5))
fig.suptitle('Hierarchical Risk Parity', fontsize=13, fontweight='bold')

ax1 = axes[0]
dendrogram(link, labels=[ASSETS[i] for i in sort_ix], ax=ax1,
           color_threshold=0.5, above_threshold_color='#9e9e9e')
ax1.set_title('Asset Dendrogram')
ax1.set_ylabel('Distance')

ax2 = axes[1]
w_ew = pd.Series(1/len(ASSETS), index=ASSETS)
x = np.arange(len(ASSETS))
ax2.bar(x - 0.2, w_ew * 100, 0.4, label='Equal Weight', color='#9e9e9e', alpha=0.8)
ax2.bar(x + 0.2, w_hrp * 100, 0.4, label='HRP', color='#1976d2', alpha=0.8)
ax2.set_xticks(x); ax2.set_xticklabels(ASSETS, rotation=45)
ax2.set_ylabel('Weight (%)')
ax2.legend(fontsize=8)
ax2.set_title('HRP vs Equal Weight')

plt.tight_layout(); plt.show()

pd.DataFrame({'Asset': ASSETS, 'HRP_Weight': w_hrp.values.round(4)}).to_csv(
    'hierarchical_risk_parity.csv', index=False)
print('Saved: hierarchical_risk_parity.csv')