//! Covariance matrix estimation for portfolio optimization

use crate::{Asset, PortfolioError, Result};
use nalgebra::DMatrix;

/// Estimate covariance matrix from asset volatilities and correlation structure
///
/// For demonstration purposes, this uses a simple correlation model.
/// In production, you would estimate from historical returns.
pub fn estimate_covariance_matrix(assets: &[Asset]) -> Result<DMatrix<f64>> {
    if assets.len() < 2 {
        return Err(PortfolioError::InsufficientData(
            "Need at least 2 assets to compute covariance".to_string(),
        ));
    }

    let n = assets.len();
    let mut cov_matrix = DMatrix::zeros(n, n);

    for i in 0..n {
        for j in 0..n {
            if i == j {
                cov_matrix[(i, j)] = assets[i].volatility * assets[i].volatility;
            } else {
                let correlation = estimate_correlation(&assets[i], &assets[j]);
                cov_matrix[(i, j)] = correlation * assets[i].volatility * assets[j].volatility;
            }
        }
    }

    Ok(cov_matrix)
}

/// Estimate correlation between two assets based on their characteristics
///
/// This is a simplified correlation model for demonstration.
/// Real implementations would use historical data.
fn estimate_correlation(asset1: &Asset, asset2: &Asset) -> f64 {
    if asset1.asset_class == asset2.asset_class {
        if asset1.sector == asset2.sector {
            0.7
        } else {
            0.5
        }
    } else if asset1.asset_class == "bond" || asset2.asset_class == "bond" {
        0.2
    } else {
        0.3
    }
}

/// Compute portfolio variance given weights and covariance matrix
pub fn portfolio_variance(weights: &[f64], covariance: &DMatrix<f64>) -> Result<f64> {
    if weights.len() != covariance.nrows() || weights.len() != covariance.ncols() {
        return Err(PortfolioError::InvalidConfig(
            "Weight vector size must match covariance matrix dimensions".to_string(),
        ));
    }

    let mut variance = 0.0;
    for i in 0..weights.len() {
        for j in 0..weights.len() {
            variance += weights[i] * weights[j] * covariance[(i, j)];
        }
    }

    if variance < 0.0 {
        return Err(PortfolioError::NumericalError(
            "Negative variance computed".to_string(),
        ));
    }

    Ok(variance)
}

