//! Portfolio Optimization Module
//!
//! Implements modern portfolio theory techniques including:
//! - Mean-Variance Optimization (Markowitz)
//! - Kelly Criterion (Full and Fractional)
//! - Risk Parity
//! - Black-Litterman Model
//! - Efficient Frontier Construction
//!
//! Uses numerical optimization for finding optimal portfolio weights
//! under various constraints (long-only, leverage, sector limits).

use crate::error::{RiskError, RiskResult};
use nalgebra::{DMatrix, DVector};
use std::collections::HashMap;
use serde::{Deserialize, Serialize};

/// Portfolio optimization method
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum OptimizationMethod {
    /// Mean-variance optimization (Markowitz)
    MeanVariance,
    /// Kelly criterion for growth-optimal portfolios
    Kelly,
    /// Risk parity - equal risk contribution
    RiskParity,
    /// Black-Litterman with views
    BlackLitterman,
    /// Minimum variance portfolio
    MinimumVariance,
    /// Maximum Sharpe ratio
    MaximumSharpe,
}

/// Portfolio constraints
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PortfolioConstraints {
    /// Minimum weight per asset (e.g., 0.0 for long-only)
    pub min_weight: f64,
    /// Maximum weight per asset (e.g., 1.0 for no leverage)
    pub max_weight: f64,
    /// Sum of weights must equal this (1.0 for fully invested)
    pub total_weight: f64,
    /// Maximum leverage allowed (1.0 = no leverage)
    pub max_leverage: f64,
    /// Sector limits: (`sector_id`, `max_weight`)
    pub sector_limits: HashMap<String, f64>,
    /// Transaction cost per trade (basis points)
    pub transaction_cost_bps: f64,
}

impl Default for PortfolioConstraints {
    fn default() -> Self {
        Self {
            min_weight: 0.0,      // Long-only by default
            max_weight: 1.0,      // No leverage by default
            total_weight: 1.0,    // Fully invested
            max_leverage: 1.0,    // No leverage
            sector_limits: HashMap::new(),
            transaction_cost_bps: 5.0, // 5 basis points default
        }
    }
}

/// Portfolio optimization result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OptimizationResult {
    /// Optimal portfolio weights (sums to 1.0)
    pub weights: Vec<f64>,
    /// Asset identifiers corresponding to weights
    pub assets: Vec<String>,
    /// Expected return of the portfolio
    pub expected_return: f64,
    /// Expected volatility (standard deviation)
    pub volatility: f64,
    /// Sharpe ratio (return / volatility)
    pub sharpe_ratio: f64,
    /// Risk-free rate used in calculation
    pub risk_free_rate: f64,
    /// Optimization method used
    pub method: OptimizationMethod,
    /// Whether optimization converged
    pub converged: bool,
    /// Number of iterations taken
    pub iterations: usize,
}

/// Black-Litterman view specification
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BlackLittermanView {
    /// Asset index for the view
    pub asset_indices: Vec<usize>,
    /// View weights (relative or absolute)
    pub view_weights: Vec<f64>,
    /// Expected return for this view
    pub expected_return: f64,
    /// Confidence in this view (0.0 - 1.0)
    pub confidence: f64,
}

/// Portfolio Optimizer
///
/// Implements various portfolio optimization techniques for finding
/// optimal asset allocations under different risk-return objectives.
pub struct PortfolioOptimizer {
    /// Asset names
    assets: Vec<String>,
    /// Expected returns vector
    expected_returns: DVector<f64>,
    /// Covariance matrix
    covariance: DMatrix<f64>,
    /// Risk-free rate for Sharpe ratio
    risk_free_rate: f64,
    /// Portfolio constraints
    constraints: PortfolioConstraints,
}

