//! ML-powered strategy execution engine for backtesting

use anyhow::{Context, Result};
use chrono::{DateTime, Utc, Timelike};
use std::collections::HashMap;
use std::sync::Arc;
use tracing::{debug, error, info, warn};
use serde::{Deserialize, Serialize};

use config::schemas::BacktestingStrategyConfig;
use crate::storage::StorageManager;
use crate::strategy_engine::{MarketData, BacktestTrade, TradeSide, TradeSignal, StrategyExecutor, Portfolio};

/// ML model prediction result for backtesting
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MLPrediction {
    /// Model identifier
    pub model_id: String,
    /// Prediction value (0.0-1.0)
    pub prediction_value: f64,
    /// Confidence score (0.0-1.0)
    pub confidence: f64,
    /// Features used for prediction
    pub features: Vec<f64>,
    /// Prediction timestamp
    pub timestamp: DateTime<Utc>,
    /// Inference latency in microseconds
    pub inference_latency_us: u64,
}

/// ML model performance tracking for backtesting
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct MLModelPerformance {
    /// Model identifier
    pub model_id: String,
    /// Total predictions made
    pub total_predictions: u64,
    /// Correct predictions (when outcome is known)
    pub correct_predictions: u64,
    /// Average inference latency
    pub avg_latency_us: f64,
    /// Average confidence score
    pub avg_confidence: f64,
    /// Model accuracy percentage
    pub accuracy_percentage: f64,
    /// Returns generated when following this model
    pub returns: Vec<f64>,
    /// Sharpe ratio for this model
    pub sharpe_ratio: f64,
    /// Maximum drawdown when following this model
    pub max_drawdown: f64,
}

/// ML feature extractor for market data
#[derive(Debug)]
pub struct MLFeatureExtractor {
    /// Lookback window for features
    pub lookback_periods: usize,
    /// Price history buffer
    price_history: Vec<f64>,
    /// Volume history buffer
    volume_history: Vec<f64>,
}

impl MLFeatureExtractor {
    /// Create new feature extractor
    pub fn new(lookback_periods: usize) -> Self {
        Self {
            lookback_periods,
            price_history: Vec::with_capacity(lookback_periods + 1),
            volume_history: Vec::with_capacity(lookback_periods + 1),
        }
    }

    /// Extract features from market data
    pub fn extract_features(&mut self, market_data: &MarketData) -> Vec<f64> {
        // Update price and volume history
        self.price_history.push(market_data.close.to_f64().unwrap_or(0.0));
        self.volume_history.push(market_data.volume.to_f64().unwrap_or(0.0));

        // Keep only the required lookback periods
        if self.price_history.len() > self.lookback_periods {
            self.price_history.remove(0);
        }
        if self.volume_history.len() > self.lookback_periods {
            self.volume_history.remove(0);
        }

        // Extract technical features
        let mut features = Vec::new();

        if self.price_history.len() >= 2 {
            // Price momentum (returns)
            let current_price = self.price_history.last().copied().unwrap_or(0.0);
            let prev_price = self.price_history.get(self.price_history.len() - 2).copied().unwrap_or(current_price);
            let price_return = if prev_price != 0.0 {
                (current_price - prev_price) / prev_price
            } else {
                0.0
            };
            features.push(price_return);

            // Short-term moving average
            if self.price_history.len() >= 5 {
                let short_ma: f64 = self.price_history.iter().rev().take(5).sum::<f64>() / 5.0;
                let ma_ratio = if short_ma != 0.0 { current_price / short_ma - 1.0 } else { 0.0 };
                features.push(ma_ratio);
            } else {
                features.push(0.0);
            }

            // Price volatility (rolling standard deviation)
            if self.price_history.len() >= 10 {
                let recent_returns: Vec<f64> = self.price_history
                    .windows(2)
                    .rev()
                    .take(9)
                    .map(|w| (w[1] - w[0]) / w[0])
                    .collect();
                
                let mean_return = recent_returns.iter().sum::<f64>() / recent_returns.len() as f64;
                let variance = recent_returns.iter()
                    .map(|&r| (r - mean_return).powi(2))
                    .sum::<f64>() / recent_returns.len() as f64;
                let volatility = variance.sqrt();
                features.push(volatility);
            } else {
                features.push(0.0);
            }
        } else {
            features.extend_from_slice(&[0.0, 0.0, 0.0]);
        }

        // Volume features
        if self.volume_history.len() >= 2 {
            let current_volume = self.volume_history.last().copied().unwrap_or(0.0);
            let prev_volume = self.volume_history.get(self.volume_history.len() - 2).copied().unwrap_or(current_volume);
            let volume_ratio = if prev_volume != 0.0 {
                current_volume / prev_volume - 1.0
            } else {
                0.0
            };
            features.push(volume_ratio);

            // Volume moving average
            if self.volume_history.len() >= 5 {
                let volume_ma = self.volume_history.iter().rev().take(5).sum::<f64>() / 5.0;
                let volume_ma_ratio = if volume_ma != 0.0 { current_volume / volume_ma - 1.0 } else { 0.0 };
                features.push(volume_ma_ratio);
            } else {
                features.push(0.0);
            }
        } else {
            features.extend_from_slice(&[0.0, 0.0]);
        }

        // Add time-based features
        let hour = market_data.timestamp.hour() as f64 / 24.0; // Normalized hour
        let day_of_week = market_data.timestamp.weekday().num_days_from_monday() as f64 / 6.0; // Normalized day
        features.push(hour);
        features.push(day_of_week);

        // Normalize all features to [-1, 1] range using tanh
        features.iter().map(|&f| f.tanh()).collect()
    }
}

