"""Hierarchical Risk Parity (López de Prado 2016). corr → distance → hierarchical linkage →
quasi-diagonalization → recursive bisection with inverse-variance allocation. No matrix inversion, no
concentrated weights. Falls back to inverse-variance on a degenerate/singular correlation."""
from __future__ import annotations

import numpy as np


def _ivp(cov: np.ndarray) -> np.ndarray:
    with np.errstate(divide="ignore", invalid="ignore"):
        iv = 1.0 / np.diag(cov)
        return iv / iv.sum()


def _cluster_var(cov: np.ndarray, items: list[int]) -> float:
    sub = cov[np.ix_(items, items)]
    w = _ivp(sub).reshape(-1, 1)
    return float((w.T @ sub @ w)[0, 0])


def hrp_weights(returns: dict[str, np.ndarray]) -> dict[str, float]:
    syms = list(returns)
    if len(syms) == 1:
        return {syms[0]: 1.0}
    n = min(len(r) for r in returns.values())
    mat = np.array([returns[s][-n:] for s in syms])          # rows = assets
    cov = np.cov(mat)
    if cov.ndim == 0:
        return {syms[0]: 1.0}
    std = np.sqrt(np.diag(cov))
    denom = np.outer(std, std)
    with np.errstate(divide="ignore", invalid="ignore"):
        corr = np.where(denom > 0, cov / denom, 0.0)
    np.fill_diagonal(corr, 1.0)
    try:
        from scipy.cluster.hierarchy import linkage, to_tree
        from scipy.spatial.distance import squareform
        dist = np.sqrt(np.clip((1.0 - corr) / 2.0, 0.0, 1.0))
        link = linkage(squareform(dist, checks=False), method="single")
        # quasi-diagonalization via scipy's tree pre-order traversal (no pandas)
        order = to_tree(link).pre_order()
        w = np.ones(len(syms))
        clusters = [order]
        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 i in range(0, len(clusters), 2):
                c0, c1 = clusters[i], clusters[i + 1]
                v0, v1 = _cluster_var(cov, c0), _cluster_var(cov, c1)
                alpha = 1.0 - v0 / (v0 + v1) if (v0 + v1) > 0 else 0.5
                w[c0] *= alpha
                w[c1] *= 1.0 - alpha
        w = w / w.sum()
        if not np.all(np.isfinite(w)) or np.any(w < 0) or w.sum() <= 0:
            w = np.full(len(syms), 1.0 / len(syms))      # last-resort uniform (always valid)
        return {syms[i]: float(w[i]) for i in range(len(syms))}
    except Exception:                                         # noqa: BLE001 — degenerate -> inverse-variance
        w = _ivp(cov)
        if not np.all(np.isfinite(w)) or np.any(w < 0) or w.sum() <= 0:
            w = np.full(len(syms), 1.0 / len(syms))      # last-resort uniform (always valid)
        return {syms[i]: float(w[i]) for i in range(len(syms))}
