//! Hierarchical Risk Parity (HRP) and Marcenko-Pastur covariance denoising.
//!
//! Implements the full HRP pipeline from Lopez de Prado's "Machine Learning
//! for Asset Managers":
//! - Correlation-distance matrix
//! - Single-linkage hierarchical clustering (from scratch)
//! - Quasi-diagonalization (seriation)
//! - Recursive bisection for inverse-variance weighting
//!
//! Plus Marcenko-Pastur eigenvalue denoising and detoning (market-mode removal).
//!
//! All matrix operations are implemented from scratch -- no external crate
//! dependencies beyond PyO3.

use pyo3::prelude::*;

// ===========================================================================
// PyO3 result types
// ===========================================================================

/// Result of Hierarchical Risk Parity allocation.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct HRPResult {
    /// Portfolio weights (sum to 1.0), one per asset.
    #[pyo3(get)]
    pub weights: Vec<f64>,
    /// Asset indices in quasi-diagonalized (seriation) order.
    #[pyo3(get)]
    pub sorted_indices: Vec<usize>,
    /// Linkage matrix: Vec of (cluster_a, cluster_b, distance).
    #[pyo3(get)]
    pub linkage: Vec<(usize, usize, f64)>,
}

#[pymethods]
impl HRPResult {
    fn __repr__(&self) -> String {
        let n = self.weights.len();
        let top3: Vec<String> = self
            .weights
            .iter()
            .take(3)
            .map(|w| format!("{:.4}", w))
            .collect();
        let preview = top3.join(", ");
        format!(
            "HRPResult(n_assets={}, weights=[{}{}])",
            n,
            preview,
            if n > 3 { ", ..." } else { "" }
        )
    }
}

/// Result of Marcenko-Pastur covariance denoising.
#[pyclass(frozen)]
#[derive(Debug, Clone)]
pub struct DenoisedCov {
    /// The denoised covariance matrix (N x N).
    #[pyo3(get)]
    pub covariance: Vec<Vec<f64>>,
    /// Eigenvalues (sorted descending).
    #[pyo3(get)]
    pub eigenvalues: Vec<f64>,
    /// Number of signal eigenvalues (above Marcenko-Pastur threshold).
    #[pyo3(get)]
    pub n_signals: usize,
    /// Number of noise eigenvalues (below Marcenko-Pastur threshold).
    #[pyo3(get)]
    pub n_noise: usize,
}

#[pymethods]
impl DenoisedCov {
    fn __repr__(&self) -> String {
        let n = self.covariance.len();
        format!(
            "DenoisedCov(n={}, n_signals={}, n_noise={})",
            n, self.n_signals, self.n_noise
        )
    }
}

// ===========================================================================
// Private math helpers
// ===========================================================================

#[inline]
fn finite_or_zero(v: f64) -> f64 {
    if v.is_finite() { v } else { 0.0 }
}

/// Validate that a matrix is square and non-empty.
fn validate_square_matrix(mat: &[Vec<f64>]) -> PyResult<usize> {
    if mat.is_empty() {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "matrix must not be empty",
        ));
    }
    let n = mat.len();
    for (i, row) in mat.iter().enumerate() {
        if row.len() != n {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "matrix must be square: row {} has length {} but expected {}",
                i,
                row.len(),
                n
            )));
        }
    }
    Ok(n)
}

/// Validate that diagonal elements are positive (positive semi-definite proxy).
fn validate_positive_diagonal(mat: &[Vec<f64>], n: usize) -> PyResult<()> {
    for i in 0..n {
        if !mat[i][i].is_finite() {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "diagonal element [{}][{}] is not finite",
                i, i
            )));
        }
        if mat[i][i] < 0.0 {
            return Err(pyo3::exceptions::PyValueError::new_err(format!(
                "diagonal element [{}][{}] = {} is negative (not a valid covariance matrix)",
                i, i, mat[i][i]
            )));
        }
    }
    Ok(())
}

/// Convert covariance matrix to correlation matrix.
/// Returns (correlation_matrix, standard_deviations).
fn cov_to_corr(cov: &[Vec<f64>], n: usize) -> (Vec<Vec<f64>>, Vec<f64>) {
    let mut stds = vec![0.0_f64; n];
    for i in 0..n {
        stds[i] = cov[i][i].max(0.0).sqrt();
        if stds[i] < 1e-15 {
            stds[i] = 1e-15; // avoid division by zero
        }
    }

    let mut corr = vec![vec![0.0_f64; n]; n];
    for i in 0..n {
        for j in 0..n {
            if i == j {
                corr[i][j] = 1.0;
            } else {
                let c = cov[i][j] / (stds[i] * stds[j]);
                corr[i][j] = c.clamp(-1.0, 1.0);
            }
        }
    }
    (corr, stds)
}