/// ML-powered strategy for backtesting
pub struct MLPoweredStrategy {
    /// Strategy name
    name: String,
    /// Available ML models
    models: HashMap<String, Box<dyn MLModelSimulator>>,
    /// Feature extractor
    feature_extractor: MLFeatureExtractor,
    /// Model performance tracking
    model_performance: HashMap<String, MLModelPerformance>,
    /// Current position size based on confidence
    confidence_based_sizing: bool,
    /// Minimum confidence threshold for trades
    min_confidence_threshold: f64,
}

/// Trait for ML model simulation in backtesting
pub trait MLModelSimulator: Send + Sync {
    /// Get model prediction
    fn predict(&self, features: &[f64]) -> Result<MLPrediction>;
    
    /// Get model identifier
    fn model_id(&self) -> &str;
    
    /// Validate prediction against actual outcome
    fn validate_prediction(&mut self, prediction: &MLPrediction, actual_outcome: bool);
}

/// Simple DQN model simulator
pub struct DQNModelSimulator {
    model_id: String,
    weights: Vec<f64>,
    predictions_made: u64,
    correct_predictions: u64,
}

impl DQNModelSimulator {
    pub fn new(model_id: String) -> Self {
        // Initialize with random weights for simulation
        let weights = vec![0.1, -0.05, 0.2, 0.15, -0.1, 0.08, 0.03];
        
        Self {
            model_id,
            weights,
            predictions_made: 0,
            correct_predictions: 0,
        }
    }
}

impl MLModelSimulator for DQNModelSimulator {
    fn predict(&self, features: &[f64]) -> Result<MLPrediction> {
        if features.len() != self.weights.len() {
            return Err(anyhow::anyhow!("Feature dimension mismatch: expected {}, got {}", 
                self.weights.len(), features.len()));
        }

        // Simple linear combination with sigmoid activation
        let linear_output: f64 = features.iter()
            .zip(self.weights.iter())
            .map(|(f, w)| f * w)
            .sum();
        
        let prediction_value = 1.0 / (1.0 + (-linear_output).exp()); // Sigmoid activation
        
        // Calculate confidence based on distance from 0.5
        let confidence = 0.5 + (prediction_value - 0.5).abs() * 0.8;
        
        Ok(MLPrediction {
            model_id: self.model_id.clone(),
            prediction_value,
            confidence,
            features: features.to_vec(),
            timestamp: Utc::now(),
            inference_latency_us: 50, // Simulated latency
        })
    }
    
    fn model_id(&self) -> &str {
        &self.model_id
    }
    
    fn validate_prediction(&mut self, _prediction: &MLPrediction, actual_outcome: bool) {
        self.predictions_made += 1;
        
        // Simple validation: if prediction > 0.5 and outcome is positive, it's correct
        let predicted_positive = _prediction.prediction_value > 0.5;
        if predicted_positive == actual_outcome {
            self.correct_predictions += 1;
        }
    }
}

/// Transformer model simulator
pub struct TransformerModelSimulator {
    model_id: String,
    attention_weights: Vec<Vec<f64>>,
    predictions_made: u64,
    correct_predictions: u64,
}

impl TransformerModelSimulator {
    pub fn new(model_id: String) -> Self {
        // Initialize with simulated attention weights
        let attention_weights = vec![
            vec![0.3, 0.2, 0.1, 0.05, 0.02, 0.01, 0.01], // Attention to recent features
            vec![0.1, 0.15, 0.2, 0.15, 0.1, 0.05, 0.05], // Attention to trend features  
        ];
        
        Self {
            model_id,
            attention_weights,
            predictions_made: 0,
            correct_predictions: 0,
        }
    }
}

