foxhunt/adaptive-strategy/src/ensemble/mod.rs

//! Enhanced ensemble coordinator for managing multiple ML models
//!
//! This module provides an advanced coordination layer for ensemble model management,
//! including sophisticated dynamic weight optimization, confidence-weighted voting,
//! uncertainty quantification, and performance-based adaptation.

// Import core types
use common::Order;
use common::Position;
use common::Symbol;
use common::Price;
use common::Quantity;
use common::error::CommonError;
use common::error::CommonResult;
use common::HftTimestamp;
use common::OrderId;
use common::TradeId;
use common::Execution;

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

// Add missing core types
use super::config::{EnsembleConfig, ModelConfig, StrategyConfig};
use super::models::{ModelPrediction, ModelTrait};

pub mod confidence_aggregator;
pub mod weight_optimizer;

use confidence_aggregator::{
    AggregationConfig, ConfidenceAggregator, EnsemblePredictionWithUncertainty, ReliabilityRecord,
};
use weight_optimizer::{OptimizedWeights, PerformanceRecord, WeightOptimizer};

/// Enhanced ensemble coordinator managing multiple ML models
///
/// The coordinator provides sophisticated capabilities including:
/// - Advanced dynamic weight optimization with multiple algorithms
/// - Confidence-weighted voting with uncertainty quantification
/// - Performance-based model selection and adaptation
/// - Real-time regime-aware weight adjustment
/// - Model health monitoring and reliability scoring
#[derive(Debug)]
pub struct EnsembleCoordinator {
    /// Configuration for the ensemble
    config: EnsembleConfig,
    /// Active models in the ensemble
    models: HashMap<String, Arc<dyn ModelTrait>>,
    /// Advanced weight optimizer with multiple algorithms
    weight_optimizer: Arc<RwLock<WeightOptimizer>>,
    /// Confidence aggregator for uncertainty-aware voting
    confidence_aggregator: Arc<RwLock<ConfidenceAggregator>>,
    /// Current optimized weights
    current_weights: Arc<RwLock<OptimizedWeights>>,
    /// Model performance tracking (legacy - migrating to weight_optimizer)
    performance_tracker: Arc<RwLock<PerformanceTracker>>,
    /// Prediction history for analysis (legacy - migrating to confidence_aggregator)
    prediction_history: Arc<RwLock<PredictionHistory>>,
}

/// Tracks performance metrics for each model
#[derive(Debug, Clone)]
pub struct PerformanceTracker {
    /// Accuracy metrics per model
    accuracy: HashMap<String, f64>,
    /// Precision metrics per model
    precision: HashMap<String, f64>,
    /// Recall metrics per model
    recall: HashMap<String, f64>,
    /// Sharpe ratio per model
    sharpe_ratio: HashMap<String, f64>,
    /// Recent prediction count per model
    prediction_counts: HashMap<String, u64>,
    /// Last update timestamp
    last_update: chrono::DateTime<chrono::Utc>,
}

/// Stores prediction history for analysis and optimization
#[derive(Debug, Clone)]
pub struct PredictionHistory {
    /// Historical predictions by model
    predictions: HashMap<String, Vec<HistoricalPrediction>>,
    /// Maximum history length to maintain
    max_history_length: usize,
}

/// Historical prediction record
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HistoricalPrediction {
    /// Timestamp of prediction
    pub timestamp: chrono::DateTime<chrono::Utc>,
    /// Model prediction
    pub prediction: ModelPrediction,
    /// Actual outcome (if available)
    pub actual_outcome: Option<f64>,
    /// Model confidence
    pub confidence: f64,
}

/// Aggregated ensemble prediction
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnsemblePrediction {
    /// Weighted average prediction
    pub prediction: f64,
    /// Ensemble confidence
    pub confidence: f64,
    /// Contributing models and their weights
    pub model_contributions: HashMap<String, ModelContribution>,
    /// Prediction timestamp
    pub timestamp: chrono::DateTime<chrono::Utc>,
    /// Prediction horizon
    pub horizon: chrono::Duration,
}

/// Individual model contribution to ensemble prediction (legacy format)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelContribution {
    /// Model's raw prediction
    pub prediction: f64,
    /// Model's confidence
    pub confidence: f64,
    /// Model's weight in ensemble
    pub weight: f64,
    /// Model's weighted contribution
    pub weighted_contribution: f64,
}

