foxhunt/services/trading_agent_service/src/allocation.rs

//! Portfolio Allocation Logic
//!
//! Determines position sizes and weights across selected assets.
//! Implements 5 allocation strategies:
//! 1. Equal Weight (Baseline)
//! 2. Risk Parity (Inverse volatility weighting)
//! 3. Mean-Variance Optimization (Markowitz)
//! 4. ML-Optimized (ML predictions as expected returns)
//! 5. Kelly Criterion (Position sizing by edge)

use anyhow::{Context, Result};
use nalgebra::{DMatrix, DVector};
use rust_decimal::Decimal;
use std::collections::HashMap;

/// Portfolio allocation engine
pub struct PortfolioAllocator {
    method: AllocationMethod,
}

/// Allocation strategy selection
#[derive(Debug, Clone)]
pub enum AllocationMethod {
    /// Equal weight allocation (1/N)
    EqualWeight,
    /// Risk parity (inverse volatility weighting)
    RiskParity,
    /// Mean-variance optimization (Markowitz)
    MeanVariance {
        /// Risk aversion parameter (higher = more conservative)
        lambda: f64,
    },
    /// ML-optimized allocation (use ML predictions as expected returns)
    MLOptimized,
    /// Kelly Criterion (fractional Kelly for risk management)
    KellyCriterion {
        /// Fraction of Kelly to use (0.25 = quarter Kelly)
        fraction: f64,
    },
}

impl PortfolioAllocator {
    /// Create new portfolio allocator with specified method
    pub fn new(method: AllocationMethod) -> Self {
        Self { method }
    }

    /// Allocate capital across assets
    ///
    /// # Arguments
    /// * `assets` - Asset information (returns, volatility, ML scores)
    /// * `total_capital` - Total capital to allocate
    ///
    /// # Returns
    /// HashMap of symbol -> allocated capital
    pub fn allocate(
        &self,
        assets: &[AssetInfo],
        total_capital: Decimal,
    ) -> Result<HashMap<String, Decimal>> {
        if assets.is_empty() {
            return Ok(HashMap::new());
        }

        match &self.method {
            AllocationMethod::EqualWeight => self.equal_weight(assets, total_capital),
            AllocationMethod::RiskParity => self.risk_parity(assets, total_capital),
            AllocationMethod::MeanVariance { lambda } => {
                self.mean_variance(assets, total_capital, *lambda)
            },
            AllocationMethod::MLOptimized => self.ml_optimized(assets, total_capital),
            AllocationMethod::KellyCriterion { fraction } => {
                self.kelly_criterion(assets, total_capital, *fraction)
            },
        }
    }

    /// Allocate capital across assets using a correlation matrix
    ///
    /// Like [`allocate`](Self::allocate), but accepts an N x N correlation matrix
    /// to build a full covariance matrix for mean-variance optimization.
    /// Only meaningful when the allocation method is `MeanVariance` or `MLOptimized`;
    /// other methods ignore the correlation matrix.
    ///
    /// # Arguments
    /// * `assets` - Asset information (returns, volatility, ML scores)
    /// * `total_capital` - Total capital to allocate
    /// * `correlations` - N x N correlation matrix (must be symmetric, 1.0 on diagonal)
    ///
    /// # Errors
    /// Returns an error if the correlation matrix dimensions do not match the asset count.
    pub fn allocate_with_correlations(
        &self,
        assets: &[AssetInfo],
        total_capital: Decimal,
        correlations: &DMatrix<f64>,
    ) -> Result<HashMap<String, Decimal>> {
        if assets.is_empty() {
            return Ok(HashMap::new());
        }

        match &self.method {
            AllocationMethod::MeanVariance { lambda } => {
                self.mean_variance_with_corr(assets, total_capital, *lambda, Some(correlations))
            }
            AllocationMethod::MLOptimized => {
                // Use ML scores as expected returns, then apply correlated mean-variance
                let ml_assets: Vec<AssetInfo> = assets
                    .iter()
                    .map(|a| {
                        let mut asset = a.clone();
                        asset.expected_return = a.ml_score;
                        asset
                    })
                    .collect();
                self.mean_variance_with_corr(&ml_assets, total_capital, 1.0, Some(correlations))
            }
            // Other methods don't use correlations — delegate to standard allocate
            _ => self.allocate(assets, total_capital),
        }
    }