/// Compute correlation distance matrix: d[i][j] = sqrt(0.5 * (1 - corr[i][j])).
fn corr_distance(corr: &[Vec<f64>], n: usize) -> Vec<Vec<f64>> {
    let mut dist = vec![vec![0.0_f64; n]; n];
    for i in 0..n {
        for j in (i + 1)..n {
            let d = (0.5 * (1.0 - corr[i][j])).max(0.0).sqrt();
            dist[i][j] = finite_or_zero(d);
            dist[j][i] = dist[i][j];
        }
    }
    dist
}

// ===========================================================================
// Single-linkage hierarchical clustering (from scratch)
// ===========================================================================

/// Single-linkage agglomerative clustering.
///
/// Returns a linkage list: Vec<(cluster_a, cluster_b, distance)>
/// where cluster indices < n are original assets, and indices >= n are
/// merged clusters formed in previous steps.
fn single_linkage_clustering(dist: &[Vec<f64>], n: usize) -> Vec<(usize, usize, f64)> {
    if n <= 1 {
        return vec![];
    }

    // Working copy of distances (condensed to active clusters only)
    let mut active_dist: Vec<Vec<f64>> = dist.to_vec();
    // Map from internal index to cluster label
    let mut cluster_label: Vec<usize> = (0..n).collect();
    // Track which indices are still active
    let mut active: Vec<bool> = vec![true; n];
    let mut linkage: Vec<(usize, usize, f64)> = Vec::with_capacity(n - 1);
    let mut next_cluster = n;

    for _ in 0..(n - 1) {
        // Find the two closest active clusters
        let mut min_dist = f64::INFINITY;
        let mut min_i = 0;
        let mut min_j = 0;

        for i in 0..active_dist.len() {
            if !active[i] {
                continue;
            }
            for j in (i + 1)..active_dist.len() {
                if !active[j] {
                    continue;
                }
                if active_dist[i][j] < min_dist {
                    min_dist = active_dist[i][j];
                    min_i = i;
                    min_j = j;
                }
            }
        }

        // Record the merge
        let label_a = cluster_label[min_i];
        let label_b = cluster_label[min_j];
        linkage.push((label_a, label_b, finite_or_zero(min_dist)));

        // Update distances using single-linkage rule: d(new, k) = min(d(i,k), d(j,k))
        for k in 0..active_dist.len() {
            if !active[k] || k == min_i || k == min_j {
                continue;
            }
            let d = active_dist[min_i][k].min(active_dist[min_j][k]);
            active_dist[min_i][k] = d;
            active_dist[k][min_i] = d;
        }

        // Deactivate min_j, keep min_i as the merged cluster
        active[min_j] = false;
        cluster_label[min_i] = next_cluster;
        next_cluster += 1;
    }

    linkage
}

// ===========================================================================
// Quasi-diagonalization (seriation)
// ===========================================================================

/// Recursively extract the leaf ordering from a linkage tree.
///
/// This is the quasi-diagonalization step: it produces an ordering of the
/// original assets such that correlated assets are placed adjacent in the
/// resulting sequence.
fn quasi_diag(linkage: &[(usize, usize, f64)], n: usize) -> Vec<usize> {
    if n == 0 {
        return vec![];
    }
    if n == 1 {
        return vec![0];
    }

    // Build a tree: for each non-leaf node, store its two children.
    // Node indices: 0..n are leaves, n..2n-1 are internal.
    let mut children: Vec<(usize, usize)> = Vec::with_capacity(linkage.len());
    for &(a, b, _) in linkage {
        children.push((a, b));
    }

    // Iterative DFS to extract leaf order (avoids stack overflow on large n).
    let root = n + linkage.len() - 1;
    let mut stack: Vec<usize> = vec![root];
    let mut order: Vec<usize> = Vec::with_capacity(n);

    while let Some(node) = stack.pop() {
        if node < n {
            // Leaf node
            order.push(node);
        } else {
            // Internal node: index into children array is (node - n)
            let idx = node - n;
            if idx < children.len() {
                let (left, right) = children[idx];
                // Push right first so left is processed first (DFS order)
                stack.push(right);
                stack.push(left);
            }
        }
    }

    order
}

// ===========================================================================
// Recursive bisection
// ===========================================================================

/// Compute inverse-variance portfolio weights on a subset of assets.
#[inline]
fn inv_var_weights(cov: &[Vec<f64>], indices: &[usize]) -> Vec<f64> {
    let n = indices.len();
    if n == 0 {
        return vec![];
    }
    if n == 1 {
        return vec![1.0];
    }

    let mut inv_vars = Vec::with_capacity(n);
    let mut total = 0.0_f64;
    for &i in indices {
        let var = cov[i][i].max(1e-15);
        let iv = 1.0 / var;
        inv_vars.push(finite_or_zero(iv));
        total += finite_or_zero(iv);
    }

    if total < 1e-30 {
        // Degenerate: equal weights
        return vec![1.0 / n as f64; n];
    }

    for w in &mut inv_vars {
        *w /= total;
    }
    inv_vars
}

