//! ML Pipeline Testing Harness
//!
//! Provides comprehensive testing capabilities for ML models including
//! inference testing, feature extraction, ensemble predictions, and
//! model performance monitoring.

use anyhow::Result;
use common::types::MarketTick;
use std::collections::HashMap;
use std::time::{Duration, Instant};
use tracing::{debug, info, warn};

/// ML model status information
#[derive(Debug, Clone)]
pub struct MLModelStatus {
    pub mamba_available: bool,
    pub dqn_available: bool,
    pub ppo_available: bool,
    pub tft_available: bool,
    pub tlob_available: bool,
    pub ensemble_available: bool,
}

impl MLModelStatus {
    /// Check if any models are available
    pub fn any_available(&self) -> bool {
        self.mamba_available
            || self.dqn_available
            || self.ppo_available
            || self.tft_available
            || self.tlob_available
    }

    /// Get count of available models
    pub fn available_count(&self) -> usize {
        let mut count = 0;
        if self.mamba_available {
            count += 1;
        }
        if self.dqn_available {
            count += 1;
        }
        if self.ppo_available {
            count += 1;
        }
        if self.tft_available {
            count += 1;
        }
        if self.tlob_available {
            count += 1;
        }
        count
    }
}

/// ML prediction result
#[derive(Debug, Clone)]
pub struct MLPrediction {
    pub signal: f64,              // Trading signal (-1.0 to 1.0)
    pub confidence: f64,          // Confidence (0.0 to 1.0)
    pub model_name: String,       // Name of the model that made the prediction
    pub inference_time: Duration, // Time taken for inference
}

/// Ensemble prediction result
#[derive(Debug, Clone)]
pub struct EnsemblePrediction {
    pub signal: f64,     // Aggregated trading signal
    pub confidence: f64, // Aggregated confidence
    pub individual_predictions: Vec<MLPrediction>,
    pub ensemble_method: String, // Method used for aggregation
    pub total_inference_time: Duration,
    pub prediction: PredictionType, // Discrete prediction type
    pub signal_strength: f64,       // Absolute signal strength
}

// Define PredictionType locally since tli crate is not available in e2e tests
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PredictionType {
    Buy,
    Sell,
    Hold,
    StrongBuy,
    StrongSell,
}

impl std::fmt::Display for PredictionType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            PredictionType::Buy => write!(f, "BUY"),
            PredictionType::Sell => write!(f, "SELL"),
            PredictionType::Hold => write!(f, "HOLD"),
            PredictionType::StrongBuy => write!(f, "STRONG_BUY"),
            PredictionType::StrongSell => write!(f, "STRONG_SELL"),
        }
    }
}

/// Model performance metrics
#[derive(Debug, Clone)]
pub struct ModelMetrics {
    pub inference_count: u64,
    pub avg_latency_ms: f64,
    pub error_rate: f64,
    pub last_prediction_time: chrono::DateTime<chrono::Utc>,
}

/// Feature vector for ML models
#[derive(Debug, Clone)]
pub struct FeatureVector {
    pub features: Vec<f64>,
    pub feature_names: Vec<String>,
    pub timestamp: chrono::DateTime<chrono::Utc>,
    pub symbol: String,
}

/// ML Pipeline Testing Harness
#[derive(Debug)]
pub struct MLPipelineTestHarness {
    model_status: MLModelStatus,
    model_metrics: HashMap<String, ModelMetrics>,
    feature_cache: HashMap<String, Vec<FeatureVector>>,
    mock_mode: bool,
}