impl EnsembleCoordinator {
    /// Create a new enhanced ensemble coordinator
    ///
    /// # Arguments
    ///
    /// * `strategy_config` - Complete strategy configuration including ensemble settings
    ///
    /// # Returns
    ///
    /// A new enhanced `EnsembleCoordinator` instance with advanced capabilities
    pub async fn new(strategy_config: &StrategyConfig) -> Result<Self> {
        info!(
            "Initializing enhanced ensemble coordinator with {} models",
            strategy_config.ensemble.models.len()
        );

        let config = strategy_config.ensemble.clone();
        let mut models = HashMap::new();
        let mut initial_weights = HashMap::new();

        // Initialize models
        for model_config in &config.models {
            if model_config.enabled {
                info!("Initializing model: {}", model_config.name);

                // Create model instance (production - would use actual model factory)
                let model = create_model_instance(model_config).await?;
                models.insert(model_config.name.clone(), model);
                initial_weights.insert(model_config.name.clone(), model_config.initial_weight);
            }
        }

        // Initialize advanced weight optimizer
        let weight_optimizer = Arc::new(RwLock::new(WeightOptimizer::new(
            Duration::from_secs(3600), // 1 hour performance window
            0.01,                      // Adaptation rate
        )));

        // Initialize confidence aggregator
        let confidence_aggregator = Arc::new(RwLock::new(ConfidenceAggregator::new(
            AggregationConfig::default(),
        )));

        // Create initial optimized weights
        let model_names: Vec<String> = initial_weights.keys().cloned().collect();
        let current_weights = Arc::new(RwLock::new(OptimizedWeights {
            weights: initial_weights,
            algorithm_used: weight_optimizer::WeightingAlgorithmType::BayesianModelAveraging,
            confidence: 0.5,
            expected_variance: 0.1,
            timestamp: chrono::Utc::now(),
        }));

        // Legacy components (for backward compatibility)
        let performance_tracker = Arc::new(RwLock::new(PerformanceTracker::new()));
        let prediction_history = Arc::new(RwLock::new(PredictionHistory::new(1000)));

        Ok(Self {
            config,
            models,
            weight_optimizer,
            confidence_aggregator,
            current_weights,
            performance_tracker,
            prediction_history,
        })
    }

    /// Generate enhanced ensemble prediction with uncertainty quantification
    ///
    /// # Arguments
    ///
    /// * `features` - Input features for prediction
    /// * `horizon` - Prediction time horizon
    /// * `market_regime` - Current market regime (optional)
    ///
    /// # Returns
    ///
    /// Enhanced ensemble prediction with uncertainty bounds and confidence intervals
    pub async fn predict_with_uncertainty(
        &self,
        features: &[f64],
        horizon: chrono::Duration,
        market_regime: Option<&str>,
    ) -> Result<EnsemblePredictionWithUncertainty> {
        debug!(
            "Generating enhanced ensemble prediction for {} features",
            features.len()
        );

        let mut model_predictions = HashMap::new();

        // Collect predictions from all models
        for (model_name, model) in &self.models {
            match model.predict(features).await {
                Ok(prediction) => {
                    model_predictions.insert(model_name.clone(), prediction);
                }
                Err(e) => {
                    warn!("Model {} failed to predict: {}", model_name, e);
                    // Continue with other models
                }
            }
        }

        if model_predictions.is_empty() {
            anyhow::bail!("No models provided valid predictions");
        }

        // Get optimized weights
        let model_names: Vec<String> = model_predictions.keys().cloned().collect();
        let optimized_weights = {
            let mut optimizer = self.weight_optimizer.write().await;
            optimizer
                .optimize_weights(&model_names, market_regime)
                .await?
        };

        // Update current weights
        {
            let mut current = self.current_weights.write().await;
            *current = optimized_weights.clone();
        }

        // Aggregate with uncertainty quantification
        let ensemble_prediction = {
            let mut aggregator = self.confidence_aggregator.write().await;
            aggregator
                .aggregate_with_uncertainty(model_predictions, &optimized_weights.weights)
                .await?
        };

        // Store prediction in history (legacy support)
        self.store_enhanced_prediction_history(&ensemble_prediction, horizon)
            .await?;

        info!(
            "Generated ensemble prediction: {:.4} (confidence: {:.3}, uncertainty: {:.3})",
            ensemble_prediction.prediction,
            ensemble_prediction.confidence,
            ensemble_prediction.uncertainty_decomposition.total
        );

        Ok(ensemble_prediction)
    }