/// Compute the cluster variance for a set of asset indices.
///
/// Uses the inverse-variance weights within the cluster to compute
/// the portfolio variance: w' * Sigma * w.
fn cluster_variance(cov: &[Vec<f64>], indices: &[usize]) -> f64 {
    let weights = inv_var_weights(cov, indices);
    let n = indices.len();
    let mut var = 0.0_f64;
    for a in 0..n {
        for b in 0..n {
            var += weights[a] * weights[b] * cov[indices[a]][indices[b]];
        }
    }
    finite_or_zero(var.max(1e-30))
}

/// Recursive bisection to allocate weights across the quasi-diagonalized order.
///
/// Splits the sorted asset list in half, allocates between halves inversely
/// proportional to their cluster variances, and recurses.
fn recursive_bisection(cov: &[Vec<f64>], sorted_indices: &[usize], n_assets: usize) -> Vec<f64> {
    let mut weights = vec![1.0_f64; n_assets];

    // Use an explicit stack to avoid recursion depth issues
    let mut stack: Vec<(Vec<usize>, f64)> = vec![(sorted_indices.to_vec(), 1.0)];

    while let Some((indices, alloc)) = stack.pop() {
        if indices.len() <= 1 {
            if let Some(&idx) = indices.first() {
                weights[idx] = alloc;
            }
            continue;
        }

        let mid = indices.len() / 2;
        let left = &indices[..mid];
        let right = &indices[mid..];

        let var_left = cluster_variance(cov, left);
        let var_right = cluster_variance(cov, right);

        // Allocate inversely proportional to cluster variance
        let total_inv = 1.0 / var_left + 1.0 / var_right;
        let alpha = if total_inv.is_finite() && total_inv > 1e-30 {
            (1.0 / var_left) / total_inv
        } else {
            0.5 // equal split if degenerate
        };

        stack.push((left.to_vec(), alloc * alpha));
        stack.push((right.to_vec(), alloc * (1.0 - alpha)));
    }

    // Normalize to sum to 1.0 (should already be close)
    let total: f64 = weights.iter().sum();
    if total > 1e-15 && total.is_finite() {
        for w in &mut weights {
            *w /= total;
        }
    }

    weights
}

// ===========================================================================
// Eigendecomposition (Jacobi method, from scratch)
// ===========================================================================

/// Jacobi eigenvalue algorithm for real symmetric matrices.
///
/// Returns (eigenvalues, eigenvectors) where eigenvectors[i][j] stores the
/// i-th component of the j-th eigenvector.  Eigenvalues are NOT sorted.
///
/// Uses the classical cyclic Jacobi method: for each off-diagonal element
/// above the tolerance, apply a Givens rotation to zero it out.  For a
/// symmetric matrix, convergence is guaranteed.
fn jacobi_eigen(mat: &[Vec<f64>], n: usize) -> (Vec<f64>, Vec<Vec<f64>>) {
    // Working copy of the matrix
    let mut a: Vec<Vec<f64>> = mat.to_vec();
    // Eigenvector matrix (starts as identity)
    let mut v = vec![vec![0.0_f64; n]; n];
    for i in 0..n {
        v[i][i] = 1.0;
    }

    let max_iters = 100 * n * n;
    let tol = 1e-12;

    for _ in 0..max_iters {
        // Find the largest off-diagonal element
        let mut max_off = 0.0_f64;
        let mut p = 0;
        let mut q = 1;
        for i in 0..n {
            for j in (i + 1)..n {
                let aij = a[i][j].abs();
                if aij > max_off {
                    max_off = aij;
                    p = i;
                    q = j;
                }
            }
        }

        if max_off < tol {
            break;
        }

        // Compute Jacobi rotation angle
        let app = a[p][p];
        let aqq = a[q][q];
        let apq = a[p][q];

        let theta = if (app - aqq).abs() < 1e-30 {
            std::f64::consts::FRAC_PI_4
        } else {
            0.5 * (2.0 * apq / (app - aqq)).atan()
        };

        let c = theta.cos();
        let s = theta.sin();

        // Apply Givens rotation: A' = G^T A G
        // Only rows/columns p and q are affected.

        // Step 1: Save the values we need before overwriting
        let mut a_ip = vec![0.0_f64; n]; // old a[i][p] for all i
        let mut a_iq = vec![0.0_f64; n]; // old a[i][q] for all i
        for i in 0..n {
            a_ip[i] = a[i][p];
            a_iq[i] = a[i][q];
        }

        // Step 2: Update off-diagonal elements in rows/cols p and q
        for i in 0..n {
            if i == p || i == q {
                continue;
            }
            let new_ip = c * a_ip[i] + s * a_iq[i];
            let new_iq = -s * a_ip[i] + c * a_iq[i];
            a[i][p] = new_ip;
            a[p][i] = new_ip;
            a[i][q] = new_iq;
            a[q][i] = new_iq;
        }

        // Step 3: Update the 2x2 diagonal block
        a[p][p] = c * c * app + 2.0 * c * s * apq + s * s * aqq;
        a[q][q] = s * s * app - 2.0 * c * s * apq + c * c * aqq;
        a[p][q] = 0.0;
        a[q][p] = 0.0;

        // Step 4: Accumulate eigenvectors
        for i in 0..n {
            let vip = v[i][p];
            let viq = v[i][q];
            v[i][p] = c * vip + s * viq;
            v[i][q] = -s * vip + c * viq;
        }
    }

    // Extract eigenvalues from diagonal
    let eigenvalues: Vec<f64> = (0..n).map(|i| finite_or_zero(a[i][i])).collect();

    (eigenvalues, v)
}

