"""
Canopy Core — Cluster Engine: Distance, Linkage, Seriation, Cluster Variance
Copyright © 2026 Anagatam Technologies. All rights reserved.

This module provides the fundamental mathematical building blocks required by
all three hierarchical allocation algorithms (HRP, HERC, NCO). It implements:

    1. Correlation-to-Distance Transformations (Codependence Metrics)
    2. Hierarchical Linkage Computation with Optimal Leaf Ordering
    3. Quasi-Diagonalization (Tree Seriation) via SciPy ClusterNode
    4. Cluster Variance Estimation under Naive Risk Parity

Mathematical Foundation:
    The key insight underlying hierarchical portfolio methods is that the sample
    correlation matrix ρ can be converted into a proper metric space (satisfying
    the triangle inequality) via distance transformations. This enables the
    application of agglomerative clustering algorithms (Ward, Single, Complete,
    Average) to discover the latent block-diagonal structure of asset returns.

    The quasi-diagonalization step (seriation) reorders assets so that correlated
    assets are adjacent in the allocation vector, which is critical for the
    recursive bisection procedure in HRP (Lopez de Prado, 2016).

Codependence Metrics Supported:
    - 'pearson'      : Standard angular distance, d = √(½(1−ρ))
    - 'abs_pearson'  : Sign-agnostic distance, d = √(1−|ρ|)
    - 'spearman'     : Rank-based distance (robust to outliers)
    - 'kendall'      : Ordinal association distance (robust to non-linearity)

    → Canopy supports Pearson, abs_Pearson, Spearman, and Kendall —
      covering the most commonly used codependence metrics in academic
      literature and institutional practice.

Numerical Stability:
    All functions include guards against:
    - Floating-point correlation values outside [-1, 1]
    - Zero-variance assets (diagonal elements ≈ 0 in covariance matrices)
    - Asymmetric distance matrices due to IEEE 754 rounding
    - Degenerate single-asset clusters

References:
    [1] Lopez de Prado, M. (2016). "Building Diversified Portfolios that
        Outperform Out of Sample." Journal of Portfolio Management, 42(4).
    [2] Mantegna, R.N. (1999). "Hierarchical Structure in Financial Markets."
        The European Physical Journal B, 11(1), 193-197.
    [3] Ward, J.H. (1963). "Hierarchical Grouping to Optimize an Objective
        Function." Journal of the American Statistical Association, 58(301).
    [4] Bar-Joseph, Z. et al. (2001). "Fast optimal leaf ordering for
        hierarchical clustering." Bioinformatics, 17(suppl_1), S22-S29.

Complexity Analysis:
    - correl_dist:     O(N²)     where N = number of assets
    - compute_linkage: O(N² log N) via scipy's optimized agglomeration
    - get_quasi_diag:  O(N)      tree pre-order traversal via ClusterNode
    - get_cluster_var: O(K³)     where K = cluster size (matrix multiply)
"""

import numpy as np
import pandas as pd
from scipy.cluster.hierarchy import linkage, to_tree, optimal_leaf_ordering
from scipy.spatial.distance import squareform


# ---------------------------------------------------------------------------
# 1. CORRELATION → DISTANCE TRANSFORMATION
# ---------------------------------------------------------------------------

