//! Neural Value-at-Risk Models for HFT Risk Management
//!
//! Implements advanced neural network architectures for VaR estimation,
//! Expected Shortfall calculation, and stress testing with canonical types.

use chrono::{DateTime, Utc};
use ndarray::{Array1, Array2};
use serde::{Deserialize, Serialize};

// Import types from crate root (which imports from common)
use crate::{MLError, MLResult as Result};
use common::types::{Price, Quantity, Symbol};

// AssetId type for VaR models
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AssetId(String);

/// Market tick data for VaR calculations
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MarketTick {
    pub symbol: Symbol,
    pub price: Price,
    pub quantity: Quantity,
    pub timestamp: DateTime<Utc>,
}

/// VaR prediction result using canonical types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VarPrediction {
    pub asset_id: AssetId,
    pub var_estimates: Vec<Price>,
    pub expected_shortfall: Vec<Price>,
    pub volatility_forecast: Price,
    pub model_confidence: Price,
    pub stress_test_results: Option<StressTestResults>,
}

/// Stress test results using canonical types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StressTestResults {
    pub stress_var: Price,
    pub stress_es: Price,
    pub scenario_name: String,
}

/// Neural VaR model configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NeuralVarConfig {
    pub confidence_levels: Vec<f64>,
    pub lookback_period: usize,
    pub lookback_days: usize,
    pub monte_carlo_simulations: usize,
    pub enable_stress_testing: bool,
    pub hidden_layers: Vec<usize>,
}

impl Default for NeuralVarConfig {
    fn default() -> Self {
        Self {
            confidence_levels: vec![0.95, 0.99, 0.999],
            lookback_period: 252,
            lookback_days: 252,
            monte_carlo_simulations: 10000,
            enable_stress_testing: true,
            hidden_layers: vec![128, 64, 32],
        }
    }
}

/// Neural VaR model
#[derive(Debug)]
pub struct NeuralVarModel {
    pub config: NeuralVarConfig,
    weights: Vec<Array2<f64>>,
    biases: Vec<Array1<f64>>,
}

impl NeuralVarModel {
    pub fn new(config: NeuralVarConfig) -> Result<Self> {
        let mut weights = Vec::new();
        let mut biases = Vec::new();

        // Initialize neural network layers
        let mut prev_size = 100; // Input features size
        for &hidden_size in &config.hidden_layers {
            weights.push(Array2::from_elem((hidden_size, prev_size), 0.1));
            biases.push(Array1::from_elem(hidden_size, 0.0));
            prev_size = hidden_size;
        }

        // Output layer for VaR and ES estimates
        let output_size = config.confidence_levels.len() * 2;
        weights.push(Array2::from_elem((output_size, prev_size), 0.1));
        biases.push(Array1::from_elem(output_size, 0.0));

        Ok(Self {
            config,
            weights,
            biases,
        })
    }

    pub async fn predict_var(
        &mut self,
        asset_id: AssetId,
        _market_data: &[MarketTick],
    ) -> Result<VarPrediction> {
        // Simple VaR calculation for now - production would use full neural network
        let mut var_estimates = Vec::new();
        let mut expected_shortfall = Vec::new();

        for confidence in &self.config.confidence_levels {
            // Production calculations - production would use trained model
            let var_value =
                Price::from_f64(*confidence * 0.01).map_err(|e| MLError::ValidationError {
                    message: format!("Invalid VaR price: {}", e),
                })?;
            let es_value =
                Price::from_f64(*confidence * 0.012).map_err(|e| MLError::ValidationError {
                    message: format!("Invalid ES price: {}", e),
                })?;

            var_estimates.push(var_value);
            expected_shortfall.push(es_value);
        }

        let stress_test_results = self
            .config
            .enable_stress_testing
            .then(|| {
                Ok::<_, MLError>(StressTestResults {
                    stress_var: Price::from_f64(0.05).map_err(|e| MLError::ValidationError {
                        message: format!("Invalid stress VaR: {}", e),
                    })?,
                    stress_es: Price::from_f64(0.08).map_err(|e| MLError::ValidationError {
                        message: format!("Invalid stress ES: {}", e),
                    })?,
                    scenario_name: "Market Crash".to_owned(),
                })
            })
            .transpose()?;

        Ok(VarPrediction {
            asset_id,
            var_estimates,
            expected_shortfall,
            volatility_forecast: Price::from_f64(0.02).map_err(|e| MLError::ValidationError {
                message: format!("Invalid volatility forecast: {}", e),
            })?,
            model_confidence: Price::from_f64(0.95).map_err(|e| MLError::ValidationError {
                message: format!("Invalid model confidence: {}", e),
            })?,
            stress_test_results,
        })
    }
}

/// VaR features extracted from market data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VarFeatures {
    pub returns: Vec<f64>,
    pub volatility: f64,
    pub volume: f64,
    pub timestamp: DateTime<Utc>,
}