    /// Strategy 1: Equal Weight (Baseline)
    ///
    /// Allocates capital equally across all assets (1/N portfolio).
    /// Simple but effective baseline strategy.
    fn equal_weight(
        &self,
        assets: &[AssetInfo],
        total_capital: Decimal,
    ) -> Result<HashMap<String, Decimal>> {
        let n = Decimal::from(assets.len());
        let weight_per_asset = Decimal::ONE / n;
        let capital_per_asset = total_capital * weight_per_asset;

        Ok(assets
            .iter()
            .map(|asset| (asset.symbol.clone(), capital_per_asset))
            .collect())
    }

    /// Strategy 2: Risk Parity (Allocate inversely to volatility)
    ///
    /// Assets with lower volatility receive higher allocation.
    /// Aims to equalize risk contribution across assets.
    fn risk_parity(
        &self,
        assets: &[AssetInfo],
        total_capital: Decimal,
    ) -> Result<HashMap<String, Decimal>> {
        // Calculate inverse volatility weights
        let inv_vols: Vec<f64> = assets.iter()
            .map(|a| 1.0 / a.volatility.max(0.001)) // Avoid division by zero
            .collect();

        let sum_inv_vols: f64 = inv_vols.iter().sum();

        let mut allocations = HashMap::new();
        for (asset, inv_vol) in assets.iter().zip(inv_vols.iter()) {
            let weight = Decimal::from_f64_retain(inv_vol / sum_inv_vols).unwrap_or(Decimal::ZERO);
            allocations.insert(asset.symbol.clone(), total_capital * weight);
        }

        Ok(allocations)
    }

    /// Strategy 3: Mean-Variance Optimization (Markowitz)
    ///
    /// Maximizes expected return for given level of risk.
    /// Solves: max (mu^T w - lambda * w^T Sigma w)
    ///
    /// # Arguments
    /// * `lambda` - Risk aversion parameter (higher = more conservative)
    /// * `correlations` - Optional N x N correlation matrix. When `None`, assumes
    ///   independent assets (diagonal covariance). When provided, builds full
    ///   covariance: `Sigma[i][j] = corr[i][j] * vol_i * vol_j`.
    fn mean_variance(
        &self,
        assets: &[AssetInfo],
        total_capital: Decimal,
        lambda: f64,
    ) -> Result<HashMap<String, Decimal>> {
        self.mean_variance_with_corr(assets, total_capital, lambda, None)
    }

    /// Mean-Variance optimization with optional correlation matrix.
    ///
    /// When `correlations` is `Some`, builds the full covariance matrix from the
    /// correlation matrix and per-asset volatilities. Falls back to diagonal
    /// covariance if the correlation matrix is ill-conditioned.
    fn mean_variance_with_corr(
        &self,
        assets: &[AssetInfo],
        total_capital: Decimal,
        lambda: f64,
        correlations: Option<&DMatrix<f64>>,
    ) -> Result<HashMap<String, Decimal>> {
        let n = assets.len();

        // Expected returns vector
        let mu = DVector::from_vec(assets.iter().map(|a| a.expected_return).collect());

        // Build covariance matrix
        let mut sigma = if let Some(corr) = correlations {
            // Validate dimensions
            if corr.nrows() != n || corr.ncols() != n {
                anyhow::bail!(
                    "Correlation matrix dimensions ({}, {}) do not match asset count {}",
                    corr.nrows(),
                    corr.ncols(),
                    n
                );
            }
            // Build full covariance: Sigma[i][j] = corr[i][j] * vol_i * vol_j
            let mut cov = DMatrix::zeros(n, n);
            for i in 0..n {
                let vol_i = assets.get(i).map(|a| a.volatility).unwrap_or(0.0);
                for j in 0..n {
                    let vol_j = assets.get(j).map(|a| a.volatility).unwrap_or(0.0);
                    let corr_ij = corr.get((i, j)).copied().unwrap_or(0.0);
                    if let Some(cell) = cov.get_mut((i, j)) {
                        *cell = corr_ij * vol_i * vol_j;
                    }
                }
            }
            cov
        } else {
            // Diagonal covariance (independent assets)
            let mut cov = DMatrix::zeros(n, n);
            for (i, asset) in assets.iter().enumerate() {
                if let Some(cell) = cov.get_mut((i, i)) {
                    *cell = asset.volatility.powi(2);
                }
            }
            cov
        };

        // Add small regularization to diagonal for numerical stability
        for i in 0..n {
            sigma[(i, i)] += 1e-6;
        }

        // Solve: maximize (mu^T w - lambda * w^T Sigma w)
        // Analytical solution: w = (1 / 2*lambda) * Sigma^-1 * mu
        let sigma_inv = sigma
            .try_inverse()
            .context("Failed to invert covariance matrix")?;

        let w_optimal = sigma_inv * mu * (1.0 / (2.0 * lambda));

        // Normalize weights to sum to 1
        let sum_weights: f64 = w_optimal.iter().map(|&x| x.abs()).sum();
        if sum_weights < 1e-10 {
            // Fallback to equal weight if optimization fails
            return self.equal_weight(assets, total_capital);
        }

        let w_normalized: Vec<f64> = w_optimal.iter().map(|&x| x / sum_weights).collect();

        // Clamp to [0, 0.20] (max 20% per asset for risk management)
        let mut allocations = HashMap::new();
        let mut total_weight = 0.0;

        for (i, asset) in assets.iter().enumerate() {
            let weight = w_normalized[i].clamp(0.0, 0.20);
            total_weight += weight;
            allocations.insert(
                asset.symbol.clone(),
                Decimal::ZERO, // Placeholder
            );
        }

        // Renormalize after clamping
        for (i, asset) in assets.iter().enumerate() {
            let weight = w_normalized[i].clamp(0.0, 0.20) / total_weight;
            let capital = total_capital * Decimal::from_f64_retain(weight).unwrap_or(Decimal::ZERO);
            allocations.insert(asset.symbol.clone(), capital);
        }

        Ok(allocations)
    }

