foxhunt/ml/src/ensemble/coordinator.rs

//! Ensemble Coordinator for Production Trading
//!
//! This module implements the core ensemble coordinator that aggregates predictions
//! from multiple ML models (DQN, PPO, TFT, MAMBA-2, Liquid, TLOB) for production trading decisions.
//! Supports dynamic weighting based on performance and model diversity metrics.

use crate::ensemble::inference_adapter::{FeatureVector, ModelInferenceAdapter};
use crate::ensemble::{EnsembleDecision, ModelVote, ModelWeight, TradingAction};
use crate::{Features, MLError, MLResult, ModelPrediction};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};

/// Minimum weight threshold per model (5%)
const MIN_WEIGHT_THRESHOLD: f64 = 0.05;

/// Maximum weight per model (40%) to prevent dominance
const MAX_WEIGHT_THRESHOLD: f64 = 0.40;

/// Ensemble coordinator for aggregating model predictions
pub struct EnsembleCoordinator {
    /// Active model registry (dual-buffer for hot-swapping)
    active_models: Arc<RwLock<ModelRegistry>>,

    /// Signal aggregator
    aggregator: Arc<SignalAggregator>,

    /// Model weights configuration
    model_weights: Arc<RwLock<HashMap<String, ModelWeight>>>,

    /// Configuration for ensemble behavior
    config: EnsembleConfig,

    /// Real model inference adapters for production predictions
    adapters: Vec<Box<dyn ModelInferenceAdapter>>,
}

impl std::fmt::Debug for EnsembleCoordinator {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("EnsembleCoordinator")
            .field("config", &self.config)
            .field("adapter_count", &self.adapters.len())
            .finish()
    }
}

/// Configuration for ensemble coordinator
#[derive(Debug, Clone)]
pub struct EnsembleConfig {
    /// Enable adaptive weighting based on correlation
    pub adaptive_weighting: bool,

    /// Minimum correlation threshold for diversity bonus
    pub min_correlation_for_diversity: f64,

    /// Weight adjustment factor for diversity
    pub diversity_weight_factor: f64,

    /// Performance window size (number of predictions)
    pub performance_window_size: usize,
}

impl EnsembleCoordinator {
    /// Create new ensemble coordinator
    pub fn new() -> Self {
        let config = EnsembleConfig {
            adaptive_weighting: true,
            min_correlation_for_diversity: 0.3,
            diversity_weight_factor: 1.2,
            performance_window_size: 100,
        };

        Self {
            active_models: Arc::new(RwLock::new(ModelRegistry::new())),
            aggregator: Arc::new(SignalAggregator::new()),
            model_weights: Arc::new(RwLock::new(HashMap::new())),
            config,
            adapters: Vec::new(),
        }
    }

    /// Add a real model inference adapter to the ensemble.
    /// When an adapter's model_name matches a registered model, its predictions
    /// are used for ensemble inference.
    pub fn add_adapter(&mut self, adapter: Box<dyn ModelInferenceAdapter>) {
        info!("Added inference adapter: {}", adapter.model_name());
        self.adapters.push(adapter);
    }

    /// Register a model in the ensemble
    pub async fn register_model(&self, model_id: String, weight: f64) -> MLResult<()> {
        let model_weight = ModelWeight::new(model_id.clone(), weight);

        let mut weights = self.model_weights.write().await;
        weights.insert(model_id.clone(), model_weight);

        info!("Registered model {} with weight {}", model_id, weight);
        Ok(())
    }