impl MLModelSimulator for TransformerModelSimulator {
    fn predict(&self, features: &[f64]) -> Result<MLPrediction> {
        if features.len() != self.attention_weights[0].len() {
            return Err(anyhow::anyhow!("Feature dimension mismatch: expected {}, got {}", 
                self.attention_weights[0].len(), features.len()));
        }

        // Simulate transformer attention mechanism
        let mut attended_features = Vec::new();
        
        for attention_head in &self.attention_weights {
            let attended_value: f64 = features.iter()
                .zip(attention_head.iter())
                .map(|(f, w)| f * w)
                .sum();
            attended_features.push(attended_value);
        }
        
        // Final prediction layer
        let prediction_value = attended_features.iter().sum::<f64>().tanh() * 0.5 + 0.5;
        let confidence = 0.6 + attended_features.iter().map(|x| x.abs()).sum::<f64>() * 0.2;
        
        Ok(MLPrediction {
            model_id: self.model_id.clone(),
            prediction_value: prediction_value.clamp(0.0, 1.0),
            confidence: confidence.clamp(0.0, 1.0),
            features: features.to_vec(),
            timestamp: Utc::now(),
            inference_latency_us: 75, // Transformer models typically slower
        })
    }
    
    fn model_id(&self) -> &str {
        &self.model_id
    }
    
    fn validate_prediction(&mut self, _prediction: &MLPrediction, actual_outcome: bool) {
        self.predictions_made += 1;
        
        let predicted_positive = _prediction.prediction_value > 0.5;
        if predicted_positive == actual_outcome {
            self.correct_predictions += 1;
        }
    }
}

impl MLPoweredStrategy {
    /// Create new ML-powered strategy
    pub fn new(name: String, lookback_periods: usize) -> Self {
        let mut models: HashMap<String, Box<dyn MLModelSimulator>> = HashMap::new();
        
        // Add DQN model
        models.insert("dqn_v1".to_string(), Box::new(DQNModelSimulator::new("dqn_v1".to_string())));
        
        // Add Transformer model
        models.insert("transformer_v1".to_string(), Box::new(TransformerModelSimulator::new("transformer_v1".to_string())));
        
        Self {
            name,
            models,
            feature_extractor: MLFeatureExtractor::new(lookback_periods),
            model_performance: HashMap::new(),
            confidence_based_sizing: true,
            min_confidence_threshold: 0.6,
        }
    }

    /// Get ensemble prediction from all models
    pub fn get_ensemble_prediction(&mut self, market_data: &MarketData) -> Result<Vec<MLPrediction>> {
        // Extract features
        let features = self.feature_extractor.extract_features(market_data);
        
        let mut predictions = Vec::new();
        
        // Get predictions from all models
        for (model_id, model) in &self.models {
            match model.predict(&features) {
                Ok(prediction) => {
                    debug!("Model {} prediction: {:.3} (confidence: {:.3})", 
                        model_id, prediction.prediction_value, prediction.confidence);
                    predictions.push(prediction);
                }
                Err(e) => {
                    warn!("Model {} failed to predict: {}", model_id, e);
                }
            }
        }
        
        Ok(predictions)
    }

    /// Calculate weighted ensemble prediction
    pub fn calculate_ensemble_vote(&self, predictions: &[MLPrediction]) -> Option<(f64, f64)> {
        if predictions.is_empty() {
            return None;
        }

        let total_confidence: f64 = predictions.iter().map(|p| p.confidence).sum();
        if total_confidence == 0.0 {
            return None;
        }

        // Weighted average by confidence
        let weighted_prediction: f64 = predictions.iter()
            .map(|p| p.prediction_value * p.confidence)
            .sum::<f64>() / total_confidence;

        let average_confidence: f64 = predictions.iter().map(|p| p.confidence).sum::<f64>() / predictions.len() as f64;

        Some((weighted_prediction, average_confidence))
    }