    /// Generate basic ensemble prediction (backward compatibility)
    ///
    /// # Arguments
    ///
    /// * `features` - Input features for prediction
    /// * `horizon` - Prediction time horizon
    ///
    /// # Returns
    ///
    /// Basic ensemble prediction (converted from enhanced prediction)
    pub async fn predict(
        &self,
        features: &[f64],
        horizon: chrono::Duration,
    ) -> Result<EnsemblePrediction> {
        let enhanced_prediction = self
            .predict_with_uncertainty(features, horizon, None)
            .await?;

        // Convert to legacy format
        Ok(EnsemblePrediction {
            prediction: enhanced_prediction.prediction,
            confidence: enhanced_prediction.confidence,
            model_contributions: enhanced_prediction
                .model_contributions
                .into_iter()
                .map(|(name, contrib)| {
                    (
                        name,
                        ModelContribution {
                            prediction: contrib.prediction,
                            confidence: contrib.confidence,
                            weight: contrib.weight,
                            weighted_contribution: contrib.weighted_contribution,
                        },
                    )
                })
                .collect(),
            timestamp: enhanced_prediction.timestamp,
            horizon,
        })
    }

    /// Update model weights using advanced optimization algorithms
    pub async fn update_weights(&self, market_regime: Option<&str>) -> Result<()> {
        info!("Updating ensemble model weights with advanced optimization");

        let model_names: Vec<String> = self.models.keys().cloned().collect();

        // Use advanced weight optimizer
        let optimized_weights = {
            let mut optimizer = self.weight_optimizer.write().await;
            optimizer
                .optimize_weights(&model_names, market_regime)
                .await?
        };

        // Update current weights
        {
            let mut current = self.current_weights.write().await;
            *current = optimized_weights.clone();
        }

        info!(
            "Advanced model weights updated successfully using {:?} algorithm",
            optimized_weights.algorithm_used
        );
        info!(
            "Weight confidence: {:.3}, Expected variance: {:.4}",
            optimized_weights.confidence, optimized_weights.expected_variance
        );

        Ok(())
    }

    /// Update model weights with default parameters
    pub async fn update_weights_default(&self) -> Result<()> {
        self.update_weights(None).await
    }

    /// Record actual outcome for enhanced performance tracking
    ///
    /// # Arguments
    ///
    /// * `prediction_timestamp` - Timestamp of the prediction
    /// * `actual_outcome` - The realized outcome
    /// * `model_predictions` - Individual model predictions for reliability tracking
    pub async fn record_outcome(
        &self,
        prediction_timestamp: chrono::DateTime<chrono::Utc>,
        actual_outcome: f64,
        model_predictions: Option<HashMap<String, f64>>,
    ) -> Result<()> {
        debug!(
            "Recording enhanced outcome: {} at {}",
            actual_outcome, prediction_timestamp
        );

        // Update performance history for weight optimizer
        if let Some(ref predictions) = model_predictions {
            let mut optimizer = self.weight_optimizer.write().await;

            for (model_name, predicted_value) in predictions {
                let accuracy =
                    1.0 - (predicted_value - actual_outcome).abs() / actual_outcome.abs().max(1e-6);
                let accuracy = accuracy.max(0.0).min(1.0);

                let performance_record = PerformanceRecord {
                    timestamp: prediction_timestamp,
                    accuracy,
                    sharpe_ratio: 0.0, // Would calculate based on returns history
                    max_drawdown: 0.0, // Would track from returns
                    volatility: 0.0,   // Would calculate from returns
                    return_value: (predicted_value - actual_outcome)
                        / actual_outcome.abs().max(1e-6),
                    confidence: 0.8, // Would get from model prediction
                    regime: None,    // Would get from regime detector
                };

                optimizer.update_performance(model_name.clone(), performance_record);
            }
        }

        // Update reliability for confidence aggregator
        if let Some(ref predictions) = model_predictions {
            let mut aggregator = self.confidence_aggregator.write().await;

            for (model_name, predicted_value) in predictions {
                let accuracy =
                    1.0 - (predicted_value - actual_outcome).abs() / actual_outcome.abs().max(1e-6);
                let accuracy = accuracy.max(0.0).min(1.0);

                let reliability_record = ReliabilityRecord {
                    timestamp: prediction_timestamp,
                    accuracy,
                    calibration: accuracy, // Simplified - would use proper calibration metrics
                    confidence_reliability: accuracy, // Simplified
                    actual_outcome: Some(actual_outcome),
                };

                aggregator
                    .update_reliability(model_name.clone(), reliability_record)
                    .await?;
            }
        }

        // Legacy support
        {
            let mut history = self.prediction_history.write().await;
            history.update_outcome(prediction_timestamp, actual_outcome)?;
        }

        // Update performance metrics
        self.update_performance_metrics().await?;

        Ok(())
    }