    /// Make ensemble prediction from features
    ///
    /// Returns an error if no model adapters are registered or all adapters fail inference.
    /// Production trading must never rely on simulated/mock predictions.
    pub async fn predict(&self, features: &Features) -> MLResult<EnsembleDecision> {
        if self.adapters.is_empty() {
            return Err(MLError::ModelError(
                "No model adapters registered".into(),
            ));
        }

        debug!(
            "Making ensemble prediction with {} features",
            features.values.len()
        );

        let predictions = self.generate_predictions(features).await?;

        if predictions.is_empty() {
            return Err(MLError::ModelError(
                "All model adapters failed inference".into(),
            ));
        }

        // Aggregate predictions
        let decision = self
            .aggregator
            .aggregate(predictions, &*self.model_weights.read().await)
            .await?;

        info!(
            "Ensemble decision: {:?}, confidence: {:.3}, disagreement: {:.3}",
            decision.action, decision.confidence, decision.disagreement_rate
        );

        Ok(decision)
    }

    /// Generate predictions from real ML model adapters
    ///
    /// Iterates over registered models and attempts inference via their adapters.
    /// Models without a ready adapter or whose adapter fails are skipped --
    /// no mock/simulated predictions are generated on the production path.
    async fn generate_predictions(
        &self,
        features: &Features,
    ) -> MLResult<Vec<ModelPrediction>> {
        let model_ids: Vec<String> = {
            let weights = self.model_weights.read().await;
            weights.keys().cloned().collect()
        };

        let mut predictions = Vec::new();

        for model_id in model_ids {
            if let Some(adapter) = self
                .adapters
                .iter()
                .find(|a| a.model_name() == model_id && a.is_ready())
            {
                let fv = FeatureVector {
                    values: features.values.clone(),
                    timestamp: features.timestamp as i64,
                };
                match adapter.predict(&fv) {
                    Ok(ensemble_pred) => {
                        let prediction = ModelPrediction::new(
                            model_id.clone(),
                            ensemble_pred.direction,
                            ensemble_pred.confidence,
                        );
                        predictions.push(prediction);
                    }
                    Err(e) => {
                        warn!(
                            "Adapter {} inference failed, skipping: {}",
                            model_id, e
                        );
                    }
                }
            } else {
                warn!(
                    "No ready adapter for model {}, skipping",
                    model_id
                );
            }
        }

        Ok(predictions)
    }

    /// Check whether any registered adapter is ready for inference
    pub fn has_models(&self) -> bool {
        self.adapters.iter().any(|a| a.is_ready())
    }

    /// Update model weights based on performance
    pub async fn update_model_weights(&self) -> MLResult<()> {
        let mut weights = self.model_weights.write().await;

        for weight in weights.values_mut() {
            weight.update_dynamic_weight();
        }

        debug!("Updated dynamic weights for {} models", weights.len());
        Ok(())
    }

    /// Get model count
    pub async fn model_count(&self) -> usize {
        self.model_weights.read().await.len()
    }

    /// Load PPO model from production checkpoint (Agent 170 validated)
    ///
    /// Example:
    /// ```ignore
    /// coordinator.load_ppo_checkpoint(
    ///     "PPO_epoch420",
    ///     "ml/trained_models/production/ppo/ppo_actor_epoch_420.safetensors",
    ///     "ml/trained_models/production/ppo/ppo_critic_epoch_420.safetensors",
    ///     0.33,
    /// ).await?;
    /// ```
    pub async fn load_ppo_checkpoint(
        &self,
        model_id: &str,
        actor_checkpoint: &str,
        critic_checkpoint: &str,
        weight: f64,
    ) -> MLResult<()> {
        info!(
            "Loading PPO checkpoint: actor={}, critic={}",
            actor_checkpoint, critic_checkpoint
        );

        // Stage checkpoints in registry (both actor and critic as single entry)
        let mut registry = self.active_models.write().await;
        registry.stage_checkpoint(
            model_id.to_string(),
            format!("actor={},critic={}", actor_checkpoint, critic_checkpoint),
        );
        registry.commit_swap(model_id)?;
        drop(registry);

        // Register model with weight
        self.register_model(model_id.to_string(), weight).await?;

        info!(
            "✅ PPO checkpoint loaded and registered: {} (weight: {:.2})",
            model_id, weight
        );

        Ok(())
    }