/// Sort eigenvalues descending and reorder eigenvectors accordingly.
fn sort_eigen_desc(
    eigenvalues: &mut Vec<f64>,
    eigenvectors: &mut Vec<Vec<f64>>,
    n: usize,
) {
    let mut indices: Vec<usize> = (0..n).collect();
    indices.sort_by(|&a, &b| {
        eigenvalues[b]
            .partial_cmp(&eigenvalues[a])
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    let old_vals = eigenvalues.clone();
    let old_vecs = eigenvectors.clone();

    for (new_idx, &old_idx) in indices.iter().enumerate() {
        eigenvalues[new_idx] = old_vals[old_idx];
        for i in 0..n {
            eigenvectors[i][new_idx] = old_vecs[i][old_idx];
        }
    }
}

/// Reconstruct a matrix from eigenvalues and eigenvectors: V * diag(lambda) * V^T.
fn reconstruct_from_eigen(
    eigenvalues: &[f64],
    eigenvectors: &[Vec<f64>],
    n: usize,
) -> Vec<Vec<f64>> {
    let mut result = vec![vec![0.0_f64; n]; n];
    for i in 0..n {
        for j in i..n {
            let mut val = 0.0_f64;
            for k in 0..n {
                val += eigenvectors[i][k] * eigenvalues[k] * eigenvectors[j][k];
            }
            let v = finite_or_zero(val);
            result[i][j] = v;
            result[j][i] = v;
        }
    }
    result
}

// ===========================================================================
// Public API: HRP
// ===========================================================================

/// Compute Hierarchical Risk Parity portfolio weights.
///
/// Full pipeline:
/// 1. Convert covariance to correlation matrix
/// 2. Compute correlation distance matrix
/// 3. Single-linkage hierarchical clustering
/// 4. Quasi-diagonalization (seriation)
/// 5. Recursive bisection (inverse-variance weighting)
///
/// Returns an `HRPResult` with weights summing to 1.0.
#[pyfunction]
pub fn hrp_weights(covariance: Vec<Vec<f64>>) -> PyResult<HRPResult> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    let n = validate_square_matrix(&covariance)?;
    validate_positive_diagonal(&covariance, n)?;

    if n == 1 {
        return Ok(HRPResult {
            weights: vec![1.0],
            sorted_indices: vec![0],
            linkage: vec![],
        });
    }

    // 1. Covariance → correlation
    let (corr, _stds) = cov_to_corr(&covariance, n);

    // 2. Correlation distance matrix
    let dist = corr_distance(&corr, n);

    // 3. Single-linkage clustering
    let linkage = single_linkage_clustering(&dist, n);

    // 4. Quasi-diagonalization
    let sorted_indices = quasi_diag(&linkage, n);

    // 5. Recursive bisection
    let weights = recursive_bisection(&covariance, &sorted_indices, n);

    Ok(HRPResult {
        weights,
        sorted_indices,
        linkage,
    })
}

// ===========================================================================
// Public API: Marcenko-Pastur denoising
// ===========================================================================

/// Denoise a covariance matrix using Marcenko-Pastur eigenvalue clipping.
///
/// 1. Compute eigendecomposition of the covariance matrix.
/// 2. Estimate the Marcenko-Pastur upper bound: lambda_+ = sigma^2 * (1 + sqrt(N/T))^2
/// 3. Replace noise eigenvalues (below lambda_+) with their average.
/// 4. Reconstruct the denoised covariance matrix.
///
/// # Arguments
/// * `covariance` - N x N covariance matrix.
/// * `n_observations` - Number of observations T used to estimate the covariance.
#[pyfunction]
pub fn denoise_covariance(
    covariance: Vec<Vec<f64>>,
    n_observations: usize,
) -> PyResult<DenoisedCov> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    let n = validate_square_matrix(&covariance)?;
    validate_positive_diagonal(&covariance, n)?;

    if n_observations < 2 {
        return Err(pyo3::exceptions::PyValueError::new_err(
            "n_observations must be at least 2",
        ));
    }

    if n == 1 {
        return Ok(DenoisedCov {
            covariance: covariance.clone(),
            eigenvalues: vec![covariance[0][0]],
            n_signals: 1,
            n_noise: 0,
        });
    }

    // Eigendecomposition
    let (mut eigenvalues, mut eigenvectors) = jacobi_eigen(&covariance, n);
    sort_eigen_desc(&mut eigenvalues, &mut eigenvectors, n);

    // Marcenko-Pastur bounds
    // q = N / T (ratio of variables to observations)
    let q = n as f64 / n_observations as f64;

    // Estimate sigma^2 as the average of eigenvalues (trace / n)
    let trace: f64 = eigenvalues.iter().sum();
    let sigma_sq = (trace / n as f64).max(1e-15);

    // MP upper bound: lambda_+ = sigma^2 * (1 + sqrt(q))^2
    let sqrt_q = q.sqrt();
    let lambda_plus = sigma_sq * (1.0 + sqrt_q).powi(2);

    // Classify eigenvalues as signal or noise
    let mut n_signals = 0usize;
    let mut noise_sum = 0.0_f64;
    let mut noise_count = 0usize;

    for &ev in &eigenvalues {
        if ev > lambda_plus {
            n_signals += 1;
        } else {
            noise_sum += ev;
            noise_count += 1;
        }
    }

    // Replace noise eigenvalues with their average
    let noise_avg = if noise_count > 0 {
        noise_sum / noise_count as f64
    } else {
        0.0
    };

    let mut denoised_eigenvalues = eigenvalues.clone();
    for ev in &mut denoised_eigenvalues {
        if *ev <= lambda_plus {
            *ev = noise_avg.max(1e-15); // keep positive for PSD
        }
    }

    // Reconstruct denoised covariance
    let denoised_cov = reconstruct_from_eigen(&denoised_eigenvalues, &eigenvectors, n);

    // Sort eigenvalues descending for output
    let mut sorted_eigenvalues = eigenvalues.clone();
    sorted_eigenvalues.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));

    Ok(DenoisedCov {
        covariance: denoised_cov,
        eigenvalues: sorted_eigenvalues,
        n_signals,
        n_noise: noise_count,
    })
}