    /// Validate predictions against actual market outcomes
    pub fn validate_predictions(&mut self, predictions: &[MLPrediction], actual_return: f64) {
        let actual_outcome = actual_return > 0.0; // Positive return = good outcome
        
        for prediction in predictions {
            if let Some(model) = self.models.get_mut(&prediction.model_id) {
                model.validate_prediction(prediction, actual_outcome);
            }
            
            // Update performance tracking
            let performance = self.model_performance.entry(prediction.model_id.clone())
                .or_insert_with(|| MLModelPerformance {
                    model_id: prediction.model_id.clone(),
                    ..Default::default()
                });
            
            performance.total_predictions += 1;
            
            let predicted_positive = prediction.prediction_value > 0.5;
            if predicted_positive == actual_outcome {
                performance.correct_predictions += 1;
            }
            
            performance.accuracy_percentage = if performance.total_predictions > 0 {
                (performance.correct_predictions as f64 / performance.total_predictions as f64) * 100.0
            } else {
                0.0
            };
            
            // Update average confidence
            let total_samples = performance.total_predictions as f64;
            performance.avg_confidence = (performance.avg_confidence * (total_samples - 1.0) + prediction.confidence) / total_samples;
            
            // Update average latency
            performance.avg_latency_us = (performance.avg_latency_us * (total_samples - 1.0) + prediction.inference_latency_us as f64) / total_samples;
        }
    }

    /// Get performance summary for all models
    pub fn get_performance_summary(&self) -> HashMap<String, MLModelPerformance> {
        self.model_performance.clone()
    }
}

impl StrategyExecutor for MLPoweredStrategy {
    fn execute(
        &self,
        market_data: &MarketData,
        _portfolio: &Portfolio,
        parameters: &HashMap<String, String>,
    ) -> Result<Vec<TradeSignal>> {
        // This is a bit tricky because we need mutable access to call predict
        // In a real implementation, you'd want to redesign this to avoid the issue
        // For now, we'll create a simplified version that doesn't update the feature extractor
        
        let mut signals = Vec::new();
        
        // Extract basic features without updating history (simplified for demo)
        let price = market_data.close.to_f64().unwrap_or(0.0);
        let volume = market_data.volume.to_f64().unwrap_or(0.0);
        
        // Create simplified features
        let features = vec![
            (price - 100.0) / 100.0, // Normalized price change from baseline
            (volume - 1000.0) / 1000.0, // Normalized volume
            0.0, 0.0, 0.0, 0.0, 0.0 // Placeholder features
        ];
        
        // Simple prediction using DQN-like logic
        let weights = vec![0.1, -0.05, 0.2, 0.15, -0.1, 0.08, 0.03];
        let linear_output: f64 = features.iter()
            .zip(weights.iter())
            .map(|(f, w)| f * w)
            .sum();
        
        let prediction_value = 1.0 / (1.0 + (-linear_output).exp());
        let confidence = 0.5 + (prediction_value - 0.5).abs() * 0.8;
        
        // Get minimum confidence from parameters
        let min_confidence = parameters.get("min_confidence")
            .and_then(|s| s.parse::<f64>().ok())
            .unwrap_or(self.min_confidence_threshold);
        
        // Generate signal if confidence is high enough
        if confidence >= min_confidence {
            let quantity = if self.confidence_based_sizing {
                // Size position based on confidence
                Decimal::try_from(confidence * 1000.0).unwrap_or(Decimal::from(100))
            } else {
                Decimal::from(100)
            };
            
            if prediction_value > 0.6 {
                signals.push(TradeSignal {
                    symbol: market_data.symbol.clone(),
                    side: TradeSide::Buy,
                    quantity,
                    strength: Decimal::try_from(confidence).unwrap_or(Decimal::try_from(0.5).unwrap()),
                    reason: format!("ML prediction: {:.3} (confidence: {:.3})", prediction_value, confidence),
                });
            } else if prediction_value < 0.4 {
                signals.push(TradeSignal {
                    symbol: market_data.symbol.clone(),
                    side: TradeSide::Sell,
                    quantity,
                    strength: Decimal::try_from(confidence).unwrap_or(Decimal::try_from(0.5).unwrap()),
                    reason: format!("ML prediction: {:.3} (confidence: {:.3})", prediction_value, confidence),
                });
            }
        }
        
        Ok(signals)
    }

    fn name(&self) -> &str {
        &self.name
    }
}

/// ML Strategy Engine with model performance tracking
pub struct MLStrategyEngine {
    /// Base strategy engine
    base_engine: crate::strategy_engine::StrategyEngine,
    /// ML-powered strategies
    ml_strategies: HashMap<String, MLPoweredStrategy>,
    /// Model performance tracking across backtests
    global_model_performance: HashMap<String, MLModelPerformance>,
}