    /// Load TFT-INT8 model from production checkpoint (Wave 9 INT8 quantization)
    ///
    /// INT8 quantization reduces TFT memory from 2,952MB → 738MB, enabling
    /// 4-model ensemble to fit in RTX 3050 Ti (4GB VRAM).
    ///
    /// Example:
    /// ```ignore
    /// coordinator.load_tft_int8_checkpoint(
    ///     "TFT-INT8",
    ///     "ml/trained_models/production/tft/tft_int8_epoch_200.safetensors",
    ///     0.15,
    /// ).await?;
    /// ```
    pub async fn load_tft_int8_checkpoint(
        &self,
        model_id: &str,
        checkpoint: &str,
        weight: f64,
    ) -> MLResult<()> {
        info!("Loading TFT-INT8 checkpoint: {}", checkpoint);

        // Stage checkpoint in registry
        let mut registry = self.active_models.write().await;
        registry.stage_checkpoint(model_id.to_string(), checkpoint.to_string());
        registry.commit_swap(model_id)?;
        drop(registry);

        // Register model with weight
        self.register_model(model_id.to_string(), weight).await?;

        info!(
            "✅ TFT-INT8 checkpoint loaded and registered: {} (weight: {:.2})",
            model_id, weight
        );

        Ok(())
    }
}

impl Default for EnsembleCoordinator {
    fn default() -> Self {
        Self::new()
    }
}

/// Model registry with dual-buffer support for hot-swapping
#[derive(Debug)]
pub struct ModelRegistry {
    /// Active models (currently serving predictions)
    active: HashMap<String, String>,

    /// Shadow models (staged for hot-swap)
    shadow: HashMap<String, String>,
}

impl ModelRegistry {
    /// Create new model registry
    pub fn new() -> Self {
        Self {
            active: HashMap::new(),
            shadow: HashMap::new(),
        }
    }

    /// Stage a checkpoint in shadow buffer
    pub fn stage_checkpoint(&mut self, model_id: String, checkpoint_path: String) {
        self.shadow
            .insert(model_id.clone(), checkpoint_path.clone());
        info!(
            "Staged checkpoint {} for model {}",
            checkpoint_path, model_id
        );
    }

    /// Commit swap (shadow becomes active)
    pub fn commit_swap(&mut self, model_id: &str) -> MLResult<()> {
        if let Some(shadow_path) = self.shadow.remove(model_id) {
            let old_path = self
                .active
                .insert(model_id.to_string(), shadow_path.clone());

            if let Some(old) = old_path {
                // Move old to shadow for potential rollback
                self.shadow.insert(model_id.to_string(), old);
            }

            info!("Committed checkpoint swap for model {}", model_id);
            Ok(())
        } else {
            Err(MLError::ModelNotFound(format!(
                "No staged checkpoint for model {}",
                model_id
            )))
        }
    }

    /// Rollback to previous checkpoint
    pub fn rollback(&mut self, model_id: &str) -> MLResult<()> {
        if let Some(previous_path) = self.shadow.remove(model_id) {
            self.active.insert(model_id.to_string(), previous_path);
            warn!("Rolled back model {} to previous checkpoint", model_id);
            Ok(())
        } else {
            Err(MLError::ModelNotFound(format!(
                "No previous checkpoint for model {}",
                model_id
            )))
        }
    }
}

impl Default for ModelRegistry {
    fn default() -> Self {
        Self::new()
    }
}

/// Signal aggregator for ensemble predictions
#[derive(Debug)]
pub struct SignalAggregator {
    /// Signal threshold for Buy/Sell actions
    signal_threshold: f64,

    /// Minimum confidence for high-confidence decisions
    min_confidence: f64,
}

impl SignalAggregator {
    /// Create new signal aggregator
    pub fn new() -> Self {
        Self {
            signal_threshold: 0.3,
            min_confidence: 0.6,
        }
    }