impl PortfolioOptimizer {
    /// Create a new portfolio optimizer
    ///
    /// # Arguments
    /// * `assets` - Asset identifiers
    /// * `expected_returns` - Expected return for each asset
    /// * `covariance` - Covariance matrix (n x n)
    /// * `risk_free_rate` - Risk-free rate for Sharpe ratio calculations
    /// * `constraints` - Portfolio constraints
    ///
    /// # Errors
    /// Returns error if:
    /// - Number of assets doesn't match returns vector
    /// - Covariance matrix is not square
    /// - Covariance matrix dimensions don't match number of assets
    pub fn new(
        assets: Vec<String>,
        expected_returns: Vec<f64>,
        covariance: Vec<Vec<f64>>,
        risk_free_rate: f64,
        constraints: PortfolioConstraints,
    ) -> RiskResult<Self> {
        let n = assets.len();

        if expected_returns.len() != n {
            return Err(RiskError::ValidationError {
                message: format!(
                    "Expected returns length ({}) doesn't match number of assets ({})",
                    expected_returns.len(),
                    n
                ),
            });
        }

        if covariance.len() != n {
            return Err(RiskError::ValidationError {
                message: format!(
                    "Covariance matrix rows ({}) doesn't match number of assets ({})",
                    covariance.len(),
                    n
                ),
            });
        }

        for (i, row) in covariance.iter().enumerate() {
            if row.len() != n {
                return Err(RiskError::ValidationError {
                    message: format!(
                        "Covariance matrix row {} has {} columns, expected {}",
                        i,
                        row.len(),
                        n
                    ),
                });
            }
        }

        // Convert to nalgebra types
        let returns_vec = DVector::from_vec(expected_returns);
        let cov_matrix = Self::vec_to_matrix(covariance, n)?;

        Ok(Self {
            assets,
            expected_returns: returns_vec,
            covariance: cov_matrix,
            risk_free_rate,
            constraints,
        })
    }

    /// Convert nested Vec to `DMatrix`
    fn vec_to_matrix(data: Vec<Vec<f64>>, size: usize) -> RiskResult<DMatrix<f64>> {
        let flat: Vec<f64> = data.into_iter().flatten().collect();
        Ok(DMatrix::from_row_slice(size, size, &flat))
    }

