"""
╔══════════════════════════════════════════════════════════════════════════════╗
║  Canopy — Nested Clustered Optimization (NCO) Allocator                     ║
║  Algorithm: Two-Level Markowitz with Tikhonov Regularization                ║
║                                                                              ║
║  Copyright © 2026 Anagatam Technologies. All rights reserved.               ║
╚══════════════════════════════════════════════════════════════════════════════╝

Algorithm Overview:
───────────────────
NCO (Lopez de Prado, 2019) addresses the "curse of dimensionality" in
Markowitz optimization by solving many small, well-conditioned problems
instead of one large, ill-conditioned one.

    Standard Markowitz:
        Solve ONE N×N optimization  →  κ(Σ_{N×N}) ≈ 10⁴ to 10⁸
        Error amplification: ‖Δw‖ ≤ κ(Σ) · ‖Δr‖  →  catastrophic

    NCO (Two-Level):
        Level 1: Solve K small (N/K × N/K) problems  →  κ(Σ_cluster) ≈ 10¹ to 10³
        Level 2: Solve ONE K×K aggregation problem    →  κ(Σ_reduced) ≈ 10¹ to 10²
        Error amplification: dramatically reduced.

    Mathematical Justification (Random Matrix Theory):
    ────────────────────────────────────────────────────
    By the Marchenko-Pastur law, for a random N×T matrix:
        κ(Σ̂) ∝ ((1 + √(N/T)) / (1 - √(N/T)))²

    For N/T = 0.5:  κ ≈ 34
    For N/T = 0.1:  κ ≈ 2.5

    Since each cluster has N/K assets, the effective ratio drops from
    N/T to (N/K)/T = N/(K·T), dramatically improving conditioning.
    With K=5 clusters: condition number drops by O(K²) ≈ 25×.

Three-Tier Regularization Strategy:
─────────────────────────────────────
When a cluster covariance matrix Σ_C is ill-conditioned:

    Tier 1: Check condition number κ(Σ_C)
            If κ < 10⁸ → standard inversion (well-conditioned)

    Tier 2: Apply Tikhonov regularization
            Σ_reg = Σ_C + ε·I,  where ε = λ_max(Σ_C) / 10⁴
            This shrinks eigenvalues toward ε, reducing condition number.
            (Ledoit & Wolf, 2004; used by risk desks at JPM, MS, GS)

    Tier 3: Moore-Penrose pseudo-inverse (np.linalg.pinv)
            If regularization still fails, use the pseudo-inverse which
            projects out the null space. Always produces a finite answer.

    Why ε = λ_max / 10⁴?
    This is the Marchenko-Pastur-inspired floor: we regularize noise
    eigenvalues to be at most 10⁴× smaller than the signal eigenvalue,
    which corresponds to ~4 digits of effective precision — sufficient
    for portfolio optimization where inputs are estimated to ~2 digits.

Speed Characteristics:
    - K cluster inversions of (N/K × N/K): O(K · (N/K)³) = O(N³/K²)
    - For K=5, N=500: O(500³/25) = O(5M) vs Markowitz O(125M) = 25× faster
    - Aggregation: O(K³) = O(125) — negligible

Numerical Properties:
    - NCO may produce negative weights (short positions allowed)
    - Weights are normalized to sum to 1.0
    - More concentrated than HRP/HERC (lower effective N)

References:
    [1] Lopez de Prado, M. (2019). "A Robust Estimator of the Efficient
        Frontier." SSRN Working Paper No. 3469961.
    [2] Ledoit, O. & Wolf, M. (2004). "A well-conditioned estimator for
        large-dimensional covariance matrices." Journal of Multivariate
        Analysis, 88(2), 365-411.
    [3] Marchenko, V. & Pastur, L. (1967). "Distribution of eigenvalues
        for some sets of random matrices." Mathematics of the USSR-Sbornik.
"""

import numpy as np
import pandas as pd
from scipy.cluster.hierarchy import fcluster


# ── NUMERICAL CONSTANTS ──────────────────────────────────────────────────
_CONDITION_NUMBER_THRESHOLD = 1e8  # Above this → apply regularization
_TIKHONOV_SCALE = 1e-4            # ε = λ_max · this scale factor