    /// Strategy 4: ML-Optimized (Use ML predictions as expected returns)
    ///
    /// Replaces expected returns with ML model predictions.
    /// Then applies mean-variance optimization.
    fn ml_optimized(
        &self,
        assets: &[AssetInfo],
        total_capital: Decimal,
    ) -> Result<HashMap<String, Decimal>> {
        // Use ML scores as expected returns
        let ml_assets: Vec<AssetInfo> = assets
            .iter()
            .map(|a| {
                let mut asset = a.clone();
                asset.expected_return = a.ml_score; // ML prediction replaces expected return
                asset
            })
            .collect();

        // Apply mean-variance with ML predictions (moderate risk aversion)
        self.mean_variance(&ml_assets, total_capital, 1.0)
    }

    /// Strategy 5: Kelly Criterion (Size positions by edge)
    ///
    /// Positions sized according to perceived edge.
    /// Uses fractional Kelly for risk management.
    ///
    /// # Arguments
    /// * `fraction` - Fraction of Kelly to use (0.25 = quarter Kelly)
    fn kelly_criterion(
        &self,
        assets: &[AssetInfo],
        total_capital: Decimal,
        fraction: f64,
    ) -> Result<HashMap<String, Decimal>> {
        let mut allocations = HashMap::new();

        // First pass: calculate Kelly fractions
        let kelly_fractions: Vec<(String, f64)> = assets
            .iter()
            .map(|asset| {
                // Kelly formula: f = (p * b - q) / b
                // Where p = win rate, q = loss rate, b = win/loss ratio
                let win_rate = asset.win_rate.max(0.01);
                let loss_rate = 1.0 - win_rate;
                let win_loss_ratio = asset.avg_win / asset.avg_loss.max(0.01);

                let kelly_fraction = (win_rate * win_loss_ratio - loss_rate) / win_loss_ratio;
                let f = (kelly_fraction * fraction).clamp(0.0, 0.20); // Clamp to [0, 20%] for risk management

                (asset.symbol.clone(), f)
            })
            .collect();

        // Calculate total fraction
        let total_fraction: f64 = kelly_fractions.iter().map(|(_, f)| f).sum();

        // Normalize if total exceeds 100%
        let normalization_factor = if total_fraction > 1.0 {
            1.0 / total_fraction
        } else {
            1.0
        };

        // Second pass: allocate capital
        for (symbol, f) in kelly_fractions {
            let normalized_f = f * normalization_factor;
            let capital =
                total_capital * Decimal::from_f64_retain(normalized_f).unwrap_or(Decimal::ZERO);
            allocations.insert(symbol, capital);
        }

                Ok(allocations)
            }
        