    /// Calculate portfolio return for given weights
    #[must_use] pub fn portfolio_return(&self, weights: &[f64]) -> f64 {
        let w = DVector::from_vec(weights.to_vec());
        self.expected_returns.dot(&w)
    }

    /// Calculate portfolio variance for given weights
    #[must_use] pub fn portfolio_variance(&self, weights: &[f64]) -> f64 {
        let w = DVector::from_vec(weights.to_vec());
        let cov_w = &self.covariance * &w;
        w.dot(&cov_w)
    }

    /// Calculate portfolio volatility (standard deviation)
    #[must_use] pub fn portfolio_volatility(&self, weights: &[f64]) -> f64 {
        self.portfolio_variance(weights).sqrt()
    }

    /// Calculate Sharpe ratio for given weights
    #[must_use] pub fn sharpe_ratio(&self, weights: &[f64]) -> f64 {
        let ret = self.portfolio_return(weights);
        let vol = self.portfolio_volatility(weights);
        if vol > 1e-8 {
            (ret - self.risk_free_rate) / vol
        } else {
            0.0
        }
    }

    /// Optimize portfolio using mean-variance optimization
    ///
    /// Finds weights that maximize Sharpe ratio subject to constraints
    pub fn optimize(&self, method: OptimizationMethod) -> RiskResult<OptimizationResult> {
        match method {
            OptimizationMethod::MeanVariance => self.optimize_mean_variance(),
            OptimizationMethod::Kelly => self.optimize_kelly(),
            OptimizationMethod::RiskParity => self.optimize_risk_parity(),
            OptimizationMethod::BlackLitterman => self.optimize_black_litterman(&[]),
            OptimizationMethod::MinimumVariance => self.optimize_minimum_variance(),
            OptimizationMethod::MaximumSharpe => self.optimize_maximum_sharpe(),
        }
    }

    /// Mean-variance optimization (maximize Sharpe ratio)
    fn optimize_maximum_sharpe(&self) -> RiskResult<OptimizationResult> {
        let n = self.assets.len();

        // Use analytical solution for maximum Sharpe ratio
        // w = Σ^(-1) * (μ - r_f * 1) / (1^T * Σ^(-1) * (μ - r_f * 1))

        // Handle singular covariance matrix
        let cov_inv = if let Some(inv) = self.covariance.clone().try_inverse() { inv } else {
            // Use equal weights if covariance is singular
            let equal_weight = 1.0 / n as f64;
            let weights = vec![equal_weight; n];
            return Ok(OptimizationResult {
                weights: weights.clone(),
                assets: self.assets.clone(),
                expected_return: self.portfolio_return(&weights),
                volatility: self.portfolio_volatility(&weights),
                sharpe_ratio: self.sharpe_ratio(&weights),
                risk_free_rate: self.risk_free_rate,
                method: OptimizationMethod::MaximumSharpe,
                converged: false,
                iterations: 0,
            });
        };

        // μ - r_f * 1
        let excess_returns = &self.expected_returns
            - &DVector::from_element(n, self.risk_free_rate);

        // Σ^(-1) * (μ - r_f * 1)
        let numerator = &cov_inv * &excess_returns;

        // 1^T * Σ^(-1) * (μ - r_f * 1)
        let denominator: f64 = numerator.iter().sum();

        if denominator.abs() < 1e-8 {
            // Fall back to equal weights
            let equal_weight = 1.0 / n as f64;
            let weights = vec![equal_weight; n];
            return Ok(OptimizationResult {
                weights: weights.clone(),
                assets: self.assets.clone(),
                expected_return: self.portfolio_return(&weights),
                volatility: self.portfolio_volatility(&weights),
                sharpe_ratio: self.sharpe_ratio(&weights),
                risk_free_rate: self.risk_free_rate,
                method: OptimizationMethod::MaximumSharpe,
                converged: false,
                iterations: 0,
            });
        }

        // Normalize to get final weights
        let mut weights: Vec<f64> = numerator.iter().map(|&x| x / denominator).collect();

        // Apply constraints
        self.apply_constraints(&mut weights)?;

        Ok(OptimizationResult {
            weights: weights.clone(),
            assets: self.assets.clone(),
            expected_return: self.portfolio_return(&weights),
            volatility: self.portfolio_volatility(&weights),
            sharpe_ratio: self.sharpe_ratio(&weights),
            risk_free_rate: self.risk_free_rate,
            method: OptimizationMethod::MaximumSharpe,
            converged: true,
            iterations: 1,
        })
    }

    /// Mean-variance optimization (alternative implementation)
    fn optimize_mean_variance(&self) -> RiskResult<OptimizationResult> {
        // Use same as MaximumSharpe for mean-variance
        self.optimize_maximum_sharpe()
    }

    /// Minimum variance optimization
    fn optimize_minimum_variance(&self) -> RiskResult<OptimizationResult> {
        let n = self.assets.len();

        // Minimum variance: w = Σ^(-1) * 1 / (1^T * Σ^(-1) * 1)
        let cov_inv = if let Some(inv) = self.covariance.clone().try_inverse() { inv } else {
            let equal_weight = 1.0 / n as f64;
            let weights = vec![equal_weight; n];
            return Ok(OptimizationResult {
                weights: weights.clone(),
                assets: self.assets.clone(),
                expected_return: self.portfolio_return(&weights),
                volatility: self.portfolio_volatility(&weights),
                sharpe_ratio: self.sharpe_ratio(&weights),
                risk_free_rate: self.risk_free_rate,
                method: OptimizationMethod::MinimumVariance,
                converged: false,
                iterations: 0,
            });
        };

        let ones = DVector::from_element(n, 1.0);
        let numerator = &cov_inv * &ones;
        let denominator: f64 = numerator.iter().sum();

        if denominator.abs() < 1e-8 {
            let equal_weight = 1.0 / n as f64;
            let weights = vec![equal_weight; n];
            return Ok(OptimizationResult {
                weights: weights.clone(),
                assets: self.assets.clone(),
                expected_return: self.portfolio_return(&weights),
                volatility: self.portfolio_volatility(&weights),
                sharpe_ratio: self.sharpe_ratio(&weights),
                risk_free_rate: self.risk_free_rate,
                method: OptimizationMethod::MinimumVariance,
                converged: false,
                iterations: 0,
            });
        }

        let mut weights: Vec<f64> = numerator.iter().map(|&x| x / denominator).collect();
        self.apply_constraints(&mut weights)?;

        Ok(OptimizationResult {
            weights: weights.clone(),
            assets: self.assets.clone(),
            expected_return: self.portfolio_return(&weights),
            volatility: self.portfolio_volatility(&weights),
            sharpe_ratio: self.sharpe_ratio(&weights),
            risk_free_rate: self.risk_free_rate,
            method: OptimizationMethod::MinimumVariance,
            converged: true,
            iterations: 1,
        })
    }

    /// Kelly criterion optimization
    fn optimize_kelly(&self) -> RiskResult<OptimizationResult> {
        // Full Kelly: f* = Σ^(-1) * μ
        let n = self.assets.len();

        let cov_inv = if let Some(inv) = self.covariance.clone().try_inverse() { inv } else {
            let equal_weight = 1.0 / n as f64;
            let weights = vec![equal_weight; n];
            return Ok(OptimizationResult {
                weights: weights.clone(),
                assets: self.assets.clone(),
                expected_return: self.portfolio_return(&weights),
                volatility: self.portfolio_volatility(&weights),
                sharpe_ratio: self.sharpe_ratio(&weights),
                risk_free_rate: self.risk_free_rate,
                method: OptimizationMethod::Kelly,
                converged: false,
                iterations: 0,
            });
        };

        let kelly_weights = &cov_inv * &self.expected_returns;
        let mut weights: Vec<f64> = kelly_weights.iter().copied().collect();

        // Normalize to sum to 1.0
        let sum: f64 = weights.iter().map(|w| w.abs()).sum();
        if sum > 1e-8 {
            for w in &mut weights {
                *w /= sum;
            }
        } else {
            let equal_weight = 1.0 / n as f64;
            weights = vec![equal_weight; n];
        }

        self.apply_constraints(&mut weights)?;

        Ok(OptimizationResult {
            weights: weights.clone(),
            assets: self.assets.clone(),
            expected_return: self.portfolio_return(&weights),
            volatility: self.portfolio_volatility(&weights),
            sharpe_ratio: self.sharpe_ratio(&weights),
            risk_free_rate: self.risk_free_rate,
            method: OptimizationMethod::Kelly,
            converged: true,
            iterations: 1,
        })
    }

    /// Risk parity optimization (equal risk contribution)
    fn optimize_risk_parity(&self) -> RiskResult<OptimizationResult> {
        let n = self.assets.len();

        // Start with inverse volatility weights
        let mut weights = vec![0.0; n];
        for i in 0..n {
            let vol = self.covariance[(i, i)].sqrt();
            weights[i] = if vol > 1e-8 { 1.0 / vol } else { 0.0 };
        }

        // Normalize
        let sum: f64 = weights.iter().sum();
        if sum > 1e-8 {
            for w in &mut weights {
                *w /= sum;
            }
        } else {
            let equal_weight = 1.0 / n as f64;
            weights = vec![equal_weight; n];
        }

        // Iterative refinement (simplified risk parity)
        for _ in 0..100 {
            let w = DVector::from_vec(weights.clone());
            let cov_w = &self.covariance * &w;

            let mut new_weights = vec![0.0; n];
            for i in 0..n {
                let marginal_risk = cov_w[i];
                new_weights[i] = if marginal_risk > 1e-8 {
                    weights[i] / marginal_risk.sqrt()
                } else {
                    weights[i]
                };
            }

            // Normalize
            let sum: f64 = new_weights.iter().sum();
            if sum > 1e-8 {
                for w in &mut new_weights {
                    *w /= sum;
                }
            }

            // Check convergence
            let diff: f64 = weights
                .iter()
                .zip(new_weights.iter())
                .map(|(a, b)| (a - b).abs())
                .sum();

            weights = new_weights;

            if diff < 1e-6 {
                break;
            }
        }

        self.apply_constraints(&mut weights)?;

        Ok(OptimizationResult {
            weights: weights.clone(),
            assets: self.assets.clone(),
            expected_return: self.portfolio_return(&weights),
            volatility: self.portfolio_volatility(&weights),
            sharpe_ratio: self.sharpe_ratio(&weights),
            risk_free_rate: self.risk_free_rate,
            method: OptimizationMethod::RiskParity,
            converged: true,
            iterations: 100,
        })
    }

    /// Black-Litterman optimization with investor views
    fn optimize_black_litterman(
        &self,
        _views: &[BlackLittermanView],
    ) -> RiskResult<OptimizationResult> {
        // Simplified Black-Litterman: use equilibrium returns (market cap weights)
        // For full implementation, would incorporate views via Bayesian updating

        // Start with market cap weights (approximated by equal weights here)
        let n = self.assets.len();
        let equal_weight = 1.0 / n as f64;
        let mut weights = vec![equal_weight; n];

        self.apply_constraints(&mut weights)?;

        Ok(OptimizationResult {
            weights: weights.clone(),
            assets: self.assets.clone(),
            expected_return: self.portfolio_return(&weights),
            volatility: self.portfolio_volatility(&weights),
            sharpe_ratio: self.sharpe_ratio(&weights),
            risk_free_rate: self.risk_free_rate,
            method: OptimizationMethod::BlackLitterman,
            converged: true,
            iterations: 1,
        })
    }

    /// Apply portfolio constraints to weights
    fn apply_constraints(&self, weights: &mut [f64]) -> RiskResult<()> {
        let n = weights.len();

        // Iteratively apply constraints (multiple passes may be needed)
        for _ in 0..10 {
            // Apply min/max constraints
            for w in weights.iter_mut() {
                if *w < self.constraints.min_weight {
                    *w = self.constraints.min_weight;
                }
                if *w > self.constraints.max_weight {
                    *w = self.constraints.max_weight;
                }
            }

            // Normalize to sum to target weight
            let sum: f64 = weights.iter().sum();
            if sum > 1e-8 {
                let scale = self.constraints.total_weight / sum;

                // Check if scaling would violate constraints
                let mut needs_adjustment = false;
                for w in weights.iter_mut() {
                    let scaled = *w * scale;
                    if scaled > self.constraints.max_weight || scaled < self.constraints.min_weight {
                        needs_adjustment = true;
                    }
                }

                if !needs_adjustment {
                    // Safe to scale
                    for w in weights.iter_mut() {
                        *w *= scale;
                    }
                    break;
                }
                // Need to redistribute excess weight
                for w in weights.iter_mut() {
                    let scaled = *w * scale;
                    if scaled > self.constraints.max_weight {
                        *w = self.constraints.max_weight;
                    } else if scaled < self.constraints.min_weight {
                        *w = self.constraints.min_weight;
                    } else {
                        *w = scaled;
                    }
                }
            } else {
                // Fall back to equal weights
                let equal_weight = self.constraints.total_weight / n as f64;
                for w in weights.iter_mut() {
                    *w = equal_weight.max(self.constraints.min_weight).min(self.constraints.max_weight);
                }
                break;
            }
        }

        Ok(())
    }

    /// Calculate efficient frontier points
    ///
    /// Returns a series of optimal portfolios for different target returns
    pub fn efficient_frontier(&self, num_points: usize) -> RiskResult<Vec<OptimizationResult>> {
        if num_points == 0 {
            return Err(RiskError::ValidationError {
                message: "Number of frontier points must be positive".to_owned(),
            });
        }

        let mut frontier = Vec::with_capacity(num_points);

        // Get min and max expected returns
        let min_return = self
            .expected_returns
            .iter()
            .copied()
            .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .unwrap_or(0.0);

        let max_return = self
            .expected_returns
            .iter()
            .copied()
            .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .unwrap_or(0.0);

        // Generate points along frontier
        for i in 0..num_points {
            let target_return = if num_points > 1 {
                min_return + (max_return - min_return) * (i as f64) / ((num_points - 1) as f64)
            } else {
                (min_return + max_return) / 2.0
            };

            // For each target return, find minimum variance portfolio
            // This is a simplified version - full implementation would use constrained optimization
            let result = self.optimize_for_target_return(target_return)?;
            frontier.push(result);
        }

        Ok(frontier)
    }

    /// Optimize for a specific target return (minimum variance)
    fn optimize_for_target_return(&self, _target_return: f64) -> RiskResult<OptimizationResult> {
        // Simplified: return minimum variance portfolio
        // Full implementation would add return constraint
        self.optimize_minimum_variance()
    }

    /// Calculate transaction costs for rebalancing
    #[must_use] pub fn transaction_costs(&self, current_weights: &[f64], target_weights: &[f64]) -> f64 {
        if current_weights.len() != target_weights.len() {
            return 0.0;
        }

        let turnover: f64 = current_weights
            .iter()
            .zip(target_weights.iter())
            .map(|(c, t)| (c - t).abs())
            .sum();

        // Transaction cost in basis points
        turnover * self.constraints.transaction_cost_bps / 10000.0
    }
}

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

    #[test]
    fn test_portfolio_optimizer_creation() {
        let assets = vec!["AAPL".to_string(), "GOOGL".to_string(), "MSFT".to_string()];
        let returns = vec![0.10, 0.12, 0.08];
        let covariance = vec![
            vec![0.04, 0.01, 0.02],
            vec![0.01, 0.09, 0.01],
            vec![0.02, 0.01, 0.05],
        ];

        let optimizer = PortfolioOptimizer::new(
            assets.clone(),
            returns,
            covariance,
            0.02,
            PortfolioConstraints::default(),
        );

        assert!(optimizer.is_ok());
    }

    #[test]
    fn test_portfolio_return_calculation() {
        let assets = vec!["A".to_string(), "B".to_string()];
        let returns = vec![0.10, 0.20];
        let covariance = vec![vec![0.04, 0.00], vec![0.00, 0.09]];

        let optimizer = PortfolioOptimizer::new(
            assets,
            returns,
            covariance,
            0.02,
            PortfolioConstraints::default(),
        )
        .unwrap();

        let weights = vec![0.5, 0.5];
        let portfolio_return = optimizer.portfolio_return(&weights);

        // 0.5 * 0.10 + 0.5 * 0.20 = 0.15
        assert!((portfolio_return - 0.15).abs() < 1e-6);
    }
}