def get_min_var_portfolio(cov: pd.DataFrame) -> pd.Series:
    """
    Computes the global minimum-variance portfolio with three-tier
    regularization for numerical stability.

    Mathematical Derivation:
    ────────────────────────
    The minimum-variance portfolio solves:
        min_{w}  w^T · Σ · w    s.t.  1^T · w = 1

    Analytical solution (Lagrange multiplier method):
        w* = Σ⁻¹ · 1 / (1^T · Σ⁻¹ · 1)

    Proof:
        L(w, λ) = ½ w^T Σ w − λ(1^T w − 1)
        ∂L/∂w = Σw − λ1 = 0  →  w = λΣ⁻¹·1
        Substituting into constraint: 1^T · λΣ⁻¹·1 = 1  →  λ = 1/(1^T Σ⁻¹ 1)
        Therefore: w* = Σ⁻¹·1 / (1^T·Σ⁻¹·1)  ∎

    Speed: O(N³) for matrix inversion, but N is small (N/K per cluster).

    Args:
        cov: Covariance matrix of cluster assets.

    Returns:
        pd.Series of minimum-variance weights (may have negative values).
    """
    n = cov.shape[0]
    cov_arr = cov.values
    ones = np.ones(n)

    # Tier 1: Check condition number
    cond = np.linalg.cond(cov_arr)

    if cond < _CONDITION_NUMBER_THRESHOLD:
        # Well-conditioned: standard inversion
        try:
            cov_inv = np.linalg.inv(cov_arr)
            w = cov_inv @ ones
            w /= np.sum(w)
            return pd.Series(w, index=cov.columns)
        except np.linalg.LinAlgError:
            pass  # Fall through to Tier 2

    # Tier 2: Tikhonov regularization
    try:
        max_eigenvalue = np.max(np.linalg.eigvalsh(cov_arr))
        epsilon = max_eigenvalue * _TIKHONOV_SCALE
        cov_reg = cov_arr + epsilon * np.eye(n)
        cov_inv = np.linalg.inv(cov_reg)
        w = cov_inv @ ones
        w /= np.sum(w)
        return pd.Series(w, index=cov.columns)
    except np.linalg.LinAlgError:
        pass  # Fall through to Tier 3

    # Tier 3: Moore-Penrose pseudo-inverse (always succeeds)
    cov_pinv = np.linalg.pinv(cov_arr)
    w = cov_pinv @ ones
    denom = np.sum(w)
    if abs(denom) < 1e-15:
        w = np.ones(n) / n  # Ultimate fallback: equal weight
    else:
        w /= denom
    return pd.Series(w, index=cov.columns)


def nco_allocation(cov: pd.DataFrame, link: np.ndarray,
                   num_clusters: int, ordered_assets: list) -> pd.Series:
    """
    Two-level nested clustered optimization.

    Level 1 — Intra-Cluster Optimization:
    ──────────────────────────────────────
    For each cluster C_k (k = 1..K):
        w_k = min_var_portfolio(Σ_{C_k})

    This gives the optimal weights WITHIN each cluster, treating
    each cluster as an independent portfolio. The (N/K × N/K) matrices
    are well-conditioned, making inversion safe.

    Level 2 — Inter-Cluster Aggregation:
    ─────────────────────────────────────
    Construct a reduced K×K covariance matrix:
        Σ_reduced[i,j] = w_i^T · Σ_{C_i, C_j} · w_j

    Then solve the min-var problem on Σ_reduced to get cluster weights ω_k.

    Final weights: w_i = ω_k · w_{i|k}  for each asset i in cluster k.

    Speed: O(N³/K² + K³) — for K=5, N=500: ~25× faster than full Markowitz.

    Args:
        cov: Full N×N covariance matrix.
        link: Linkage matrix Z.
        num_clusters: Number of clusters K.
        ordered_assets: Seriated asset ordering.

    Returns:
        pd.Series of weights (may include negative values, always sums to 1.0).
    """
    asset_names = cov.columns.tolist()
    labels = fcluster(link, num_clusters, criterion='maxclust')
    cluster_labels = pd.Series(labels, index=asset_names)

    # ── Level 1: Intra-cluster min-var portfolios ──
    intra_weights = {}
    cluster_ids = sorted(cluster_labels.unique())

    for k in cluster_ids:
        members = cluster_labels[cluster_labels == k].index.tolist()
        if len(members) == 1:
            intra_weights[k] = pd.Series(1.0, index=members)
        else:
            cluster_cov = cov.loc[members, members]
            intra_weights[k] = get_min_var_portfolio(cluster_cov)

    # ── Level 2: Inter-cluster aggregation ──
    n_clusters = len(cluster_ids)
    reduced_cov = np.zeros((n_clusters, n_clusters))

    for i_idx, k_i in enumerate(cluster_ids):
        members_i = cluster_labels[cluster_labels == k_i].index.tolist()
        w_i = intra_weights[k_i].values

        for j_idx, k_j in enumerate(cluster_ids):
            members_j = cluster_labels[cluster_labels == k_j].index.tolist()
            w_j = intra_weights[k_j].values

            # Cross-covariance block
            cross_cov = cov.loc[members_i, members_j].values
            reduced_cov[i_idx, j_idx] = w_i @ cross_cov @ w_j

    # Solve for cluster-level weights
    reduced_cov_df = pd.DataFrame(reduced_cov, index=cluster_ids, columns=cluster_ids)
    cluster_w = get_min_var_portfolio(reduced_cov_df)

    # ── Combine: Final weights ──
    final_weights = pd.Series(0.0, index=asset_names)
    for k_idx, k in enumerate(cluster_ids):
        members = cluster_labels[cluster_labels == k].index.tolist()
        omega_k = cluster_w.iloc[k_idx]
        for m in members:
            final_weights[m] = omega_k * intra_weights[k][m]

    # Normalize to sum to 1.0
    final_weights /= final_weights.sum()

    return final_weights