// ===========================================================================
// Public API: Detoning
// ===========================================================================

/// Remove the market factor (first principal component) from a covariance matrix.
///
/// 1. Compute eigendecomposition.
/// 2. Zero out the largest eigenvalue's contribution.
/// 3. Reconstruct the covariance without the market mode.
///
/// This is useful for removing the dominant market factor that drives all
/// assets together, revealing the idiosyncratic structure.
#[pyfunction]
pub fn detone_covariance(covariance: Vec<Vec<f64>>) -> PyResult<Vec<Vec<f64>>> {
    crate::auth::require_pro().map_err(|e| pyo3::exceptions::PyPermissionError::new_err(e.to_string()))?;
    let n = validate_square_matrix(&covariance)?;
    validate_positive_diagonal(&covariance, n)?;

    if n == 1 {
        // Single asset: removing the only component leaves zero
        return Ok(vec![vec![0.0]]);
    }

    // Eigendecomposition
    let (mut eigenvalues, mut eigenvectors) = jacobi_eigen(&covariance, n);
    sort_eigen_desc(&mut eigenvalues, &mut eigenvectors, n);

    // Zero out the largest eigenvalue (market mode)
    let mut detoned_eigenvalues = eigenvalues.clone();
    detoned_eigenvalues[0] = 0.0;

    // Reconstruct without market mode
    let detoned_cov = reconstruct_from_eigen(&detoned_eigenvalues, &eigenvectors, n);

    Ok(detoned_cov)
}

// ===========================================================================
// Registration
// ===========================================================================

/// Register all HRP and denoising functions/types in the Python module.
pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<HRPResult>()?;
    m.add_class::<DenoisedCov>()?;
    m.add_function(wrap_pyfunction!(hrp_weights, m)?)?;
    m.add_function(wrap_pyfunction!(denoise_covariance, m)?)?;
    m.add_function(wrap_pyfunction!(detone_covariance, m)?)?;
    Ok(())
}

// ===========================================================================
// Tests
// ===========================================================================

#[cfg(test)]
mod tests {
    use super::*;

    /// Helper: build a simple 2x2 covariance matrix.
    fn cov_2x2() -> Vec<Vec<f64>> {
        vec![vec![1.0, 0.5], vec![0.5, 1.0]]
    }

    /// Helper: build a 3x3 covariance matrix with varying correlations.
    fn cov_3x3() -> Vec<Vec<f64>> {
        vec![
            vec![0.04, 0.006, 0.002],
            vec![0.006, 0.09, 0.009],
            vec![0.002, 0.009, 0.01],
        ]
    }

    /// Helper: build a 4x4 covariance with block structure.
    fn cov_4x4_block() -> Vec<Vec<f64>> {
        // Two blocks: (0,1) correlated, (2,3) correlated, cross-block weak
        vec![
            vec![0.04, 0.03, 0.001, 0.001],
            vec![0.03, 0.04, 0.001, 0.001],
            vec![0.001, 0.001, 0.09, 0.06],
            vec![0.001, 0.001, 0.06, 0.09],
        ]
    }

    /// Helper: identity covariance (uncorrelated).
    fn cov_identity(n: usize) -> Vec<Vec<f64>> {
        let mut m = vec![vec![0.0; n]; n];
        for i in 0..n {
            m[i][i] = 1.0;
        }
        m
    }

    // --- Validation ---