impl MLPipelineTestHarness {
    /// Create a new ML pipeline test harness
    pub async fn new() -> Result<Self> {
        info!("🤖 Initializing ML pipeline test harness");

        // Check if we're in mock mode (for CI/testing environments without GPU)
        let mock_mode = std::env::var("ML_MOCK_MODE").unwrap_or_default() == "true";

        if mock_mode {
            info!("🎭 Running in mock mode - ML predictions will be simulated");
        }

        let model_status = Self::check_model_availability().await?;

        info!("ML Models Status:");
        info!(
            "  MAMBA: {}",
            if model_status.mamba_available {
                "✅"
            } else {
                "❌"
            }
        );
        info!(
            "  DQN: {}",
            if model_status.dqn_available {
                "✅"
            } else {
                "❌"
            }
        );
        info!(
            "  PPO: {}",
            if model_status.ppo_available {
                "✅"
            } else {
                "❌"
            }
        );
        info!(
            "  TFT: {}",
            if model_status.tft_available {
                "✅"
            } else {
                "❌"
            }
        );
        info!(
            "  TLOB: {}",
            if model_status.tlob_available {
                "✅"
            } else {
                "❌"
            }
        );

        Ok(Self {
            model_status,
            model_metrics: HashMap::new(),
            feature_cache: HashMap::new(),
            mock_mode,
        })
    }

    /// Check model health status
    pub async fn check_models_health(&self) -> Result<MLModelStatus> {
        Ok(self.model_status.clone())
    }

    /// Extract features from market data
    pub async fn extract_features(
        &mut self,
        market_data: &[MarketTick],
    ) -> Result<Vec<FeatureVector>> {
        debug!(
            "🔧 Extracting features from {} market ticks",
            market_data.len()
        );

        if market_data.is_empty() {
            return Ok(Vec::new());
        }

        let mut features = Vec::new();

        // Group by symbol for feature extraction
        let mut symbol_data: HashMap<String, Vec<&MarketTick>> = HashMap::new();
        for tick in market_data {
            symbol_data
                .entry(tick.symbol.to_string())
                .or_default()
                .push(tick);
        }

        for (symbol, ticks) in symbol_data {
            let feature_vector = self.extract_features_for_symbol(&symbol, ticks).await?;
            features.push(feature_vector);
        }

        // Cache features for later use
        for feature in &features {
            self.feature_cache
                .entry(feature.symbol.clone())
                .or_default()
                .push(feature.clone());
        }

        debug!("✅ Extracted {} feature vectors", features.len());
        Ok(features)
    }

    /// Extract features for a specific symbol
    async fn extract_features_for_symbol(
        &self,
        symbol: &str,
        ticks: Vec<&MarketTick>,
    ) -> Result<FeatureVector> {
        if ticks.is_empty() {
            return Err(anyhow::anyhow!("No ticks provided for feature extraction"));
        }

        // Sort by timestamp
        let mut sorted_ticks = ticks;
        sorted_ticks.sort_by_key(|t| t.timestamp);

        let prices: Vec<f64> = sorted_ticks.iter().map(|t| t.price.as_f64()).collect();
        let volumes: Vec<f64> = sorted_ticks.iter().map(|t| t.size.as_f64()).collect();

        // Calculate technical indicators
        let mut features = Vec::new();
        let mut feature_names = Vec::new();

        // Price-based features
        if !prices.is_empty() {
            let current_price = prices[prices.len() - 1];
            let avg_price = prices.iter().sum::<f64>() / prices.len() as f64;
            let price_std = Self::calculate_std(&prices);

            features.push(current_price);
            features.push(avg_price);
            features.push(price_std);
            features.push(current_price / avg_price - 1.0); // Price relative to average

            feature_names.extend([
                "current_price".to_string(),
                "avg_price".to_string(),
                "price_std".to_string(),
                "price_rel_avg".to_string(),
            ]);
        }

        // Volume-based features
        if !volumes.is_empty() {
            let current_volume = volumes[volumes.len() - 1];
            let avg_volume = volumes.iter().sum::<f64>() / volumes.len() as f64;

            features.push(current_volume);
            features.push(avg_volume);
            features.push(current_volume / avg_volume.max(1.0)); // Volume relative to average

            feature_names.extend([
                "current_volume".to_string(),
                "avg_volume".to_string(),
                "volume_rel_avg".to_string(),
            ]);
        }

        // Returns-based features
        if prices.len() >= 2 {
            let returns: Vec<f64> = prices.windows(2).map(|w| (w[1] / w[0]) - 1.0).collect();

            if !returns.is_empty() {
                let last_return = returns[returns.len() - 1];
                let avg_return = returns.iter().sum::<f64>() / returns.len() as f64;
                let return_std = Self::calculate_std(&returns);

                features.push(last_return);
                features.push(avg_return);
                features.push(return_std);

                feature_names.extend([
                    "last_return".to_string(),
                    "avg_return".to_string(),
                    "return_std".to_string(),
                ]);
            }
        }

        // Trend features
        if prices.len() >= 5 {
            let short_ma = prices[prices.len() - 5..].iter().sum::<f64>() / 5.0;
            let long_ma = prices.iter().sum::<f64>() / prices.len() as f64;
            let trend_strength = (short_ma / long_ma) - 1.0;

            features.push(trend_strength);
            feature_names.push("trend_strength".to_string());
        }

        Ok(FeatureVector {
            features,
            feature_names,
            timestamp: chrono::Utc::now(),
            symbol: symbol.to_string(),
        })
    }

