use std::collections::HashMap;
use nalgebra::{DMatrix, DVector};
use crate::positions::Position;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PortfolioRisk {
    pub total_var_95: f64,
    pub total_var_99: f64,
    pub total_cvar_95: f64,
    pub marginal_var: HashMap<String, f64>,
    pub component_var: HashMap<String, f64>,
    pub correlation_matrix: CorrelationMatrix,
    pub concentration_risk: ConcentrationMetrics,
    pub stress_test_results: HashMap<String, f64>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CorrelationMatrix {
    pub symbols: Vec<String>,
    pub matrix: Vec<Vec<f64>>,
    pub average_correlation: f64,
    pub max_correlation: (String, String, f64),
    pub clustering_score: f64, // How clustered the portfolio is
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConcentrationMetrics {
    pub herfindahl_index: f64,
    pub effective_number_of_positions: f64,
    pub top_5_concentration: f64,
    pub sector_concentration: HashMap<String, f64>,
}

#[derive(Debug, Clone)]
pub struct RiskModel {
    returns_history: HashMap<String, Vec<f64>>,
    lookback_period: usize,
    confidence_levels: Vec<f64>,
}

impl RiskModel {
    pub fn new(lookback_period: usize) -> Self {
        Self {
            returns_history: HashMap::new(),
            lookback_period,
            confidence_levels: vec![0.95, 0.99],
        }
    }

    pub fn update_returns(&mut self, symbol: &str, returns: Vec<f64>) {
        self.returns_history.insert(symbol.to_string(), returns);
    }

    pub fn calculate_portfolio_risk(
        &self,
        positions: &HashMap<String, Position>,
        current_prices: &HashMap<String, f64>,
    ) -> Result<PortfolioRisk, String> {
        // Get active symbols and their weights
        let (symbols, weights) = self.get_portfolio_weights(positions, current_prices)?;
        
        // Calculate covariance matrix
        let cov_matrix = self.calculate_covariance_matrix(&symbols)?;
        
        // Portfolio variance
        let portfolio_variance = self.calculate_portfolio_variance(&weights, &cov_matrix);
        let portfolio_vol = portfolio_variance.sqrt();
        
        // VaR calculations
        let total_portfolio_value = self.calculate_total_value(positions, current_prices);
        let total_var_95 = total_portfolio_value * portfolio_vol * 1.645; // 95% confidence
        let total_var_99 = total_portfolio_value * portfolio_vol * 2.326; // 99% confidence
        
        // CVaR (Conditional VaR)
        let total_cvar_95 = self.calculate_cvar(&symbols, &weights, 0.95)?;
        
        // Marginal and Component VaR
        let (marginal_var, component_var) = self.calculate_var_decomposition(
            &symbols,
            &weights,
            &cov_matrix,
            portfolio_vol,
            total_portfolio_value,
        )?;
        
        // Correlation analysis
        let correlation_matrix = self.calculate_correlation_matrix(&symbols)?;
        
        // Concentration metrics
        let concentration_risk = self.calculate_concentration_metrics(positions, current_prices);
        
        // Stress tests
        let stress_test_results = self.run_stress_tests(positions, current_prices, &cov_matrix)?;
        
        Ok(PortfolioRisk {
            total_var_95,
            total_var_99,
            total_cvar_95,
            marginal_var,
            component_var,
            correlation_matrix,
            concentration_risk,
            stress_test_results,
        })
    }

    fn get_portfolio_weights(
        &self,
        positions: &HashMap<String, Position>,
        current_prices: &HashMap<String, f64>,
    ) -> Result<(Vec<String>, DVector<f64>), String> {
        let mut symbols = Vec::new();
        let mut values = Vec::new();
        let mut total_value = 0.0;
        
        for (symbol, position) in positions {
            if let Some(&price) = current_prices.get(symbol) {
                let position_value = position.quantity * price;
                symbols.push(symbol.clone());
                values.push(position_value.abs());
                total_value += position_value.abs();
            }
        }
        
        if total_value == 0.0 {
            return Err("Portfolio has zero value".to_string());
        }
        
        let weights = DVector::from_vec(values.iter().map(|v| v / total_value).collect());
        
        Ok((symbols, weights))
    }

    fn calculate_covariance_matrix(&self, symbols: &[String]) -> Result<DMatrix<f64>, String> {
        let n = symbols.len();
        let mut matrix = DMatrix::zeros(n, n);
        
        for (i, symbol_i) in symbols.iter().enumerate() {
            for (j, symbol_j) in symbols.iter().enumerate() {
                if i <= j {
                    let cov = self.calculate_covariance(symbol_i, symbol_j)?;
                    matrix[(i, j)] = cov;
                    matrix[(j, i)] = cov; // Symmetric
                }
            }
        }
        
        Ok(matrix)
    }

    fn calculate_covariance(&self, symbol1: &str, symbol2: &str) -> Result<f64, String> {
        let returns1 = self.returns_history.get(symbol1)
            .ok_or_else(|| format!("No returns data for {}", symbol1))?;
        let returns2 = self.returns_history.get(symbol2)
            .ok_or_else(|| format!("No returns data for {}", symbol2))?;
        
        if returns1.len() != returns2.len() {
            return Err("Mismatched returns length".to_string());
        }
        
        let n = returns1.len() as f64;
        let mean1 = returns1.iter().sum::<f64>() / n;
        let mean2 = returns2.iter().sum::<f64>() / n;
        
        let covariance = returns1.iter()
            .zip(returns2.iter())
            .map(|(r1, r2)| (r1 - mean1) * (r2 - mean2))
            .sum::<f64>() / (n - 1.0);
        
        Ok(covariance)
    }

    fn calculate_portfolio_variance(&self, weights: &DVector<f64>, cov_matrix: &DMatrix<f64>) -> f64 {
        let variance = weights.transpose() * cov_matrix * weights;
        variance[(0, 0)]
    }

    fn calculate_var_decomposition(
        &self,
        symbols: &[String],
        weights: &DVector<f64>,
        cov_matrix: &DMatrix<f64>,
        portfolio_vol: f64,
        total_value: f64,
    ) -> Result<(HashMap<String, f64>, HashMap<String, f64>), String> {
        let mut marginal_var = HashMap::new();
        let mut component_var = HashMap::new();
        
        // Marginal VaR = ∂VaR/∂w_i
        let cov_weights = cov_matrix * weights;
        
        for (i, symbol) in symbols.iter().enumerate() {
            let marginal = (cov_weights[i] / portfolio_vol) * 1.645 * total_value;
            let component = marginal * weights[i];
            
            marginal_var.insert(symbol.clone(), marginal);
            component_var.insert(symbol.clone(), component);
        }
        
        Ok((marginal_var, component_var))
    }

    fn calculate_cvar(&self, symbols: &[String], weights: &DVector<f64>, confidence: f64) -> Result<f64, String> {
        // Simulate portfolio returns
        let portfolio_returns = self.simulate_portfolio_returns(symbols, weights, 10000)?;
        
        // Sort returns
        let mut sorted_returns = portfolio_returns.clone();
        sorted_returns.sort_by(|a, b| a.partial_cmp(b).unwrap());
        
        // Find VaR threshold
        let var_index = ((1.0 - confidence) * sorted_returns.len() as f64) as usize;
        let var_threshold = sorted_returns[var_index];
        
        // Calculate expected loss beyond VaR
        let tail_losses: Vec<f64> = sorted_returns.iter()
            .take(var_index)
            .cloned()
            .collect();
        
        if tail_losses.is_empty() {
            return Ok(0.0);
        }
        
        let cvar = -tail_losses.iter().sum::<f64>() / tail_losses.len() as f64;
        
        Ok(cvar)
    }

    fn simulate_portfolio_returns(
        &self,
        symbols: &[String],
        weights: &DVector<f64>,
        num_simulations: usize,
    ) -> Result<Vec<f64>, String> {
        let mut portfolio_returns = Vec::with_capacity(num_simulations);
        
        // Use historical simulation
        let min_length = symbols.iter()
            .map(|s| self.returns_history.get(s).map(|r| r.len()).unwrap_or(0))
            .min()
            .unwrap_or(0);
        
        if min_length == 0 {
            return Err("No returns data available".to_string());
        }
        
        // Bootstrap from historical returns
        use rand::prelude::*;
        let mut rng = thread_rng();
        
        for _ in 0..num_simulations {
            let idx = rng.gen_range(0..min_length);
            let mut portfolio_return = 0.0;
            
            for (i, symbol) in symbols.iter().enumerate() {
                if let Some(returns) = self.returns_history.get(symbol) {
                    portfolio_return += weights[i] * returns[idx];
                }
            }
            
            portfolio_returns.push(portfolio_return);
        }
        
        Ok(portfolio_returns)
    }

    fn calculate_correlation_matrix(&self, symbols: &[String]) -> Result<CorrelationMatrix, String> {
        let n = symbols.len();
        let mut matrix = vec![vec![0.0; n]; n];
        let mut sum_correlation = 0.0;
        let mut count = 0;
        let mut max_correlation = ("".to_string(), "".to_string(), 0.0);
        
        for (i, symbol_i) in symbols.iter().enumerate() {
            for (j, symbol_j) in symbols.iter().enumerate() {
                if i <= j {
                    let corr = if i == j {
                        1.0
                    } else {
                        self.calculate_correlation(symbol_i, symbol_j)?
                    };
                    
                    matrix[i][j] = corr;
                    matrix[j][i] = corr;
                    
                    if i != j {
                        sum_correlation += corr.abs();
                        count += 1;
                        
                        if corr.abs() > max_correlation.2 {
                            max_correlation = (symbol_i.clone(), symbol_j.clone(), corr);
                        }
                    }
                }
            }
        }
        
        let average_correlation = if count > 0 {
            sum_correlation / count as f64
        } else {
            0.0
        };
        
        // Calculate clustering score (higher = more clustered)
        let clustering_score = self.calculate_clustering_score(&matrix);
        
        Ok(CorrelationMatrix {
            symbols: symbols.to_vec(),
            matrix,
            average_correlation,
            max_correlation,
            clustering_score,
        })
    }

    fn calculate_correlation(&self, symbol1: &str, symbol2: &str) -> Result<f64, String> {
        let returns1 = self.returns_history.get(symbol1)
            .ok_or_else(|| format!("No returns data for {}", symbol1))?;
        let returns2 = self.returns_history.get(symbol2)
            .ok_or_else(|| format!("No returns data for {}", symbol2))?;
        
        let cov = self.calculate_covariance(symbol1, symbol2)?;
        let std1 = self.calculate_std_dev(returns1);
        let std2 = self.calculate_std_dev(returns2);
        
        if std1 == 0.0 || std2 == 0.0 {
            return Ok(0.0);
        }
        
        Ok(cov / (std1 * std2))
    }

    fn calculate_std_dev(&self, returns: &[f64]) -> f64 {
        let n = returns.len() as f64;
        let mean = returns.iter().sum::<f64>() / n;
        let variance = returns.iter()
            .map(|r| (r - mean).powi(2))
            .sum::<f64>() / (n - 1.0);
        
        variance.sqrt()
    }

    fn calculate_clustering_score(&self, correlation_matrix: &[Vec<f64>]) -> f64 {
        // Use average linkage clustering metric
        let n = correlation_matrix.len();
        if n < 2 {
            return 0.0;
        }
        
        let mut cluster_sum = 0.0;
        let mut cluster_count = 0;
        
        // Look for groups of highly correlated assets
        for i in 0..n {
            for j in i+1..n {
                for k in j+1..n {
                    let corr_ij = correlation_matrix[i][j].abs();
                    let corr_ik = correlation_matrix[i][k].abs();
                    let corr_jk = correlation_matrix[j][k].abs();
                    
                    // If all three are highly correlated, they form a cluster
                    let min_corr = corr_ij.min(corr_ik).min(corr_jk);
                    if min_corr > 0.5 {
                        cluster_sum += min_corr;
                        cluster_count += 1;
                    }
                }
            }
        }
        
        if cluster_count > 0 {
            cluster_sum / cluster_count as f64
        } else {
            0.0
        }
    }

    fn calculate_concentration_metrics(
        &self,
        positions: &HashMap<String, Position>,
        current_prices: &HashMap<String, f64>,
    ) -> ConcentrationMetrics {
        let mut position_values: Vec<(String, f64)> = Vec::new();
        let mut total_value = 0.0;
        
        for (symbol, position) in positions {
            if let Some(&price) = current_prices.get(symbol) {
                let value = (position.quantity * price).abs();
                position_values.push((symbol.clone(), value));
                total_value += value;
            }
        }
        
        // Sort by value descending
        position_values.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
        
        // Herfindahl Index
        let herfindahl_index = position_values.iter()
            .map(|(_, value)| {
                let weight = value / total_value;
                weight * weight
            })
            .sum();
        
        // Effective number of positions
        let effective_number_of_positions = if herfindahl_index > 0.0 {
            1.0 / herfindahl_index
        } else {
            0.0
        };
        
        // Top 5 concentration
        let top_5_value: f64 = position_values.iter()
            .take(5)
            .map(|(_, value)| value)
            .sum();
        let top_5_concentration = top_5_value / total_value;
        
        // Sector concentration (simplified - would need sector mapping)
        let sector_concentration = HashMap::new(); // TODO: Implement with sector data
        
        ConcentrationMetrics {
            herfindahl_index,
            effective_number_of_positions,
            top_5_concentration,
            sector_concentration,
        }
    }

    fn calculate_total_value(
        &self,
        positions: &HashMap<String, Position>,
        current_prices: &HashMap<String, f64>,
    ) -> f64 {
        positions.iter()
            .filter_map(|(symbol, position)| {
                current_prices.get(symbol).map(|&price| (position.quantity * price).abs())
            })
            .sum()
    }

    fn run_stress_tests(
        &self,
        positions: &HashMap<String, Position>,
        current_prices: &HashMap<String, f64>,
        cov_matrix: &DMatrix<f64>,
    ) -> Result<HashMap<String, f64>, String> {
        let mut results = HashMap::new();
        
        // Market crash scenario
        let market_crash_loss = self.calculate_scenario_loss(
            positions,
            current_prices,
            -0.20, // 20% market drop
        );
        results.insert("market_crash_20pct".to_string(), market_crash_loss);
        
        // Flight to quality
        let flight_to_quality = self.calculate_correlation_stress(
            positions,
            current_prices,
            cov_matrix,
            0.9, // High correlation scenario
        )?;
        results.insert("flight_to_quality".to_string(), flight_to_quality);
        
        // Volatility spike
        let vol_spike = self.calculate_volatility_stress(
            positions,
            current_prices,
            cov_matrix,
            2.0, // Double volatility
        )?;
        results.insert("volatility_spike_2x".to_string(), vol_spike);
        
        Ok(results)
    }

    fn calculate_scenario_loss(
        &self,
        positions: &HashMap<String, Position>,
        current_prices: &HashMap<String, f64>,
        shock: f64,
    ) -> f64 {
        let current_value = self.calculate_total_value(positions, current_prices);
        current_value * shock.abs()
    }

    fn calculate_correlation_stress(
        &self,
        positions: &HashMap<String, Position>,
        current_prices: &HashMap<String, f64>,
        cov_matrix: &DMatrix<f64>,
        target_correlation: f64,
    ) -> Result<f64, String> {
        // Adjust correlation matrix to stress scenario
        let n = cov_matrix.nrows();
        let mut stressed_cov = cov_matrix.clone();
        
        for i in 0..n {
            for j in 0..n {
                if i != j {
                    let current_corr = cov_matrix[(i, j)] / (cov_matrix[(i, i)].sqrt() * cov_matrix[(j, j)].sqrt());
                    let stress_factor = target_correlation / current_corr.abs().max(0.1);
                    stressed_cov[(i, j)] *= stress_factor;
                }
            }
        }
        
        let (_, weights) = self.get_portfolio_weights(positions, current_prices)?;
        let stressed_variance = self.calculate_portfolio_variance(&weights, &stressed_cov);
        let stressed_vol = stressed_variance.sqrt();
        
        Ok(self.calculate_total_value(positions, current_prices) * stressed_vol * 2.326) // 99% VaR
    }

    fn calculate_volatility_stress(
        &self,
        positions: &HashMap<String, Position>,
        current_prices: &HashMap<String, f64>,
        cov_matrix: &DMatrix<f64>,
        vol_multiplier: f64,
    ) -> Result<f64, String> {
        let stressed_cov = cov_matrix * vol_multiplier.powi(2);
        let (_, weights) = self.get_portfolio_weights(positions, current_prices)?;
        let stressed_variance = self.calculate_portfolio_variance(&weights, &stressed_cov);
        let stressed_vol = stressed_variance.sqrt();
        
        Ok(self.calculate_total_value(positions, current_prices) * stressed_vol * 2.326) // 99% VaR
    }
}