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 rust_decimal::Decimal;
use std::collections::HashMap;
use nalgebra::{DMatrix, DVector};

/// 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),
        }
    }

    /// 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)
    fn mean_variance(
        &self,
        assets: &[AssetInfo],
        total_capital: Decimal,
        lambda: 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()
        );

        // Covariance matrix (simplified: diagonal with volatilities)
        // TODO: Add correlations for full covariance matrix
        let mut sigma = DMatrix::zeros(n, n);
        for (i, asset) in assets.iter().enumerate() {
            sigma[(i, i)] = asset.volatility.powi(2);
        }

        // 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].max(0.0).min(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].max(0.0).min(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)
                    .max(0.0)
                    .min(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)
    }
}

/// 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)]
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
                );
            }
        }
    }
}