    /// Predict with MAMBA model
    pub async fn predict_with_mamba(&mut self, features: &[FeatureVector]) -> Result<MLPrediction> {
        self.predict_with_model("mamba", features).await
    }

    /// Predict with DQN model
    pub async fn predict_with_dqn(&mut self, features: &[FeatureVector]) -> Result<MLPrediction> {
        self.predict_with_model("dqn", features).await
    }

    /// Predict with TFT model
    pub async fn predict_with_tft(&mut self, features: &[FeatureVector]) -> Result<MLPrediction> {
        self.predict_with_model("tft", features).await
    }

    /// Predict with TLOB model
    pub async fn predict_with_tlob(&mut self, features: &[FeatureVector]) -> Result<MLPrediction> {
        self.predict_with_model("tlob", features).await
    }

    /// Generic model prediction
    async fn predict_with_model(
        &mut self,
        model_name: &str,
        features: &[FeatureVector],
    ) -> Result<MLPrediction> {
        let start_time = Instant::now();

        if features.is_empty() {
            return Err(anyhow::anyhow!("No features provided for prediction"));
        }

        let prediction = if self.mock_mode {
            self.mock_prediction(model_name, features).await?
        } else {
            self.real_prediction(model_name, features).await?
        };

        let inference_time = start_time.elapsed();

        // Update model metrics
        self.update_model_metrics(model_name, inference_time, true);

        Ok(MLPrediction {
            signal: prediction.0,
            confidence: prediction.1,
            model_name: model_name.to_string(),
            inference_time,
        })
    }

    /// Mock prediction for testing
    async fn mock_prediction(
        &self,
        model_name: &str,
        features: &[FeatureVector],
    ) -> Result<(f64, f64)> {
        use rand::Rng;
        let mut rng = rand::thread_rng();

        // Simulate some processing time
        tokio::time::sleep(Duration::from_millis(rng.gen_range(10..50))).await;

        // Generate reasonable mock predictions based on features
        let feature_sum: f64 = features.iter().flat_map(|f| &f.features).sum();

        let normalized_sum = (feature_sum / 1000.0).tanh(); // Normalize to [-1, 1]

        let signal = match model_name {
            "mamba" => normalized_sum * 0.8 + rng.gen_range(-0.1..0.1),
            "dqn" => normalized_sum * 0.6 + rng.gen_range(-0.2..0.2),
            "tft" => normalized_sum * 0.9 + rng.gen_range(-0.05..0.05),
            "tlob" => normalized_sum * 0.7 + rng.gen_range(-0.15..0.15),
            _ => rng.gen_range(-0.5..0.5),
        };

        let confidence = rng.gen_range(0.6..0.95);

        Ok((signal.clamp(-1.0, 1.0), confidence))
    }