    #[test]
    fn test_validate_empty() {
        assert!(validate_square_matrix(&vec![]).is_err());
    }

    #[test]
    fn test_validate_non_square() {
        let m = vec![vec![1.0, 2.0], vec![3.0, 4.0, 5.0]];
        assert!(validate_square_matrix(&m).is_err());
    }

    #[test]
    fn test_validate_negative_diagonal() {
        let m = vec![vec![-1.0, 0.0], vec![0.0, 1.0]];
        assert!(validate_positive_diagonal(&m, 2).is_err());
    }

    // --- cov_to_corr ---

    #[test]
    fn test_cov_to_corr_identity() {
        let c = cov_identity(3);
        let (corr, stds) = cov_to_corr(&c, 3);
        for i in 0..3 {
            assert!((corr[i][i] - 1.0).abs() < 1e-10);
            assert!((stds[i] - 1.0).abs() < 1e-10);
            for j in 0..3 {
                if i != j {
                    assert!(corr[i][j].abs() < 1e-10);
                }
            }
        }
    }

    #[test]
    fn test_cov_to_corr_2x2() {
        let (corr, stds) = cov_to_corr(&cov_2x2(), 2);
        assert!((corr[0][1] - 0.5).abs() < 1e-10);
        assert!((stds[0] - 1.0).abs() < 1e-10);
    }

    // --- corr_distance ---

    #[test]
    fn test_corr_distance_identity() {
        let corr = cov_identity(3); // identity = zero correlation off-diagonal
        let dist = corr_distance(&corr, 3);
        for i in 0..3 {
            assert!(dist[i][i].abs() < 1e-10);
            for j in 0..3 {
                if i != j {
                    // d = sqrt(0.5 * (1 - 0)) = sqrt(0.5) ≈ 0.7071
                    assert!((dist[i][j] - (0.5_f64).sqrt()).abs() < 1e-10);
                }
            }
        }
    }

    #[test]
    fn test_corr_distance_perfect_corr() {
        let corr = vec![vec![1.0, 1.0], vec![1.0, 1.0]];
        let dist = corr_distance(&corr, 2);
        assert!(dist[0][1].abs() < 1e-10); // perfect correlation → zero distance
    }

    // --- Single-linkage clustering ---

    #[test]
    fn test_clustering_2x2() {
        let dist = vec![vec![0.0, 0.5], vec![0.5, 0.0]];
        let linkage = single_linkage_clustering(&dist, 2);
        assert_eq!(linkage.len(), 1);
        assert!((linkage[0].2 - 0.5).abs() < 1e-10);
    }

    #[test]
    fn test_clustering_3x3() {
        // 0-1 close, 2 far
        let dist = vec![
            vec![0.0, 0.1, 0.9],
            vec![0.1, 0.0, 0.8],
            vec![0.9, 0.8, 0.0],
        ];
        let linkage = single_linkage_clustering(&dist, 3);
        assert_eq!(linkage.len(), 2);
        // First merge should be 0 and 1 (distance 0.1)
        assert!((linkage[0].2 - 0.1).abs() < 1e-10);
    }

    #[test]
    fn test_clustering_single_asset() {
        let linkage = single_linkage_clustering(&vec![vec![0.0]], 1);
        assert!(linkage.is_empty());
    }

    // --- Quasi-diagonalization ---

    #[test]
    fn test_quasi_diag_2() {
        let linkage = vec![(0, 1, 0.5)];
        let order = quasi_diag(&linkage, 2);
        assert_eq!(order.len(), 2);
        // Should contain both indices
        assert!(order.contains(&0));
        assert!(order.contains(&1));
    }

    #[test]
    fn test_quasi_diag_3() {
        // Linkage: first merge 0+1→3, then merge 3+2→4
        let linkage = vec![(0, 1, 0.1), (3, 2, 0.5)];
        let order = quasi_diag(&linkage, 3);
        assert_eq!(order.len(), 3);
        assert!(order.contains(&0));
        assert!(order.contains(&1));
        assert!(order.contains(&2));
    }

    #[test]
    fn test_quasi_diag_preserves_adjacency() {
        // 0 and 1 are merged first → should be adjacent in output
        let linkage = vec![(0, 1, 0.1), (3, 2, 0.5)];
        let order = quasi_diag(&linkage, 3);
        let pos0 = order.iter().position(|&x| x == 0).unwrap();
        let pos1 = order.iter().position(|&x| x == 1).unwrap();
        assert_eq!((pos0 as i64 - pos1 as i64).abs(), 1);
    }

    // --- HRP weights ---