impl VarFeatures {
    pub fn from_market_data(market_data: &[MarketTick], lookback_period: usize) -> Result<Self> {
        if market_data.is_empty() {
            return Err(MLError::InvalidInput("Empty market data".to_owned()));
        }

        let mut returns = Vec::new();
        let data_len = market_data.len().min(lookback_period);

        // Calculate returns from price data
        for i in 1..data_len {
            let prev_price = market_data
                .get(i - 1)
                .ok_or_else(|| MLError::ValidationError {
                    message: format!("Index {} out of bounds", i - 1),
                })?
                .price
                .to_f64();
            let curr_price = market_data
                .get(i)
                .ok_or_else(|| MLError::ValidationError {
                    message: format!("Index {} out of bounds", i),
                })?
                .price
                .to_f64();
            let return_val = (curr_price - prev_price) / prev_price;
            returns.push(return_val);
        }

        // Calculate rolling volatility
        let mean_return = returns.iter().sum::<f64>() / returns.len() as f64;
        let variance = returns
            .iter()
            .map(|r| (r - mean_return).powi(2))
            .sum::<f64>()
            / returns.len() as f64;
        let volatility = variance.sqrt();

        // Calculate average volume
        let volume = market_data
            .iter()
            .take(data_len)
            .map(|tick| tick.quantity.to_f64())
            .sum::<f64>()
            / data_len as f64;

        Ok(Self {
            returns,
            volatility,
            volume,
            timestamp: market_data
                .last()
                .ok_or_else(|| MLError::InvalidInput("No market data provided".to_owned()))?
                .timestamp,
        })
    }

    pub fn to_feature_vector(&self) -> Array1<f64> {
        let mut features = Vec::new();

        // Add statistical features
        features.push(self.volatility);
        features.push(self.volume);

        // Add recent returns (up to 10)
        let recent_returns = self
            .returns
            .iter()
            .rev()
            .take(10)
            .cloned()
            .collect::<Vec<_>>();
        features.extend(recent_returns);

        // Pad with zeros if needed
        while features.len() < 100 {
            features.push(0.0);
        }

        Array1::from_vec(features)
    }
}

/// Linear layer for neural network
#[derive(Debug)]
pub struct LinearLayer {
    weights: Array2<f64>,
    bias: Array1<f64>,
}

impl LinearLayer {
    pub fn new(input_size: usize, output_size: usize) -> Result<Self> {
        Ok(Self {
            weights: Array2::from_elem((output_size, input_size), 0.1),
            bias: Array1::from_elem(output_size, 0.0),
        })
    }

    pub fn forward(&self, input: &Array1<f64>) -> Result<Array1<f64>> {
        let output = self.weights.dot(input) + &self.bias;
        Ok(output)
    }
}

/// Feature scaler for normalization
#[derive(Debug)]
pub struct FeatureScaler {
    mean: Option<Array1<f64>>,
    std: Option<Array1<f64>>,
}

impl FeatureScaler {
    pub fn new() -> Self {
        Self {
            mean: None,
            std: None,
        }
    }

    pub fn fit(&mut self, data: &Array2<f64>) -> Result<()> {
        let mean = data
            .mean_axis(ndarray::Axis(0))
            .ok_or_else(|| MLError::InvalidInput("Cannot compute mean".to_owned()))?;

        let std = data.std_axis(ndarray::Axis(0), 0.0);

        self.mean = Some(mean);
        self.std = Some(std);

        Ok(())
    }

    pub fn transform(&self, data: &Array1<f64>) -> Result<Array1<f64>> {
        match (&self.mean, &self.std) {
            (Some(mean), Some(std)) => {
                let normalized = (data - mean) / std;
                Ok(normalized)
            },
            _ => Err(MLError::InvalidInput("Scaler not fitted".to_owned())),
        }
    }
}

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


    #[test]
    fn test_neural_var_model_creation() -> Result<()> {
        let config = NeuralVarConfig::default();
        let model = NeuralVarModel::new(config);
        assert!(model.is_ok());
        Ok(())
    }

    #[test]
    fn test_var_features_from_market_data() -> Result<()> {
        let mut market_data = Vec::new();
        let symbol = Symbol::from("AAPL");

        for i in 0..10 {
            market_data.push(MarketTick {
                symbol: symbol.clone(),
                price: Price::from_f64(100.0 + i as f64).unwrap(),
                quantity: Quantity::from_f64(1000.0).unwrap(),
                timestamp: Utc::now(),
            });
        }

        let features = VarFeatures::from_market_data(&market_data, 252);
        assert!(features.is_ok());

        let features = features?;
        assert!(!features.returns.is_empty());
        assert!(features.volatility > 0.0);
        assert!(features.volume > 0.0);
        Ok(())
    }

    #[test]
    fn test_linear_layer() -> Result<()> {
        let layer = LinearLayer::new(10, 5)?;
        let input = Array1::from_elem(10, 1.0);
        let output = layer.forward(&input);

        assert!(output.is_ok());
        let output = output?;
        assert_eq!(output.len(), 5);
        Ok(())
    }

    #[test]
    fn test_feature_scaler() {
        let mut scaler = FeatureScaler::new();
        let data = Array2::from_elem((100, 10), 1.0);

        let fit_result = scaler.fit(&data);
        assert!(fit_result.is_ok());

        let input = Array1::from_elem(10, 1.0);
        let transformed = scaler.transform(&input);
        assert!(transformed.is_ok());
    }
}