    /// Real prediction (would integrate with actual ML models)
    async fn real_prediction(
        &self,
        model_name: &str,
        _features: &[FeatureVector],
    ) -> Result<(f64, f64)> {
        // This would integrate with the actual ML models in the ml crate
        warn!(
            "Real ML prediction not implemented for {}, using mock",
            model_name
        );

        // For now, fall back to mock prediction
        self.mock_prediction(model_name, _features).await
    }

    /// Test ensemble prediction with raw feature array
    pub async fn test_ensemble_prediction(
        &mut self,
        raw_features: Vec<f64>,
    ) -> Result<EnsemblePrediction> {
        // Convert raw features to FeatureVector format
        let feature_count = raw_features.len();
        let feature_vector = FeatureVector {
            features: raw_features,
            feature_names: (0..feature_count)
                .map(|i| format!("feature_{}", i))
                .collect(),
            timestamp: chrono::Utc::now(),
            symbol: "TEST".to_string(),
        };

        self.predict_ensemble(&[feature_vector]).await
    }

    /// Ensemble prediction aggregating multiple models
    pub async fn predict_ensemble(
        &mut self,
        features: &[FeatureVector],
    ) -> Result<EnsemblePrediction> {
        let start_time = Instant::now();
        let mut predictions = Vec::new();

        // Get predictions from available models
        if self.model_status.mamba_available {
            match self.predict_with_mamba(features).await {
                Ok(pred) => predictions.push(pred),
                Err(e) => warn!("MAMBA prediction failed: {}", e),
            }
        }

        if self.model_status.dqn_available {
            match self.predict_with_dqn(features).await {
                Ok(pred) => predictions.push(pred),
                Err(e) => warn!("DQN prediction failed: {}", e),
            }
        }

        if self.model_status.tft_available {
            match self.predict_with_tft(features).await {
                Ok(pred) => predictions.push(pred),
                Err(e) => warn!("TFT prediction failed: {}", e),
            }
        }

        if self.model_status.tlob_available {
            match self.predict_with_tlob(features).await {
                Ok(pred) => predictions.push(pred),
                Err(e) => warn!("TLOB prediction failed: {}", e),
            }
        }

        if predictions.is_empty() {
            return Err(anyhow::anyhow!(
                "No models available for ensemble prediction"
            ));
        }

        // Aggregate predictions using weighted average
        let total_weight: f64 = predictions.iter().map(|p| p.confidence).sum();
        let weighted_signal: f64 = predictions
            .iter()
            .map(|p| p.signal * p.confidence)
            .sum::<f64>()
            / total_weight;

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

        let total_inference_time = start_time.elapsed();

        // Convert signal to discrete prediction type
        let prediction = if weighted_signal > 0.6 {
            PredictionType::StrongBuy
        } else if weighted_signal > 0.2 {
            PredictionType::Buy
        } else if weighted_signal < -0.6 {
            PredictionType::StrongSell
        } else if weighted_signal < -0.2 {
            PredictionType::Sell
        } else {
            PredictionType::Hold
        };

        let signal_strength = weighted_signal.abs();

        Ok(EnsemblePrediction {
            signal: weighted_signal.clamp(-1.0, 1.0),
            confidence: avg_confidence.clamp(0.0, 1.0),
            individual_predictions: predictions,
            ensemble_method: "weighted_average".to_string(),
            total_inference_time,
            prediction,
            signal_strength,
        })
    }

    /// Get model performance metrics
    pub async fn get_model_metrics(&self) -> Result<HashMap<String, ModelMetrics>> {
        Ok(self.model_metrics.clone())
    }