    #[test]
    fn test_hrp_single_asset() {
        let cov = vec![vec![0.04]];
        let result = hrp_weights(cov).unwrap();
        assert_eq!(result.weights.len(), 1);
        assert!((result.weights[0] - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_hrp_2_assets() {
        let result = hrp_weights(cov_2x2()).unwrap();
        assert_eq!(result.weights.len(), 2);
        let total: f64 = result.weights.iter().sum();
        assert!((total - 1.0).abs() < 1e-10, "total={}", total);
        // Equal variance → should be roughly equal weights
        assert!((result.weights[0] - result.weights[1]).abs() < 0.1);
    }

    #[test]
    fn test_hrp_3_assets() {
        let result = hrp_weights(cov_3x3()).unwrap();
        assert_eq!(result.weights.len(), 3);
        let total: f64 = result.weights.iter().sum();
        assert!((total - 1.0).abs() < 1e-10);
        for &w in &result.weights {
            assert!(w > 0.0, "weight {} should be positive", w);
        }
    }

    #[test]
    fn test_hrp_4_assets_block() {
        let result = hrp_weights(cov_4x4_block()).unwrap();
        assert_eq!(result.weights.len(), 4);
        let total: f64 = result.weights.iter().sum();
        assert!((total - 1.0).abs() < 1e-10);
        // Lower variance assets (0,1 with var=0.04) should get more weight
        // than higher variance assets (2,3 with var=0.09)
        let low_var_weight = result.weights[0] + result.weights[1];
        let high_var_weight = result.weights[2] + result.weights[3];
        assert!(
            low_var_weight > high_var_weight,
            "low_var={} should > high_var={}",
            low_var_weight,
            high_var_weight
        );
    }

    #[test]
    fn test_hrp_identity_cov() {
        let result = hrp_weights(cov_identity(5)).unwrap();
        assert_eq!(result.weights.len(), 5);
        let total: f64 = result.weights.iter().sum();
        assert!((total - 1.0).abs() < 1e-10);
        // Identical variance → equal weights
        for &w in &result.weights {
            assert!((w - 0.2).abs() < 0.05, "w={} expected ~0.2", w);
        }
    }

    #[test]
    fn test_hrp_linkage_length() {
        let result = hrp_weights(cov_3x3()).unwrap();
        assert_eq!(result.linkage.len(), 2); // n-1 merges for n=3
    }

    #[test]
    fn test_hrp_sorted_indices() {
        let result = hrp_weights(cov_4x4_block()).unwrap();
        assert_eq!(result.sorted_indices.len(), 4);
        let mut sorted = result.sorted_indices.clone();
        sorted.sort();
        assert_eq!(sorted, vec![0, 1, 2, 3]);
    }

    #[test]
    fn test_hrp_empty() {
        assert!(hrp_weights(vec![]).is_err());
    }

    #[test]
    fn test_hrp_non_square() {
        let m = vec![vec![1.0, 0.5], vec![0.5, 1.0, 0.3]];
        assert!(hrp_weights(m).is_err());
    }

    #[test]
    fn test_hrp_negative_diagonal() {
        let m = vec![vec![-1.0, 0.0], vec![0.0, 1.0]];
        assert!(hrp_weights(m).is_err());
    }

    #[test]
    fn test_hrp_repr() {
        let result = hrp_weights(cov_2x2()).unwrap();
        let repr = result.__repr__();
        assert!(repr.contains("HRPResult"));
        assert!(repr.contains("n_assets=2"));
    }

    // --- Jacobi eigendecomposition ---

    #[test]
    fn test_jacobi_identity() {
        let m = cov_identity(3);
        let (eigenvalues, _eigenvectors) = jacobi_eigen(&m, 3);
        for &ev in &eigenvalues {
            assert!((ev - 1.0).abs() < 1e-8, "ev={}", ev);
        }
    }

    #[test]
    fn test_jacobi_diagonal() {
        let m = vec![
            vec![4.0, 0.0, 0.0],
            vec![0.0, 2.0, 0.0],
            vec![0.0, 0.0, 1.0],
        ];
        let (mut eigenvalues, _) = jacobi_eigen(&m, 3);
        eigenvalues.sort_by(|a, b| b.partial_cmp(a).unwrap());
        assert!((eigenvalues[0] - 4.0).abs() < 1e-8);
        assert!((eigenvalues[1] - 2.0).abs() < 1e-8);
        assert!((eigenvalues[2] - 1.0).abs() < 1e-8);
    }

    #[test]
    fn test_jacobi_2x2() {
        let m = vec![vec![2.0, 1.0], vec![1.0, 2.0]];
        let (mut eigenvalues, _) = jacobi_eigen(&m, 2);
        eigenvalues.sort_by(|a, b| b.partial_cmp(a).unwrap());
        // Eigenvalues of [[2,1],[1,2]] are 3 and 1
        assert!((eigenvalues[0] - 3.0).abs() < 1e-8, "ev0={}", eigenvalues[0]);
        assert!((eigenvalues[1] - 1.0).abs() < 1e-8, "ev1={}", eigenvalues[1]);
    }

    #[test]
    fn test_jacobi_reconstruction() {
        let m = cov_3x3();
        let (eigenvalues, eigenvectors) = jacobi_eigen(&m, 3);
        let reconstructed = reconstruct_from_eigen(&eigenvalues, &eigenvectors, 3);
        for i in 0..3 {
            for j in 0..3 {
                assert!(
                    (reconstructed[i][j] - m[i][j]).abs() < 1e-6,
                    "mismatch at [{},{}]: {} vs {}",
                    i,
                    j,
                    reconstructed[i][j],
                    m[i][j]
                );
            }
        }
    }

    // --- Marcenko-Pastur denoising ---

    #[test]
    fn test_denoise_basic() {
        let cov = cov_3x3();
        let result = denoise_covariance(cov, 100).unwrap();
        assert_eq!(result.covariance.len(), 3);
        assert_eq!(result.eigenvalues.len(), 3);
        assert!(result.n_signals + result.n_noise == 3);
        // Eigenvalues should be descending
        for i in 1..result.eigenvalues.len() {
            assert!(
                result.eigenvalues[i] <= result.eigenvalues[i - 1] + 1e-10,
                "eigenvalues not descending: {} > {}",
                result.eigenvalues[i],
                result.eigenvalues[i - 1]
            );
        }
    }

    #[test]
    fn test_denoise_preserves_symmetry() {
        let cov = cov_3x3();
        let result = denoise_covariance(cov, 100).unwrap();
        let n = result.covariance.len();
        for i in 0..n {
            for j in 0..n {
                assert!(
                    (result.covariance[i][j] - result.covariance[j][i]).abs() < 1e-10,
                    "not symmetric at [{},{}]",
                    i,
                    j
                );
            }
        }
    }

    #[test]
    fn test_denoise_positive_diagonal() {
        let cov = cov_4x4_block();
        let result = denoise_covariance(cov, 200).unwrap();
        for i in 0..4 {
            assert!(
                result.covariance[i][i] > 0.0,
                "diagonal [{}][{}] = {} not positive",
                i,
                i,
                result.covariance[i][i]
            );
        }
    }

    #[test]
    fn test_denoise_identity() {
        // Identity covariance: with many observations, all eigenvalues are signal
        let result = denoise_covariance(cov_identity(3), 1000).unwrap();
        assert_eq!(result.n_signals, 3);
        assert_eq!(result.n_noise, 0);
    }

    #[test]
    fn test_denoise_single_asset() {
        let result = denoise_covariance(vec![vec![0.04]], 100).unwrap();
        assert_eq!(result.covariance.len(), 1);
        assert_eq!(result.n_signals, 1);
    }

    #[test]
    fn test_denoise_too_few_obs() {
        assert!(denoise_covariance(cov_2x2(), 1).is_err());
    }

    #[test]
    fn test_denoise_empty() {
        assert!(denoise_covariance(vec![], 100).is_err());
    }

    #[test]
    fn test_denoise_repr() {
        let result = denoise_covariance(cov_3x3(), 100).unwrap();
        let repr = result.__repr__();
        assert!(repr.contains("DenoisedCov"));
        assert!(repr.contains("n=3"));
    }

    // --- Detoning ---

    #[test]
    fn test_detone_basic() {
        let cov = cov_3x3();
        let result = detone_covariance(cov.clone()).unwrap();
        assert_eq!(result.len(), 3);
        // Detoned covariance should have smaller magnitude than original
        // (we removed the dominant eigenvalue)
        let orig_trace: f64 = (0..3).map(|i| cov[i][i]).sum();
        let det_trace: f64 = (0..3).map(|i| result[i][i]).sum();
        assert!(
            det_trace < orig_trace + 1e-10,
            "detoned trace {} should be < original trace {}",
            det_trace,
            orig_trace
        );
    }

    #[test]
    fn test_detone_preserves_symmetry() {
        let result = detone_covariance(cov_3x3()).unwrap();
        for i in 0..3 {
            for j in 0..3 {
                assert!(
                    (result[i][j] - result[j][i]).abs() < 1e-10,
                    "not symmetric at [{},{}]",
                    i,
                    j
                );
            }
        }
    }

    #[test]
    fn test_detone_single_asset() {
        let result = detone_covariance(vec![vec![0.04]]).unwrap();
        assert!((result[0][0] - 0.0).abs() < 1e-10);
    }

    #[test]
    fn test_detone_identity() {
        // Identity: all eigenvalues = 1, removing one leaves n-1 eigenvalues at 1
        let result = detone_covariance(cov_identity(3)).unwrap();
        let trace: f64 = (0..3).map(|i| result[i][i]).sum();
        // Removed one eigenvalue of 1.0 → trace should be ~2.0
        assert!((trace - 2.0).abs() < 0.1, "trace={}", trace);
    }

    #[test]
    fn test_detone_empty() {
        assert!(detone_covariance(vec![]).is_err());
    }

    // --- Integration: HRP + denoised covariance ---

    #[test]
    fn test_hrp_on_denoised_cov() {
        let cov = cov_4x4_block();
        let denoised = denoise_covariance(cov, 200).unwrap();
        let result = hrp_weights(denoised.covariance).unwrap();
        assert_eq!(result.weights.len(), 4);
        let total: f64 = result.weights.iter().sum();
        assert!((total - 1.0).abs() < 1e-10);
    }
}