            /// Strategy 5b: Kelly Criterion with Regime Adaptation
            ///
            /// Extends Kelly Criterion with regime-aware position sizing.
            /// Applies regime-specific multipliers to base Kelly allocations:
            /// - Crisis/Volatile: 0.2x-0.5x (reduce position size)
            /// - Ranging: 0.8x (reduce position size in choppy markets)
            /// - Normal: 1.0x (full Kelly)
            /// - Trending: 1.5x (increase size in trends)
            ///
            /// # Arguments
            /// * `assets` - Asset information (returns, volatility, ML scores)
            /// * `total_capital` - Total capital to allocate
            /// * `fraction` - Fraction of Kelly to use (0.25 = quarter Kelly)
            /// * `pool` - Database connection pool for regime queries
            ///
            /// # Returns
            /// HashMap of symbol -> regime-adjusted allocated capital
            ///
            /// # Algorithm
            /// 1. Calculate base Kelly allocations
            /// 2. Query regime state for each symbol
            /// 3. Apply regime multiplier (0.2x-1.5x)
            /// 4. Normalize if total exceeds 100%
            /// 5. Cap individual positions at 20%
            pub async fn kelly_criterion_regime_adaptive(
                &self,
                assets: &[AssetInfo],
                total_capital: Decimal,
                fraction: f64,
                pool: &sqlx::PgPool,
            ) -> Result<HashMap<String, Decimal>> {
                // Step 1: Calculate base Kelly allocations
                let base_allocations = self.kelly_criterion(assets, total_capital, fraction)?;
        
                // Step 2 & 3: Query regime states and apply multipliers
                let mut regime_adjusted = HashMap::new();
        
                for (symbol, base_capital) in &base_allocations {
                    // Query regime state (fallback to Normal if unavailable)
                    let regime = match crate::regime::get_regime_for_symbol(pool, symbol).await {
                        Ok(r) => r.regime,
                        Err(_) => {
                            // Regime data unavailable - use Normal (1.0x multiplier)
                            "Normal".to_string()
                        }
                    };
        
                    // Get regime-specific position multiplier
                    let multiplier = crate::regime::regime_to_position_multiplier(&regime);
        
                    // Apply multiplier to base allocation
                    let adjusted_capital = *base_capital * Decimal::from_f64_retain(multiplier)
                        .unwrap_or(Decimal::ONE);
        
                    regime_adjusted.insert(symbol.clone(), adjusted_capital);
                }
        
                // Step 4: Normalize if total exceeds capital
                let total_adjusted: Decimal = regime_adjusted.values().sum();
                if total_adjusted > total_capital {
                    let normalization_factor = total_capital / total_adjusted;
                    for capital in regime_adjusted.values_mut() {
                        *capital *= normalization_factor;
                    }
                }
        
                // Step 5: Cap individual positions at 20%
                let max_per_asset = total_capital * Decimal::from_f64_retain(0.20).unwrap_or(Decimal::ZERO);
                for capital in regime_adjusted.values_mut() {
                    *capital = (*capital).min(max_per_asset);
                }
        
                Ok(regime_adjusted)
            }
        }
        
        /// Asset information for allocation
#[derive(Debug, Clone)]
pub struct AssetInfo {
    /// Symbol identifier
    pub symbol: String,
    /// Expected return (annualized)
    pub expected_return: f64,
    /// Volatility (annualized standard deviation)
    pub volatility: f64,
    /// ML model prediction score (0-1)
    pub ml_score: f64,
    /// Historical win rate (0-1)
    pub win_rate: f64,
    /// Average winning trade size
    pub avg_win: f64,
    /// Average losing trade size
    pub avg_loss: f64,
}