    /// Record actual outcome with basic parameters
    pub async fn record_outcome_legacy(
        &self,
        prediction_timestamp: chrono::DateTime<chrono::Utc>,
        actual_outcome: f64,
    ) -> Result<()> {
        self.record_outcome(prediction_timestamp, actual_outcome, None)
            .await
    }

    /// Get current optimized model weights
    pub async fn get_weights(&self) -> HashMap<String, f64> {
        self.current_weights.read().await.weights.clone()
    }

    /// Get detailed weight information including algorithm and confidence
    pub async fn get_detailed_weights(&self) -> OptimizedWeights {
        self.current_weights.read().await.clone()
    }

    /// Get performance metrics for all models
    pub async fn get_performance(&self) -> PerformanceTracker {
        self.performance_tracker.read().await.clone()
    }

    /// Aggregate predictions from multiple models
    async fn aggregate_predictions(
        &self,
        model_predictions: HashMap<String, ModelPrediction>,
        weights: &HashMap<String, f64>,
        horizon: chrono::Duration,
    ) -> Result<EnsemblePrediction> {
        let mut weighted_sum = 0.0;
        let mut total_weight = 0.0;
        let mut weighted_confidence = 0.0;
        let mut model_contributions = HashMap::new();

        for (model_name, prediction) in &model_predictions {
            if let Some(&weight) = weights.get(model_name) {
                let weighted_prediction = prediction.value * weight;
                let weighted_conf = prediction.confidence * weight;

                weighted_sum += weighted_prediction;
                weighted_confidence += weighted_conf;
                total_weight += weight;

                model_contributions.insert(
                    model_name.clone(),
                    ModelContribution {
                        prediction: prediction.value,
                        confidence: prediction.confidence,
                        weight,
                        weighted_contribution: weighted_prediction,
                    },
                );
            }
        }

        if total_weight == 0.0 {
            anyhow::bail!("Total weight is zero - cannot aggregate predictions");
        }

        let ensemble_prediction = weighted_sum / total_weight;
        let ensemble_confidence = weighted_confidence / total_weight;

        Ok(EnsemblePrediction {
            prediction: ensemble_prediction,
            confidence: ensemble_confidence,
            model_contributions,
            timestamp: chrono::Utc::now(),
            horizon,
        })
    }

    /// Store enhanced prediction in history for analysis
    async fn store_enhanced_prediction_history(
        &self,
        ensemble_prediction: &EnsemblePredictionWithUncertainty,
        horizon: chrono::Duration,
    ) -> Result<()> {
        // Store prediction history
        let mut history = self.prediction_history.write().await;

        for (model_name, contribution) in &ensemble_prediction.model_contributions {
            let historical_prediction = HistoricalPrediction {
                timestamp: ensemble_prediction.timestamp,
                prediction: ModelPrediction {
                    value: contribution.prediction,
                    confidence: contribution.confidence,
                    features_used: vec![], // Would store actual features in production
                    metadata: None,
                },
                actual_outcome: None,
                confidence: contribution.confidence,
            };

            history.add_prediction(model_name.clone(), historical_prediction);
        }

        Ok(())
    }

    /// Store prediction in history for later analysis (legacy)
    async fn store_prediction_history(
        &self,
        ensemble_prediction: &EnsemblePrediction,
    ) -> Result<()> {
        let mut history = self.prediction_history.write().await;

        for (model_name, contribution) in &ensemble_prediction.model_contributions {
            let historical_prediction = HistoricalPrediction {
                timestamp: ensemble_prediction.timestamp,
                prediction: ModelPrediction {
                    value: contribution.prediction,
                    confidence: contribution.confidence,
                    features_used: vec![], // Would store actual features in production
                    metadata: None,
                },
                actual_outcome: None,
                confidence: contribution.confidence,
            };

            history.add_prediction(model_name.clone(), historical_prediction);
        }

        Ok(())
    }

    /// Update performance metrics based on historical predictions
    async fn update_performance_metrics(&self) -> Result<()> {
        let history = self.prediction_history.read().await;
        let mut performance = self.performance_tracker.write().await;

        for (model_name, predictions) in &history.predictions {
            let recent_predictions: Vec<_> = predictions
                .iter()
                .filter(|p| p.actual_outcome.is_some())
                .collect();

            if !recent_predictions.is_empty() {
                let accuracy = self.calculate_accuracy(&recent_predictions);
                let sharpe = self.calculate_sharpe_ratio(&recent_predictions);

                performance.accuracy.insert(model_name.clone(), accuracy);
                performance.sharpe_ratio.insert(model_name.clone(), sharpe);
                performance
                    .prediction_counts
                    .insert(model_name.clone(), recent_predictions.len() as u64);
            }
        }

        performance.last_update = chrono::Utc::now();
        Ok(())
    }