    /// Aggregate model predictions into ensemble decision
    pub async fn aggregate(
        &self,
        predictions: Vec<ModelPrediction>,
        weights: &HashMap<String, ModelWeight>,
    ) -> MLResult<EnsembleDecision> {
        if predictions.is_empty() {
            return Err(MLError::ValidationError {
                message: "No predictions to aggregate".to_string(),
            });
        }

        // Calculate weighted average signal
        let (weighted_signal, _total_weight) =
            self.calculate_weighted_signal(&predictions, weights);

        // Calculate ensemble confidence
        let confidence = self.calculate_ensemble_confidence(&predictions, weights);

        // Calculate disagreement rate
        let disagreement_rate = self.calculate_disagreement_rate(&predictions);

        // Determine trading action
        let action = TradingAction::from_signal(weighted_signal, self.signal_threshold);

        // Build model votes
        let model_votes = self.build_model_votes(&predictions, weights);

        let decision = EnsembleDecision::new(
            action,
            confidence,
            weighted_signal,
            disagreement_rate,
            model_votes,
        );

        Ok(decision)
    }

    /// Calculate weighted average signal
    fn calculate_weighted_signal(
        &self,
        predictions: &[ModelPrediction],
        weights: &HashMap<String, ModelWeight>,
    ) -> (f64, f64) {
        let mut weighted_sum = 0.0;
        let mut total_weight = 0.0;

        for pred in predictions {
            let weight = weights
                .get(&pred.model_id)
                .map(|w| w.effective_weight())
                .unwrap_or(1.0 / predictions.len() as f64);

            weighted_sum += pred.value * pred.confidence * weight;
            total_weight += weight * pred.confidence;
        }

        let signal = if total_weight > 0.0 {
            weighted_sum / total_weight
        } else {
            0.0
        };

        (signal, total_weight)
    }

    /// Calculate ensemble confidence
    fn calculate_ensemble_confidence(
        &self,
        predictions: &[ModelPrediction],
        weights: &HashMap<String, ModelWeight>,
    ) -> f64 {
        let mut confidence_sum = 0.0;
        let mut weight_sum = 0.0;

        for pred in predictions {
            let weight = weights
                .get(&pred.model_id)
                .map(|w| w.effective_weight())
                .unwrap_or(1.0 / predictions.len() as f64);

            confidence_sum += pred.confidence * weight;
            weight_sum += weight;
        }

        if weight_sum > 0.0 {
            confidence_sum / weight_sum
        } else {
            0.0
        }
    }

    /// Calculate disagreement rate (% models disagree with ensemble)
    fn calculate_disagreement_rate(&self, predictions: &[ModelPrediction]) -> f64 {
        if predictions.len() < 2 {
            return 0.0;
        }

        // Calculate mean signal
        let mean_signal: f64 =
            predictions.iter().map(|p| p.value).sum::<f64>() / predictions.len() as f64;

        // Count models with opposite sign from mean
        let disagreements = predictions
            .iter()
            .filter(|p| (p.value * mean_signal) < 0.0)
            .count();

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

    /// Build model votes map
    fn build_model_votes(
        &self,
        predictions: &[ModelPrediction],
        weights: &HashMap<String, ModelWeight>,
    ) -> HashMap<String, ModelVote> {
        let mut votes = HashMap::new();

        for pred in predictions {
            let weight = weights
                .get(&pred.model_id)
                .map(|w| w.effective_weight())
                .unwrap_or(1.0 / predictions.len() as f64);

            let vote = ModelVote::new(pred.model_id.clone(), pred.value, pred.confidence, weight);

            votes.insert(pred.model_id.clone(), vote);
        }

        votes
    }
}

impl Default for SignalAggregator {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[tokio::test]
    async fn test_ensemble_coordinator_creation() {
        let coordinator = EnsembleCoordinator::new();
        assert_eq!(coordinator.model_count().await, 0);
    }