impl Default for AssetInfo {
    fn default() -> Self {
        Self {
            symbol: String::new(),
            expected_return: 0.0,
            volatility: 0.15, // 15% default volatility
            ml_score: 0.5,
            win_rate: 0.5,
            avg_win: 100.0,
            avg_loss: 100.0,
        }
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;

    fn create_test_assets() -> Vec<AssetInfo> {
        vec![
            AssetInfo {
                symbol: "ES.FUT".to_string(),
                expected_return: 0.08,
                volatility: 0.15,
                ml_score: 0.65,
                win_rate: 0.55,
                avg_win: 100.0,
                avg_loss: 80.0,
            },
            AssetInfo {
                symbol: "NQ.FUT".to_string(),
                expected_return: 0.10,
                volatility: 0.20,
                ml_score: 0.70,
                win_rate: 0.52,
                avg_win: 150.0,
                avg_loss: 100.0,
            },
            AssetInfo {
                symbol: "ZN.FUT".to_string(),
                expected_return: 0.04,
                volatility: 0.10,
                ml_score: 0.55,
                win_rate: 0.53,
                avg_win: 50.0,
                avg_loss: 45.0,
            },
        ]
    }

    #[test]
    fn test_equal_weight() {
        let allocator = PortfolioAllocator::new(AllocationMethod::EqualWeight);
        let assets = create_test_assets();
        let total_capital = Decimal::from(100_000);

        let alloc = allocator.allocate(&assets, total_capital).unwrap();

        assert_eq!(alloc.len(), 3);

        // Calculate expected allocation per asset
        let expected_per_asset = Decimal::from(100_000) / Decimal::from(3);

        // Check each allocation (with small tolerance for rounding)
        for (symbol, capital) in &alloc {
            let diff = (*capital - expected_per_asset).abs();
            assert!(
                diff < Decimal::from_f64_retain(0.01).unwrap(),
                "{} allocation {} differs from expected {} by {}",
                symbol,
                capital,
                expected_per_asset,
                diff
            );
        }

        // Verify sum equals total capital (within rounding)
        let sum: Decimal = alloc.values().sum();
        assert!((sum - total_capital).abs() < Decimal::from(1));
    }

    #[test]
    fn test_risk_parity() {
        let allocator = PortfolioAllocator::new(AllocationMethod::RiskParity);
        let assets = create_test_assets();
        let total_capital = Decimal::from(100_000);

        let alloc = allocator.allocate(&assets, total_capital).unwrap();

        assert_eq!(alloc.len(), 3);

        // Lower volatility assets should get higher allocation
        // ZN.FUT (10% vol) > ES.FUT (15% vol) > NQ.FUT (20% vol)
        assert!(alloc["ZN.FUT"] > alloc["ES.FUT"]);
        assert!(alloc["ES.FUT"] > alloc["NQ.FUT"]);

        // Verify sum equals total capital (within rounding)
        let sum: Decimal = alloc.values().sum();
        assert!((sum - total_capital).abs() < Decimal::from(1));
    }

    #[test]
    fn test_mean_variance() {
        let allocator = PortfolioAllocator::new(AllocationMethod::MeanVariance { lambda: 2.0 });
        let assets = create_test_assets();
        let total_capital = Decimal::from(100_000);

        let alloc = allocator.allocate(&assets, total_capital).unwrap();

        assert_eq!(alloc.len(), 3);

        // Should allocate based on return/risk tradeoff
        // All allocations should be non-negative
        for (symbol, capital) in &alloc {
            assert!(
                *capital >= Decimal::ZERO,
                "{} has negative allocation: {}",
                symbol,
                capital
            );
        }

        // Verify sum equals total capital (within rounding)
        let sum: Decimal = alloc.values().sum();
        assert!(
            (sum - total_capital).abs() < Decimal::from(10),
            "Sum {} differs from total {} by more than 10",
            sum,
            total_capital
        );
    }

    #[test]
    fn test_ml_optimized() {
        let allocator = PortfolioAllocator::new(AllocationMethod::MLOptimized);
        let assets = create_test_assets();
        let total_capital = Decimal::from(100_000);

        let alloc = allocator.allocate(&assets, total_capital).unwrap();

        assert_eq!(alloc.len(), 3);

        // Should favor higher ML scores
        // NQ.FUT (0.70) should get more than ES.FUT (0.65) > ZN.FUT (0.55)
        // (accounting for volatility adjustments)

        // All allocations should be non-negative
        for (symbol, capital) in &alloc {
            assert!(
                *capital >= Decimal::ZERO,
                "{} has negative allocation: {}",
                symbol,
                capital
            );
        }

        // Verify sum equals total capital (within rounding)
        let sum: Decimal = alloc.values().sum();
        assert!(
            (sum - total_capital).abs() < Decimal::from(10),
            "Sum {} differs from total {} by more than 10",
            sum,
            total_capital
        );
    }

    #[test]
    fn test_kelly_criterion() {
        let allocator =
            PortfolioAllocator::new(AllocationMethod::KellyCriterion { fraction: 0.25 });
        let assets = create_test_assets();
        let total_capital = Decimal::from(100_000);

        let alloc = allocator.allocate(&assets, total_capital).unwrap();

        assert_eq!(alloc.len(), 3);

        // All allocations should be non-negative
        for (symbol, capital) in &alloc {
            assert!(
                *capital >= Decimal::ZERO,
                "{} has negative allocation: {}",
                symbol,
                capital
            );
        }

        // No single position should exceed 20% (max clamp)
        for (symbol, capital) in &alloc {
            let weight = *capital / total_capital;
            assert!(
                weight <= Decimal::from_f64_retain(0.20).unwrap(),
                "{} exceeds 20% allocation: {}",
                symbol,
                weight
            );
        }

        // Verify sum doesn't exceed total capital
        let sum: Decimal = alloc.values().sum();
        assert!(
            sum <= total_capital,
            "Sum {} exceeds total {}",
            sum,
            total_capital
        );
    }

    #[test]
    fn test_empty_assets() {
        let allocator = PortfolioAllocator::new(AllocationMethod::EqualWeight);
        let assets = vec![];
        let total_capital = Decimal::from(100_000);

        let alloc = allocator.allocate(&assets, total_capital).unwrap();

        assert_eq!(alloc.len(), 0);
    }

    #[test]
    fn test_single_asset() {
        let allocator = PortfolioAllocator::new(AllocationMethod::EqualWeight);
        let assets = vec![AssetInfo {
            symbol: "ES.FUT".to_string(),
            expected_return: 0.08,
            volatility: 0.15,
            ml_score: 0.65,
            win_rate: 0.55,
            avg_win: 100.0,
            avg_loss: 80.0,
        }];
        let total_capital = Decimal::from(100_000);

        let alloc = allocator.allocate(&assets, total_capital).unwrap();

        assert_eq!(alloc.len(), 1);
        assert_eq!(alloc["ES.FUT"], total_capital);
    }

    #[test]
    fn test_allocation_methods_consistency() {
        let assets = create_test_assets();
        let total_capital = Decimal::from(100_000);

        let methods = vec![
            AllocationMethod::EqualWeight,
            AllocationMethod::RiskParity,
            AllocationMethod::MeanVariance { lambda: 1.0 },
            AllocationMethod::MLOptimized,
            AllocationMethod::KellyCriterion { fraction: 0.25 },
        ];

        for method in methods {
            let allocator = PortfolioAllocator::new(method);
            let alloc = allocator.allocate(&assets, total_capital).unwrap();

            // All methods should allocate to all assets
            assert_eq!(alloc.len(), 3, "Method allocates to all assets");

            // All allocations should be non-negative
            for (symbol, capital) in &alloc {
                assert!(
                    *capital >= Decimal::ZERO,
                    "{} has negative allocation",
                    symbol
                );
            }
        }
    }

    /// Identity correlation matrix (diagonal = 1.0) should produce the same result
    /// as the default diagonal covariance path (no correlations).
    #[test]
    fn test_mean_variance_identity_correlation_matches_diagonal() {
        let allocator = PortfolioAllocator::new(AllocationMethod::MeanVariance { lambda: 2.0 });
        let assets = create_test_assets();
        let total_capital = Decimal::from(100_000);
        let n = assets.len();

        // Identity correlation matrix
        let identity = DMatrix::identity(n, n);

        let alloc_diagonal = allocator.allocate(&assets, total_capital).unwrap();
        let alloc_identity = allocator
            .allocate_with_correlations(&assets, total_capital, &identity)
            .unwrap();

        // Both should produce identical allocations
        for asset in &assets {
            let diag_val = alloc_diagonal.get(&asset.symbol).unwrap();
            let ident_val = alloc_identity.get(&asset.symbol).unwrap();
            let diff = (*diag_val - *ident_val).abs();
            assert!(
                diff < Decimal::from_f64_retain(0.01).unwrap(),
                "Symbol {} differs: diagonal={}, identity={}",
                asset.symbol,
                diag_val,
                ident_val,
            );
        }
    }

    /// When two assets are highly correlated, the optimizer should allocate
    /// differently compared to the uncorrelated (diagonal) case.
    #[test]
    fn test_correlated_allocation_differs_from_diagonal() {
        let allocator = PortfolioAllocator::new(AllocationMethod::MeanVariance { lambda: 2.0 });
        let assets = create_test_assets(); // ES, NQ, ZN
        let total_capital = Decimal::from(100_000);
        let n = assets.len();

        // High correlation between ES and NQ (both equity futures), low with ZN (bonds)
        let corr_data = vec![
            1.0, 0.90, 0.10, // ES row
            0.90, 1.0, 0.10, // NQ row
            0.10, 0.10, 1.0, // ZN row
        ];
        let corr = DMatrix::from_row_slice(n, n, &corr_data);

        let alloc_diagonal = allocator.allocate(&assets, total_capital).unwrap();
        let alloc_correlated = allocator
            .allocate_with_correlations(&assets, total_capital, &corr)
            .unwrap();

        // Correlated allocation should differ from diagonal
        let mut any_differs = false;
        for asset in &assets {
            let diag_val = alloc_diagonal.get(&asset.symbol).unwrap();
            let corr_val = alloc_correlated.get(&asset.symbol).unwrap();
            if (*diag_val - *corr_val).abs() > Decimal::from_f64_retain(1.0).unwrap() {
                any_differs = true;
            }
        }
        assert!(
            any_differs,
            "Correlated allocation should differ from diagonal allocation"
        );

        // With high ES-NQ correlation, ZN (diversifier) should get relatively more weight
        // compared to the diagonal case
        let zn_diag = alloc_diagonal.get("ZN.FUT").unwrap();
        let zn_corr = alloc_correlated.get("ZN.FUT").unwrap();
        assert!(
            zn_corr > zn_diag,
            "ZN (uncorrelated diversifier) should get more weight with correlations: corr={}, diag={}",
            zn_corr,
            zn_diag,
        );
    }

    /// Correlation matrix with wrong dimensions should return an error.
    #[test]
    fn test_invalid_correlation_matrix_dimensions() {
        let allocator = PortfolioAllocator::new(AllocationMethod::MeanVariance { lambda: 2.0 });
        let assets = create_test_assets(); // 3 assets
        let total_capital = Decimal::from(100_000);

        // 2x2 matrix for 3 assets — wrong dimensions
        let bad_corr = DMatrix::identity(2, 2);
        let result = allocator.allocate_with_correlations(&assets, total_capital, &bad_corr);
        assert!(result.is_err(), "Should fail with mismatched dimensions");
        let err_msg = format!("{}", result.unwrap_err());
        assert!(
            err_msg.contains("do not match"),
            "Error should mention dimension mismatch: {}",
            err_msg
        );

        // 4x4 matrix for 3 assets — also wrong
        let bad_corr_large = DMatrix::identity(4, 4);
        let result = allocator.allocate_with_correlations(&assets, total_capital, &bad_corr_large);
        assert!(result.is_err(), "Should fail with oversized dimensions");
    }

    /// Non-square correlation matrix should also fail.
    #[test]
    fn test_non_square_correlation_matrix() {
        let allocator = PortfolioAllocator::new(AllocationMethod::MeanVariance { lambda: 2.0 });
        let assets = create_test_assets();
        let total_capital = Decimal::from(100_000);

        // 3x2 matrix — not square
        let bad_corr = DMatrix::zeros(3, 2);
        let result = allocator.allocate_with_correlations(&assets, total_capital, &bad_corr);
        assert!(result.is_err(), "Should fail with non-square matrix");
    }

    /// Allocate with correlations on non-MeanVariance methods should delegate
    /// to standard allocate (correlations ignored).
    #[test]
    fn test_correlations_ignored_for_equal_weight() {
        let allocator = PortfolioAllocator::new(AllocationMethod::EqualWeight);
        let assets = create_test_assets();
        let total_capital = Decimal::from(100_000);
        let n = assets.len();

        let corr = DMatrix::identity(n, n);
        let alloc_std = allocator.allocate(&assets, total_capital).unwrap();
        let alloc_corr = allocator
            .allocate_with_correlations(&assets, total_capital, &corr)
            .unwrap();

        for asset in &assets {
            assert_eq!(
                alloc_std.get(&asset.symbol),
                alloc_corr.get(&asset.symbol),
                "EqualWeight should ignore correlations for {}",
                asset.symbol
            );
        }
    }
}