def correl_dist(corr: pd.DataFrame, method: str = 'pearson') -> pd.DataFrame:
    """
    Transforms a correlation matrix into a proper distance matrix suitable
    for hierarchical clustering.

    The transformation maps pairwise correlations ρ_{i,j} ∈ [-1, 1] into
    distances d_{i,j} ∈ [0, 1] that satisfy the properties of a metric space
    (non-negativity, identity of indiscernibles, symmetry, triangle inequality).

    Supported codependence metrics:
        'pearson':
            d_{i,j} = √(½ · (1 − ρ_{i,j}))
            Standard angular distance (Lopez de Prado, 2016; PyPortfolioOpt).
            ρ = +1 → d = 0, ρ = 0 → d ≈ 0.707, ρ = -1 → d = 1.

        'abs_pearson':
            d_{i,j} = √(1 − |ρ_{i,j}|)
            Sign-agnostic distance.
            Appropriate when short-selling is permitted.
            |ρ| = 1 → d = 0, |ρ| = 0 → d = 1.

        'spearman':
            Computes Spearman rank correlation internally, then applies
            the standard angular distance transformation.
            Robust to outliers and non-linear monotonic relationships.

        'kendall':
            Computes Kendall τ rank correlation internally, then applies
            the standard angular distance transformation.
            Robust to ties and ordinal data.

    Numerical Safeguards:
        - Clips correlation values to [-1, 1] to handle floating-point drift
        - Clips the argument of √ to [0, 1] to prevent NaN from negative rounding
        - Forces diagonal to exactly 0 (self-distance)

    Note:
        For 'spearman' and 'kendall', this function expects the raw returns
        DataFrame to be accessible so it can compute the rank correlation
        internally. However, since the facade (MasterCanopy) handles this
        by computing the appropriate correlation matrix before calling
        correl_dist, the corr argument already contains the Spearman/Kendall
        matrix. The method parameter here controls the distance formula.

    Args:
        corr: Square correlation matrix (N × N) with asset names as index/columns.
              For 'spearman'/'kendall', this should already be the rank correlation.
        method: Codependence metric — 'pearson', 'abs_pearson', 'spearman', 'kendall'.

    Returns:
        Distance matrix (N × N) as a pandas DataFrame preserving asset labels.

    Raises:
        ValueError: If an unsupported codependence method is requested.
    """
    _SUPPORTED = {'pearson', 'abs_pearson', 'spearman', 'kendall'}

    # Guard: ensure correlation values are bounded due to IEEE 754 precision
    corr_clipped = np.clip(corr.values, -1.0, 1.0)

    if method == 'abs_pearson':
        # Absolute correlation distance (symmetric w.r.t. sign of ρ)
        dist_matrix = np.sqrt(np.clip(1.0 - np.abs(corr_clipped), 0.0, 1.0))
    elif method in ('pearson', 'spearman', 'kendall'):
        # Standard angular distance — applied to Pearson, Spearman, or Kendall ρ
        dist_matrix = np.sqrt(np.clip(0.5 * (1.0 - corr_clipped), 0.0, 1.0))
    else:
        raise ValueError(
            f"Unsupported codependence method '{method}'. "
            f"Supported: {_SUPPORTED}"
        )

    # Enforce exact zero on diagonal (self-distance identity)
    np.fill_diagonal(dist_matrix, 0.0)

    return pd.DataFrame(dist_matrix, index=corr.columns, columns=corr.columns)


# ---------------------------------------------------------------------------
# 2. HIERARCHICAL LINKAGE COMPUTATION
# ---------------------------------------------------------------------------

def compute_linkage(dist_matrix: pd.DataFrame, method: str = 'ward',
                    optimal_ordering: bool = True) -> np.ndarray:
    """
    Computes the agglomerative hierarchical clustering linkage matrix from
    a precomputed distance matrix, with optional optimal leaf ordering.

    The linkage matrix Z is an (N-1) × 4 array where each row represents
    a merge operation: [cluster_i, cluster_j, distance, count]. This is
    the standard output format of scipy.cluster.hierarchy.linkage.

    Optimal Leaf Ordering (Bar-Joseph et al., 2001):
        When optimal_ordering=True, the leaf order of the dendrogram is
        optimized to minimize the sum of distances between adjacent leaves.
        This produces a more interpretable dendrogram and is used by
        Riskfolio-Lib (leaf_order=True). PyPortfolioOpt does NOT use this.
        → Canopy uses optimal leaf ordering by default for superior tree quality.

    Linkage Methods:
        'ward'    : Minimizes total within-cluster variance (Ward, 1963).
                    Produces compact, spherical clusters. Default in literature.
        'single'  : Nearest-point algorithm. Sensitive to noise ("chaining").
        'complete': Farthest-point algorithm. Produces tight clusters.
        'average' : UPGMA — Unweighted Pair Group Method with Arithmetic Mean.

    Numerical Safeguards:
        - Enforces exact symmetry via (D + D^T) / 2 before condensing
        - Uses checks=False in squareform to bypass scipy's symmetry assertion

    Args:
        dist_matrix: Square distance matrix (N × N) as pandas DataFrame.
        method: Agglomerative linkage criterion.
        optimal_ordering: If True, applies optimal leaf ordering to the
                         dendrogram (default: True). Increases computation
                         from O(N² log N) to O(N³) but yields superior tree.

    Returns:
        Linkage matrix Z of shape (N-1, 4) as numpy ndarray.

    Complexity:
        O(N² log N) without optimal ordering; O(N³) with.
    """
    dist_arr = dist_matrix.values

    # Enforce perfect symmetry to prevent squareform errors
    dist_sym = (dist_arr + dist_arr.T) / 2.0

    # Convert to condensed (upper-triangular flat) format for scipy
    condensed_dist = squareform(dist_sym, checks=False)

    Z = linkage(condensed_dist, method=method)

    # Optimal leaf ordering: minimizes sum of adjacent-leaf distances
    # (Bar-Joseph et al., 2001).
    if optimal_ordering:
        try:
            Z = optimal_leaf_ordering(Z, condensed_dist)
        except Exception:
            pass  # Graceful fallback if ordering fails (e.g., tiny N)

    return Z