    #[tokio::test]
    async fn test_register_models() {
        let coordinator = EnsembleCoordinator::new();

        coordinator
            .register_model("DQN".to_string(), 0.33)
            .await
            .unwrap();
        coordinator
            .register_model("PPO".to_string(), 0.33)
            .await
            .unwrap();
        coordinator
            .register_model("TFT".to_string(), 0.34)
            .await
            .unwrap();

        assert_eq!(coordinator.model_count().await, 3);
    }

    #[tokio::test]
    async fn test_ensemble_prediction() {
        use crate::ensemble::inference_adapter::{
            EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
        };

        struct StubAdapter {
            name: &'static str,
            direction: f64,
            confidence: f64,
        }
        impl ModelInferenceAdapter for StubAdapter {
            fn model_name(&self) -> &str {
                self.name
            }
            fn predict(
                &self,
                _features: &FeatureVector,
            ) -> crate::MLResult<EnsemblePrediction> {
                Ok(EnsemblePrediction {
                    model_name: self.name.to_string(),
                    direction: self.direction,
                    confidence: self.confidence,
                    metadata: PredictionMeta::default(),
                })
            }
            fn is_ready(&self) -> bool {
                true
            }
        }
        unsafe impl Send for StubAdapter {}
        unsafe impl Sync for StubAdapter {}

        let mut coordinator = EnsembleCoordinator::new();
        coordinator.add_adapter(Box::new(StubAdapter { name: "DQN", direction: 0.6, confidence: 0.9 }));
        coordinator.add_adapter(Box::new(StubAdapter { name: "PPO", direction: 0.5, confidence: 0.85 }));
        coordinator.add_adapter(Box::new(StubAdapter { name: "TFT", direction: 0.4, confidence: 0.8 }));

        coordinator
            .register_model("DQN".to_string(), 0.33)
            .await
            .unwrap();
        coordinator
            .register_model("PPO".to_string(), 0.33)
            .await
            .unwrap();
        coordinator
            .register_model("TFT".to_string(), 0.34)
            .await
            .unwrap();

        let features = Features::new(
            vec![0.5, 0.6, 0.7, 0.8, 0.9],
            vec![
                "f1".to_string(),
                "f2".to_string(),
                "f3".to_string(),
                "f4".to_string(),
                "f5".to_string(),
            ],
        );

        let decision = coordinator.predict(&features).await.unwrap();

        assert!(decision.confidence >= 0.0 && decision.confidence <= 1.0);
        assert!(decision.signal >= -1.0 && decision.signal <= 1.0);
        assert_eq!(decision.model_count(), 3);
    }

    #[tokio::test]
    async fn test_disagreement_detection() {
        let aggregator = SignalAggregator::new();

        let predictions = vec![
            ModelPrediction::new("DQN".to_string(), 0.8, 0.9),
            ModelPrediction::new("PPO".to_string(), -0.7, 0.85),
            ModelPrediction::new("TFT".to_string(), 0.1, 0.8),
        ];

        let weights = HashMap::new();
        let decision = aggregator.aggregate(predictions, &weights).await.unwrap();

        assert!(decision.disagreement_rate > 0.3);
    }

    #[tokio::test]
    async fn test_weighted_voting() {
        let aggregator = SignalAggregator::new();

        let predictions = vec![
            ModelPrediction::new("DQN".to_string(), 0.8, 0.9),
            ModelPrediction::new("PPO".to_string(), 0.7, 0.85),
            ModelPrediction::new("TFT".to_string(), 0.6, 0.8),
        ];

        let mut weights = HashMap::new();
        weights.insert("DQN".to_string(), ModelWeight::new("DQN".to_string(), 0.5));
        weights.insert("PPO".to_string(), ModelWeight::new("PPO".to_string(), 0.3));
        weights.insert("TFT".to_string(), ModelWeight::new("TFT".to_string(), 0.2));

        let decision = aggregator.aggregate(predictions, &weights).await.unwrap();

        assert_eq!(decision.action, TradingAction::Buy);
        assert!(decision.signal > 0.6);
    }