/// Compute risk contribution of each asset to total portfolio risk
pub fn risk_contributions(weights: &[f64], covariance: &DMatrix<f64>) -> Result<Vec<f64>> {
    let variance = portfolio_variance(weights, covariance)?;
    if variance == 0.0 {
        return Ok(vec![0.0; weights.len()]);
    }

    let volatility = variance.sqrt();
    let mut contributions = vec![0.0; weights.len()];

    for i in 0..weights.len() {
        let mut marginal_contribution = 0.0;
        for j in 0..weights.len() {
            marginal_contribution += weights[j] * covariance[(i, j)];
        }
        contributions[i] = (weights[i] * marginal_contribution) / volatility;
    }

    Ok(contributions)
}

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

    fn create_test_assets() -> Vec<Asset> {
        vec![
            Asset {
                ticker: "STOCK1".to_string(),
                name: "Stock 1".to_string(),
                asset_class: "equity".to_string(),
                sector: Some("Technology".to_string()),
                expected_return: 0.10,
                volatility: 0.20,
            },
            Asset {
                ticker: "STOCK2".to_string(),
                name: "Stock 2".to_string(),
                asset_class: "equity".to_string(),
                sector: Some("Healthcare".to_string()),
                expected_return: 0.08,
                volatility: 0.15,
            },
            Asset {
                ticker: "BOND1".to_string(),
                name: "Bond 1".to_string(),
                asset_class: "bond".to_string(),
                sector: Some("Fixed Income".to_string()),
                expected_return: 0.03,
                volatility: 0.05,
            },
        ]
    }

    #[test]
    fn test_estimate_covariance_matrix() {
        let assets = create_test_assets();
        let cov_matrix = estimate_covariance_matrix(&assets).unwrap();

        assert_eq!(cov_matrix.nrows(), 3);
        assert_eq!(cov_matrix.ncols(), 3);

        assert!((cov_matrix[(0, 0)] - 0.04).abs() < 1e-10);
        assert!((cov_matrix[(1, 1)] - 0.0225).abs() < 1e-10);
        assert!((cov_matrix[(2, 2)] - 0.0025).abs() < 1e-10);
        assert!(cov_matrix[(0, 1)] > 0.0);
        assert_eq!(cov_matrix[(0, 1)], cov_matrix[(1, 0)]);
    }

    #[test]
    fn test_estimate_covariance_insufficient_data() {
        let assets = vec![Asset {
            ticker: "STOCK1".to_string(),
            name: "Stock 1".to_string(),
            asset_class: "equity".to_string(),
            sector: None,
            expected_return: 0.10,
            volatility: 0.20,
        }];

        let result = estimate_covariance_matrix(&assets);
        assert!(result.is_err());
    }

    #[test]
    fn test_estimate_correlation() {
        let stock1 = Asset {
            ticker: "STOCK1".to_string(),
            name: "Stock 1".to_string(),
            asset_class: "equity".to_string(),
            sector: Some("Technology".to_string()),
            expected_return: 0.10,
            volatility: 0.20,
        };

        let stock2_same_sector = Asset {
            ticker: "STOCK2".to_string(),
            name: "Stock 2".to_string(),
            asset_class: "equity".to_string(),
            sector: Some("Technology".to_string()),
            expected_return: 0.12,
            volatility: 0.22,
        };

        let stock3_diff_sector = Asset {
            ticker: "STOCK3".to_string(),
            name: "Stock 3".to_string(),
            asset_class: "equity".to_string(),
            sector: Some("Healthcare".to_string()),
            expected_return: 0.08,
            volatility: 0.15,
        };

        let bond = Asset {
            ticker: "BOND1".to_string(),
            name: "Bond 1".to_string(),
            asset_class: "bond".to_string(),
            sector: Some("Fixed Income".to_string()),
            expected_return: 0.03,
            volatility: 0.05,
        };

        assert_eq!(estimate_correlation(&stock1, &stock2_same_sector), 0.7);
        assert_eq!(estimate_correlation(&stock1, &stock3_diff_sector), 0.5);
        assert_eq!(estimate_correlation(&stock1, &bond), 0.2);
    }

    #[test]
    fn test_portfolio_variance() {
        let assets = create_test_assets();
        let cov_matrix = estimate_covariance_matrix(&assets).unwrap();
        let weights = vec![0.6, 0.3, 0.1];

        let variance = portfolio_variance(&weights, &cov_matrix).unwrap();
        assert!(variance > 0.0);
    }

    #[test]
    fn test_portfolio_variance_dimension_mismatch() {
        let assets = create_test_assets();
        let cov_matrix = estimate_covariance_matrix(&assets).unwrap();
        let weights = vec![0.5, 0.5];

        let result = portfolio_variance(&weights, &cov_matrix);
        assert!(result.is_err());
    }

    #[test]
    fn test_risk_contributions() {
        let assets = create_test_assets();
        let cov_matrix = estimate_covariance_matrix(&assets).unwrap();
        let weights = vec![0.6, 0.3, 0.1];

        let contributions = risk_contributions(&weights, &cov_matrix).unwrap();
        assert_eq!(contributions.len(), 3);
        assert!(contributions.iter().all(|&c| c >= 0.0));
        let total_contribution: f64 = contributions.iter().sum();
        assert!(
            (total_contribution - portfolio_variance(&weights, &cov_matrix).unwrap().sqrt()).abs()
                < 1e-10
        );
    }
}