# ---------------------------------------------------------------------------
# 3. QUASI-DIAGONALIZATION (TREE SERIATION)
# ---------------------------------------------------------------------------

def get_quasi_diag(link: np.ndarray) -> list:
    """
    Performs quasi-diagonalization (seriation) of the hierarchical clustering
    tree to reorder assets so that correlated items are adjacent.

    This is a depth-first pre-order traversal of the dendrogram tree that
    produces a leaf ordering consistent with the hierarchical structure.
    The resulting permutation σ quasi-diagonalizes the correlation matrix:

        C̃ = P_σ · C · P_σ^T

    Implementation Note:
        We use scipy.cluster.hierarchy.to_tree() to convert the linkage
        matrix into a ClusterNode tree, then call .pre_order() for a clean
        depth-first traversal. This is more robust than the manual
        iterative expansion used in earlier implementations,
        which can have subtle indexing issues with large asset universes.

    Args:
        link: Linkage matrix Z of shape (N-1, 4) from compute_linkage().

    Returns:
        List of original asset indices in seriated (quasi-diagonal) order.

    Complexity:
        O(N) — single pre-order traversal of the binary tree.
    """
    # Convert linkage matrix to a tree and extract leaf ordering
    # This is cleaner than the iterative approach
    root_node = to_tree(link, rd=False)
    return root_node.pre_order()


# ---------------------------------------------------------------------------
# 4. CLUSTER VARIANCE ESTIMATION
# ---------------------------------------------------------------------------

_EPSILON_VARIANCE = 1e-10  # Floor for zero-variance assets

def get_cluster_var(cov: pd.DataFrame, c_items: list) -> float:
    """
    Computes the variance of a cluster under the Naive Risk Parity (inverse
    variance) weighting scheme.

    For a cluster C of K assets, the inverse-variance weights are:

        w_i = (1 / σ²_i) / Σ_{j ∈ C} (1 / σ²_j)

    and the cluster variance is:

        Var_C = w^T · Σ_C · w

    where Σ_C is the K × K covariance sub-matrix of the cluster.

    This quantity represents the irreducible risk of the cluster when
    capital is allocated proportionally to inverse individual variance.
    It is used by HRP during recursive bisection to determine the split
    ratio α between left and right sub-trees.

    Numerical Safeguards:
        - Applies an epsilon floor (1e-10) to individual variances to prevent
          division-by-zero for constant-price assets (σ² ≈ 0).
        - Handles single-asset clusters gracefully (returns the asset's variance).

    Args:
        cov: Full covariance matrix (N × N) as pandas DataFrame.
        c_items: List of asset names (column labels) belonging to the cluster.

    Returns:
        Scalar cluster variance (float).

    Complexity:
        O(K²) for the matrix-vector product, where K = |c_items|.
    """
    cov_slice = cov.loc[c_items, c_items]

    # Extract diagonal variances with epsilon floor for numerical stability
    diag_var = np.maximum(np.diag(cov_slice), _EPSILON_VARIANCE)

    # Inverse-variance weights, normalized to sum to 1
    w = 1.0 / diag_var
    w /= w.sum()

    # Cluster variance = w^T Σ w
    cluster_var = float(np.dot(w, np.dot(cov_slice.values, w)))

    return max(cluster_var, _EPSILON_VARIANCE)