    #[test]
    fn test_model_registry_swap() {
        let mut registry = ModelRegistry::new();

        registry.stage_checkpoint(
            "DQN".to_string(),
            "checkpoint_epoch_100.safetensors".to_string(),
        );

        registry.commit_swap("DQN").unwrap();

        assert!(registry.active.contains_key("DQN"));
    }

    #[tokio::test]
    async fn test_ensemble_with_real_adapter() {
        use crate::ensemble::inference_adapter::{
            EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
        };

        struct TestAdapter;
        impl ModelInferenceAdapter for TestAdapter {
            fn model_name(&self) -> &str {
                "DQN"
            }
            fn predict(
                &self,
                _features: &FeatureVector,
            ) -> crate::MLResult<EnsemblePrediction> {
                Ok(EnsemblePrediction {
                    model_name: "DQN".to_string(),
                    direction: 0.6,
                    confidence: 0.85,
                    metadata: PredictionMeta::default(),
                })
            }
            fn is_ready(&self) -> bool {
                true
            }
        }
        // Safety: TestAdapter has no mutable state, safe to share across threads
        unsafe impl Send for TestAdapter {}
        unsafe impl Sync for TestAdapter {}

        let mut coordinator = EnsembleCoordinator::new();
        coordinator.add_adapter(Box::new(TestAdapter));
        coordinator
            .register_model("DQN".to_string(), 1.0)
            .await
            .unwrap();

        let features = Features::new(
            vec![0.5; 10],
            vec!["f1", "f2", "f3", "f4", "f5", "f6", "f7", "f8", "f9", "f10"]
                .into_iter()
                .map(String::from)
                .collect(),
        );
        let decision = coordinator.predict(&features).await.unwrap();
        assert!(decision.confidence > 0.0);
        assert_eq!(decision.model_count(), 1);
    }

    #[tokio::test]
    async fn test_full_ensemble_with_dqn_adapter() {
        use crate::ensemble::adapters::DqnInferenceAdapter;
        use crate::dqn::dqn::DQNConfig;

        let dqn_config = DQNConfig {
            state_dim: 51,
            num_actions: 45,
            hidden_dims: vec![64, 64],
            ..Default::default()
        };

        let adapter = DqnInferenceAdapter::new(dqn_config).unwrap();

        let mut coordinator = EnsembleCoordinator::new();
        coordinator.add_adapter(Box::new(adapter));
        coordinator
            .register_model("DQN".to_string(), 1.0)
            .await
            .unwrap();

        let features = Features::new(
            vec![
                0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 0.1, 0.2, 0.3, 0.4, 0.5,
                0.6, 0.7, 0.8, 0.9, 1.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0,
                0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 0.1, 0.2, 0.3, 0.4, 0.5,
                0.6, 0.7, 0.8, 0.9, 1.0, 0.55,
            ],
            (0..51).map(|i| format!("f{}", i)).collect(),
        );

        let decision = coordinator.predict(&features).await.unwrap();

        assert!(decision.confidence > 0.0);
        assert!(decision.signal >= -1.0 && decision.signal <= 1.0);
        assert_eq!(decision.model_count(), 1);
    }

    #[tokio::test]
    async fn test_ensemble_rejects_no_adapters() {
        let coordinator = EnsembleCoordinator::new();
        coordinator
            .register_model("DQN".to_string(), 0.5)
            .await
            .unwrap();
        coordinator
            .register_model("PPO".to_string(), 0.5)
            .await
            .unwrap();

        let features = Features::new(
            vec![0.5; 5],
            vec!["f1", "f2", "f3", "f4", "f5"]
                .into_iter()
                .map(String::from)
                .collect(),
        );
        let result = coordinator.predict(&features).await;
        assert!(result.is_err(), "predict() must fail when no adapters are registered");
        let err_msg = format!("{}", result.unwrap_err());
        assert!(
            err_msg.contains("No model adapters registered"),
            "Expected 'No model adapters registered' error, got: {}",
            err_msg
        );
    }
}