    /// Calculate accuracy for a set of predictions
    fn calculate_accuracy(&self, predictions: &[&HistoricalPrediction]) -> f64 {
        if predictions.is_empty() {
            return 0.0;
        }

        let correct_predictions = predictions
            .iter()
            .filter(|p| {
                if let Some(actual) = p.actual_outcome {
                    // Simple accuracy: correct direction prediction
                    (p.prediction.value > 0.0 && actual > 0.0)
                        || (p.prediction.value < 0.0 && actual < 0.0)
                } else {
                    false
                }
            })
            .count();

        correct_predictions as f64 / predictions.len() as f64
    }

    /// Calculate Sharpe ratio for a set of predictions
    fn calculate_sharpe_ratio(&self, predictions: &[&HistoricalPrediction]) -> f64 {
        if predictions.len() < 2 {
            return 0.0;
        }

        let returns: Vec<f64> = predictions
            .iter()
            .filter_map(|p| p.actual_outcome)
            .collect();

        if returns.is_empty() {
            return 0.0;
        }

        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;

        if variance == 0.0 {
            return 0.0;
        }

        mean_return / variance.sqrt()
    }
}

impl PerformanceTracker {
    /// Create a new performance tracker
    pub fn new() -> Self {
        Self {
            accuracy: HashMap::new(),
            precision: HashMap::new(),
            recall: HashMap::new(),
            sharpe_ratio: HashMap::new(),
            prediction_counts: HashMap::new(),
            last_update: chrono::Utc::now(),
        }
    }
}

impl PredictionHistory {
    /// Create a new prediction history
    pub fn new(max_length: usize) -> Self {
        Self {
            predictions: HashMap::new(),
            max_history_length: max_length,
        }
    }

    /// Add a prediction to the history
    pub fn add_prediction(&mut self, model_name: String, prediction: HistoricalPrediction) {
        let predictions = self.predictions.entry(model_name).or_insert_with(Vec::new);
        predictions.push(prediction);

        // Maintain maximum history length
        if predictions.len() > self.max_history_length {
            predictions.remove(0);
        }
    }

    /// Update a prediction with actual outcome
    pub fn update_outcome(
        &mut self,
        timestamp: chrono::DateTime<chrono::Utc>,
        actual_outcome: f64,
    ) -> Result<()> {
        let mut updated = false;

        for predictions in self.predictions.values_mut() {
            for prediction in predictions.iter_mut() {
                if (prediction.timestamp - timestamp).num_seconds().abs() < 60
                    && prediction.actual_outcome.is_none()
                {
                    prediction.actual_outcome = Some(actual_outcome);
                    updated = true;
                }
            }
        }

        if !updated {
            warn!(
                "Could not find prediction to update with timestamp: {}",
                timestamp
            );
        }

        Ok(())
    }
}

/// Factory function to create model instances
/// This would be implemented with actual model constructors in production
async fn create_model_instance(config: &ModelConfig) -> Result<Arc<dyn ModelTrait>> {
    // Production implementation - would create actual models based on config.model_type
    use crate::models::MockModel;

    info!("Creating mock model instance for: {}", config.name);
    Ok(Arc::new(MockModel::new(config.name.clone())))
}

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

    #[tokio::test]
    async fn test_ensemble_coordinator_creation() {
        let config = StrategyConfig::default();
        let coordinator = EnsembleCoordinator::new(&config).await;
        assert!(coordinator.is_ok());
    }

    #[tokio::test]
    async fn test_performance_tracker() {
        let tracker = PerformanceTracker::new();
        assert!(tracker.accuracy.is_empty());
        assert!(tracker.sharpe_ratio.is_empty());
    }

    #[tokio::test]
    async fn test_prediction_history() {
        let mut history = PredictionHistory::new(10);

        let prediction = HistoricalPrediction {
            timestamp: chrono::Utc::now(),
            prediction: ModelPrediction {
                value: 0.5,
                confidence: 0.8,
                features_used: vec![],
            },
            actual_outcome: None,
            confidence: 0.8,
        };

        history.add_prediction("test_model".to_string(), prediction);
        assert_eq!(history.predictions.get("test_model").unwrap().len(), 1);
    }
}