impl MLStrategyEngine {
    /// Create new ML strategy engine
    pub async fn new(
        config: &BacktestingStrategyConfig,
        storage_manager: Arc<StorageManager>,
    ) -> Result<Self> {
        let base_engine = crate::strategy_engine::StrategyEngine::new(config, storage_manager).await?;
        
        let mut ml_strategies = HashMap::new();
        
        // Add ML-powered strategies
        ml_strategies.insert(
            "ml_momentum".to_string(),
            MLPoweredStrategy::new("ml_momentum".to_string(), 20)
        );
        
        ml_strategies.insert(
            "ml_ensemble".to_string(), 
            MLPoweredStrategy::new("ml_ensemble".to_string(), 50)
        );
        
        Ok(Self {
            base_engine,
            ml_strategies,
            global_model_performance: HashMap::new(),
        })
    }

    /// Execute backtest with ML model validation
    pub async fn execute_ml_backtest(
        &mut self,
        context: &crate::service::BacktestContext,
    ) -> Result<(Vec<BacktestTrade>, HashMap<String, MLModelPerformance>)> {
        info!("Executing ML-powered backtest {} for strategy {}", context.id, context.strategy_name);

        // Check if this is an ML strategy
        if let Some(ml_strategy) = self.ml_strategies.get_mut(&context.strategy_name) {
            // Execute ML-powered backtest with model validation
            self.execute_ml_strategy_backtest(ml_strategy, context).await
        } else {
            // Fall back to base strategy engine
            let trades = self.base_engine.execute_backtest(context).await?;
            Ok((trades, HashMap::new()))
        }
    }

    /// Execute backtest for ML strategy with model performance tracking
    async fn execute_ml_strategy_backtest(
        &mut self,
        ml_strategy: &mut MLPoweredStrategy,
        context: &crate::service::BacktestContext,
    ) -> Result<(Vec<BacktestTrade>, HashMap<String, MLModelPerformance>)> {
        // Load market data for the backtest period
        let market_data = self.base_engine
            .load_market_data(
                &context.symbols,
                context.started_at,
                context.completed_at.unwrap_or(chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0)),
            )
            .await?;

        let mut trades = Vec::new();
        let mut previous_price = None;

        // Process each data point with ML predictions
        for (i, data_point) in market_data.iter().enumerate() {
            // Get ML predictions
            let predictions = ml_strategy.get_ensemble_prediction(data_point)?;
            
            // Calculate ensemble vote
            if let Some((ensemble_prediction, ensemble_confidence)) = ml_strategy.calculate_ensemble_vote(&predictions) {
                debug!("Ensemble prediction: {:.3} (confidence: {:.3})", ensemble_prediction, ensemble_confidence);
                
                // Validate predictions against future returns if we have next price
                if let Some(prev_price) = previous_price {
                    let actual_return = (data_point.close.to_f64().unwrap_or(0.0) - prev_price) / prev_price;
                    ml_strategy.validate_predictions(&predictions, actual_return);
                }
            }
            
            previous_price = Some(data_point.close.to_f64().unwrap_or(0.0));
            
            // Generate and execute trades using base strategy logic
            // (This would integrate with the existing strategy execution logic)
            if i % 100 == 0 {
                let progress = (i as f64 / market_data.len() as f64) * 100.0;
                debug!("ML backtest progress: {:.1}%", progress);
            }
        }

        // Get final model performance
        let model_performance = ml_strategy.get_performance_summary();
        
        // Update global performance tracking
        for (model_id, perf) in &model_performance {
            self.global_model_performance.insert(model_id.clone(), perf.clone());
        }

        info!("ML backtest completed with {} trades and {} model evaluations", 
            trades.len(), model_performance.len());
        
        Ok((trades, model_performance))
    }

    /// Get model performance across all backtests
    pub fn get_global_model_performance(&self) -> &HashMap<String, MLModelPerformance> {
        &self.global_model_performance
    }

    /// Generate model performance report
    pub fn generate_performance_report(&self) -> String {
        let mut report = String::new();
        report.push_str("=== ML Model Performance Report ===\n\n");

        for (model_id, performance) in &self.global_model_performance {
            report.push_str(&format!("Model: {}\n", model_id));
            report.push_str(&format!("  Total Predictions: {}\n", performance.total_predictions));
            report.push_str(&format!("  Accuracy: {:.2}%\n", performance.accuracy_percentage));
            report.push_str(&format!("  Average Confidence: {:.3}\n", performance.avg_confidence));
            report.push_str(&format!("  Average Latency: {:.1}μs\n", performance.avg_latency_us));
            if performance.sharpe_ratio != 0.0 {
                report.push_str(&format!("  Sharpe Ratio: {:.3}\n", performance.sharpe_ratio));
            }
            if performance.max_drawdown != 0.0 {
                report.push_str(&format!("  Max Drawdown: {:.2}%\n", performance.max_drawdown * 100.0));
            }
            report.push_str("\n");
        }

        report
    }
}