    /// Update model metrics
    fn update_model_metrics(&mut self, model_name: &str, inference_time: Duration, success: bool) {
        let metrics = self
            .model_metrics
            .entry(model_name.to_string())
            .or_insert_with(|| ModelMetrics {
                inference_count: 0,
                avg_latency_ms: 0.0,
                error_rate: 0.0,
                last_prediction_time: chrono::Utc::now(),
            });

        metrics.inference_count += 1;

        // Update average latency
        let new_latency_ms = inference_time.as_millis() as f64;
        if metrics.inference_count == 1 {
            metrics.avg_latency_ms = new_latency_ms;
        } else {
            metrics.avg_latency_ms =
                (metrics.avg_latency_ms * (metrics.inference_count - 1) as f64 + new_latency_ms)
                    / metrics.inference_count as f64;
        }

        // Update error rate
        if !success {
            let error_count = (metrics.error_rate * (metrics.inference_count - 1) as f64) + 1.0;
            metrics.error_rate = error_count / metrics.inference_count as f64;
        } else {
            let error_count = metrics.error_rate * (metrics.inference_count - 1) as f64;
            metrics.error_rate = error_count / metrics.inference_count as f64;
        }

        metrics.last_prediction_time = chrono::Utc::now();
    }

    /// Disable a model for failover testing
    pub async fn disable_model(&mut self, model_name: &str) -> Result<()> {
        match model_name {
            "mamba" => self.model_status.mamba_available = false,
            "dqn" => self.model_status.dqn_available = false,
            "ppo" => self.model_status.ppo_available = false,
            "tft" => self.model_status.tft_available = false,
            "tlob" => self.model_status.tlob_available = false,
            _ => return Err(anyhow::anyhow!("Unknown model: {}", model_name)),
        }

        info!("🚫 Disabled model: {}", model_name);
        Ok(())
    }

    /// Enable a model for failover testing
    pub async fn enable_model(&mut self, model_name: &str) -> Result<()> {
        match model_name {
            "mamba" => self.model_status.mamba_available = true,
            "dqn" => self.model_status.dqn_available = true,
            "ppo" => self.model_status.ppo_available = true,
            "tft" => self.model_status.tft_available = true,
            "tlob" => self.model_status.tlob_available = true,
            _ => return Err(anyhow::anyhow!("Unknown model: {}", model_name)),
        }

        info!("✅ Enabled model: {}", model_name);
        Ok(())
    }

    /// Check model availability
    async fn check_model_availability() -> Result<MLModelStatus> {
        // In a real implementation, this would check for model files,
        // GPU availability, etc. For testing, we'll assume models are available.
        Ok(MLModelStatus {
            mamba_available: true,
            dqn_available: true,
            ppo_available: true,
            tft_available: true,
            tlob_available: true,
            ensemble_available: true,
        })
    }

    /// Calculate standard deviation
    fn calculate_std(values: &[f64]) -> f64 {
        if values.len() < 2 {
            return 0.0;
        }

        let mean = values.iter().sum::<f64>() / values.len() as f64;
        let variance =
            values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (values.len() - 1) as f64;

        variance.sqrt()
    }
}

/// Type alias for compatibility with workflow code
pub type MLTestPipeline = MLPipelineTestHarness;

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

    #[tokio::test]
    async fn test_ml_harness_creation() {
        let harness = MLPipelineTestHarness::new().await;
        assert!(harness.is_ok());

        let harness = harness.unwrap();
        assert!(harness.model_status.any_available());
    }

    #[test]
    fn test_model_status() {
        let status = MLModelStatus {
            mamba_available: true,
            dqn_available: true,
            ppo_available: false,
            tft_available: false,
            tlob_available: false,
            ensemble_available: true,
        };

        assert!(status.any_available());
        assert_eq!(status.available_count(), 2);
    }

    #[test]
    fn test_std_calculation() {
        let values = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let std = MLPipelineTestHarness::calculate_std(&values);
        assert!((std - 1.581).abs() < 0.01); // Approximately sqrt(2.5)
    }
}