foxhunt/crates/ml-ensemble/src/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::conviction_gates::{
    ConvictionGateConfig, ConvictionGateEvaluator, ConvictionGateOutcome, GateInput, TradingSession,
};
use crate::inference_adapter::{FeatureVector, ModelInferenceAdapter};
use crate::{EnsembleDecision, ModelVote, ModelWeight, TradingAction};
use crate::{Features, MLError, MLResult, ModelPrediction};
use chrono::{DateTime, Timelike, Utc};
use chrono_tz::America::New_York;
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>>,

    /// Conviction gate evaluator (optional -- None means no gating)
    conviction_gates: Option<ConvictionGateEvaluator>,
}

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(),
            conviction_gates: Some(ConvictionGateEvaluator::new(ConvictionGateConfig {
                // Allow all sessions by default -- trading service narrows to Regular
                allowed_sessions: vec![
                    TradingSession::PreMarket,
                    TradingSession::Regular,
                    TradingSession::AfterHours,
                ],
                ..ConvictionGateConfig::default()
            })),
        }
    }

    /// Replace conviction gate configuration
    pub fn set_conviction_gates(&mut self, config: ConvictionGateConfig) {
        self.conviction_gates = Some(ConvictionGateEvaluator::new(config));
    }

    /// Get mutable reference to conviction gate evaluator
    pub fn conviction_gates_mut(&mut self) -> Option<&mut ConvictionGateEvaluator> {
        self.conviction_gates.as_mut()
    }

    /// 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?;

        // Apply conviction gates if configured
        if let Some(ref gates) = self.conviction_gates {
            let gate_input = GateInput {
                confidence: decision.confidence,
                disagreement_rate: decision.disagreement_rate,
                quorum_ratio: Self::calculate_quorum_ratio(&decision),
                healthy_models: self.adapters.iter().filter(|a| a.is_ready()).count(),
                total_models: self.adapters.len(),
                current_session: Self::current_trading_session(),
                regime_volatility: Self::extract_regime_volatility(features),
            };

            match gates.evaluate(&gate_input) {
                ConvictionGateOutcome::Passed(pass_result) => {
                    info!(
                        "Conviction gates passed: score={:.3}, {} gates evaluated",
                        pass_result.conviction_score,
                        pass_result.gate_details.len()
                    );
                    let mut gated_decision = decision;
                    gated_decision.metadata.insert(
                        "conviction_score".into(),
                        serde_json::Value::from(pass_result.conviction_score),
                    );
                    return Ok(gated_decision);
                }
                ConvictionGateOutcome::Rejected(rejection) => {
                    info!("Conviction gate rejected: {:?} -- forcing HOLD", rejection);
                    let mut hold_decision = EnsembleDecision::new(
                        TradingAction::Hold,
                        decision.confidence,
                        0.0,
                        decision.disagreement_rate,
                        decision.model_votes,
                    );
                    hold_decision.metadata.insert(
                        "gate_rejection".into(),
                        serde_json::Value::String(format!("{:?}", rejection)),
                    );
                    return Ok(hold_decision);
                }
            }
        }

        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) => {
                        if !ensemble_pred.direction.is_finite()
                            || !ensemble_pred.confidence.is_finite()
                        {
                            warn!(
                                "Adapter {} returned non-finite prediction \
                                 (direction={}, confidence={}), skipping",
                                model_id, ensemble_pred.direction, ensemble_pred.confidence
                            );
                            continue;
                        }
                        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 model from production checkpoint (BF16 precision)
    ///
    /// Example:
    /// ```ignore
    /// coordinator.load_tft_checkpoint(
    ///     "TFT",
    ///     "ml/trained_models/production/tft/tft_epoch_200.safetensors",
    ///     0.15,
    /// ).await?;
    /// ```
    pub async fn load_tft_checkpoint(
        &self,
        model_id: &str,
        checkpoint: &str,
        weight: f64,
    ) -> MLResult<()> {
        info!("Loading TFT 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 checkpoint loaded and registered: {} (weight: {:.2})",
            model_id, weight
        );

        Ok(())
    }

    /// Calculate quorum ratio (fraction of models agreeing on majority direction)
    fn calculate_quorum_ratio(decision: &EnsembleDecision) -> f64 {
        if decision.model_votes.is_empty() {
            return 0.0;
        }
        let majority_action = &decision.action;
        let agreeing = decision
            .model_votes
            .values()
            .filter(|v| TradingAction::from_signal(v.signal, 0.3) == *majority_action)
            .count();
        agreeing as f64 / decision.model_votes.len() as f64
    }

    /// Determine current trading session (Eastern Time, DST-aware)
    fn current_trading_session() -> TradingSession {
        Self::trading_session_at(Utc::now())
    }

    /// Determine trading session for an arbitrary UTC timestamp.
    ///
    /// Uses `chrono-tz` America/New_York so the boundaries are correct
    /// during both EST (UTC-5, Nov-Mar) and EDT (UTC-4, Mar-Nov).
    fn trading_session_at(utc_time: DateTime<Utc>) -> TradingSession {
        let et = utc_time.with_timezone(&New_York);
        let et_time = et.hour() * 60 + et.minute();

        match et_time {
            t if t < 240 => TradingSession::AfterHours, // 00:00-04:00 ET
            t if t < 570 => TradingSession::PreMarket,  // 04:00-09:30 ET
            t if t < 960 => TradingSession::Regular,     // 09:30-16:00 ET
            _ => TradingSession::AfterHours,             // 16:00-24:00 ET
        }
    }

    /// Extract regime volatility from feature vector
    /// Regime features are at indices 48-50 in the 51-dim standard feature vector
    fn extract_regime_volatility(features: &Features) -> f64 {
        features
            .values
            .get(48)
            .copied()
            .unwrap_or(0.01)
            .abs()
    }
}

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);

            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_owned(),
            });
        }

        // Normalize effective weights so they sum to 1.0 before any calculations.
        // Raw effective_weight = static_weight * dynamic_weight can range 0.5-1.5
        // per model, so the sum across models is not guaranteed to be 1.0.
        let normalized = Self::normalize_effective_weights(&predictions, weights);

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

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

        // 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 (uses normalized weights)
        let model_votes = self.build_model_votes(&predictions, &normalized);

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

        Ok(decision)
    }

    /// Normalize effective weights across all predictions so they sum to 1.0.
    ///
    /// Models without an explicit weight entry receive equal share (1/N).
    /// If total weight is zero or non-finite, all models get uniform 1/N weight.
    fn normalize_effective_weights(
        predictions: &[ModelPrediction],
        weights: &HashMap<String, ModelWeight>,
    ) -> HashMap<String, f64> {
        let n = predictions.len();
        let uniform = if n > 0 { 1.0 / n as f64 } else { 0.0 };

        // Collect raw effective weights
        let raw: Vec<(String, f64)> = predictions
            .iter()
            .map(|pred| {
                let w = weights
                    .get(&pred.model_id)
                    .map(|mw| mw.effective_weight())
                    .unwrap_or(uniform);
                (pred.model_id.clone(), w)
            })
            .collect();

        let total: f64 = raw.iter().map(|(_, w)| w).sum();

        let mut normalized = HashMap::new();
        if total > 0.0 && total.is_finite() {
            for (id, w) in raw {
                normalized.insert(id, w / total);
            }
        } else {
            // Fallback: uniform weights
            for (id, _) in raw {
                normalized.insert(id, uniform);
            }
        }

        normalized
    }

    /// Calculate weighted average signal using pre-normalized weights
    fn calculate_weighted_signal(
        &self,
        predictions: &[ModelPrediction],
        normalized_weights: &HashMap<String, f64>,
    ) -> (f64, f64) {
        let n = predictions.len();
        let uniform = if n > 0 { 1.0 / n as f64 } else { 0.0 };
        let mut weighted_sum = 0.0;
        let mut total_weight = 0.0;

        for pred in predictions {
            let weight = normalized_weights
                .get(&pred.model_id)
                .copied()
                .unwrap_or(uniform);

            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 using pre-normalized weights
    fn calculate_ensemble_confidence(
        &self,
        predictions: &[ModelPrediction],
        normalized_weights: &HashMap<String, f64>,
    ) -> f64 {
        let n = predictions.len();
        let uniform = if n > 0 { 1.0 / n as f64 } else { 0.0 };
        let mut confidence_sum = 0.0;
        let mut weight_sum = 0.0;

        for pred in predictions {
            let weight = normalized_weights
                .get(&pred.model_id)
                .copied()
                .unwrap_or(uniform);

            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 using pre-normalized weights
    fn build_model_votes(
        &self,
        predictions: &[ModelPrediction],
        normalized_weights: &HashMap<String, f64>,
    ) -> HashMap<String, ModelVote> {
        let n = predictions.len();
        let uniform = if n > 0 { 1.0 / n as f64 } else { 0.0 };
        let mut votes = HashMap::new();

        for pred in predictions {
            let weight = normalized_weights
                .get(&pred.model_id)
                .copied()
                .unwrap_or(uniform);

            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)]
#[allow(
    unsafe_code,
    clippy::undocumented_unsafe_blocks,
    clippy::ok_expect,
    clippy::for_kv_map
)]
mod tests {
    use super::*;
    use chrono::TimeZone;

    #[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_owned(), 0.33)
            .await
            .unwrap();
        coordinator
            .register_model("PPO".to_owned(), 0.33)
            .await
            .unwrap();
        coordinator
            .register_model("TFT".to_owned(), 0.34)
            .await
            .unwrap();

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

    #[tokio::test]
    async fn test_ensemble_prediction() {
        use crate::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_owned(), 0.33)
            .await
            .unwrap();
        coordinator
            .register_model("PPO".to_owned(), 0.33)
            .await
            .unwrap();
        coordinator
            .register_model("TFT".to_owned(), 0.34)
            .await
            .unwrap();

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

        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_owned(), 0.8, 0.9),
            ModelPrediction::new("PPO".to_owned(), -0.7, 0.85),
            ModelPrediction::new("TFT".to_owned(), 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_owned(), 0.8, 0.9),
            ModelPrediction::new("PPO".to_owned(), 0.7, 0.85),
            ModelPrediction::new("TFT".to_owned(), 0.6, 0.8),
        ];

        let mut weights = HashMap::new();
        weights.insert("DQN".to_owned(), ModelWeight::new("DQN".to_owned(), 0.5));
        weights.insert("PPO".to_owned(), ModelWeight::new("PPO".to_owned(), 0.3));
        weights.insert("TFT".to_owned(), ModelWeight::new("TFT".to_owned(), 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_owned(),
            "checkpoint_epoch_100.safetensors".to_owned(),
        );

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

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

    #[tokio::test]
    async fn test_ensemble_with_real_adapter() {
        use crate::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_owned(),
                    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_owned(), 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);
    }

    // NOTE: test_full_ensemble_with_dqn_adapter lives in ml crate (depends on DqnInferenceAdapter)

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

        struct NanAdapter;
        impl ModelInferenceAdapter for NanAdapter {
            fn model_name(&self) -> &str {
                "NAN_MODEL"
            }
            fn predict(
                &self,
                _features: &FeatureVector,
            ) -> crate::MLResult<EnsemblePrediction> {
                Ok(EnsemblePrediction {
                    model_name: "NAN_MODEL".to_owned(),
                    direction: f64::NAN,
                    confidence: 0.9,
                    metadata: PredictionMeta::default(),
                })
            }
            fn is_ready(&self) -> bool {
                true
            }
        }
        unsafe impl Send for NanAdapter {}
        unsafe impl Sync for NanAdapter {}

        struct InfAdapter;
        impl ModelInferenceAdapter for InfAdapter {
            fn model_name(&self) -> &str {
                "INF_MODEL"
            }
            fn predict(
                &self,
                _features: &FeatureVector,
            ) -> crate::MLResult<EnsemblePrediction> {
                Ok(EnsemblePrediction {
                    model_name: "INF_MODEL".to_owned(),
                    direction: 0.5,
                    confidence: f64::INFINITY,
                    metadata: PredictionMeta::default(),
                })
            }
            fn is_ready(&self) -> bool {
                true
            }
        }
        unsafe impl Send for InfAdapter {}
        unsafe impl Sync for InfAdapter {}

        struct GoodAdapter;
        impl ModelInferenceAdapter for GoodAdapter {
            fn model_name(&self) -> &str {
                "GOOD_MODEL"
            }
            fn predict(
                &self,
                _features: &FeatureVector,
            ) -> crate::MLResult<EnsemblePrediction> {
                Ok(EnsemblePrediction {
                    model_name: "GOOD_MODEL".to_owned(),
                    direction: 0.6,
                    confidence: 0.85,
                    metadata: PredictionMeta::default(),
                })
            }
            fn is_ready(&self) -> bool {
                true
            }
        }
        unsafe impl Send for GoodAdapter {}
        unsafe impl Sync for GoodAdapter {}

        let mut coordinator = EnsembleCoordinator::new();
        coordinator.add_adapter(Box::new(NanAdapter));
        coordinator.add_adapter(Box::new(InfAdapter));
        coordinator.add_adapter(Box::new(GoodAdapter));

        coordinator
            .register_model("NAN_MODEL".to_owned(), 0.33)
            .await
            .unwrap();
        coordinator
            .register_model("INF_MODEL".to_owned(), 0.33)
            .await
            .unwrap();
        coordinator
            .register_model("GOOD_MODEL".to_owned(), 0.34)
            .await
            .unwrap();

        let features = Features::new(
            vec![0.5; 5],
            vec!["f1", "f2", "f3", "f4", "f5"]
                .into_iter()
                .map(String::from)
                .collect(),
        );

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

        // Only the good model should survive filtering
        assert_eq!(
            decision.model_count(),
            1,
            "NaN and Inf predictions should be filtered out"
        );
        assert!(
            decision.confidence.is_finite(),
            "Ensemble confidence must be finite"
        );
        assert!(
            decision.signal.is_finite(),
            "Ensemble signal must be finite"
        );
    }

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

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

        let mut coordinator = EnsembleCoordinator::new();
        coordinator.add_adapter(Box::new(AllNanAdapter { name: "A" }));
        coordinator.add_adapter(Box::new(AllNanAdapter { name: "B" }));

        coordinator
            .register_model("A".to_owned(), 0.5)
            .await
            .unwrap();
        coordinator
            .register_model("B".to_owned(), 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(),
            "All-NaN predictions should produce an error"
        );
    }

    #[tokio::test]
    async fn test_ensemble_rejects_no_adapters() {
        let coordinator = EnsembleCoordinator::new();
        coordinator
            .register_model("DQN".to_owned(), 0.5)
            .await
            .unwrap();
        coordinator
            .register_model("PPO".to_owned(), 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
        );
    }

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

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

        // Create weights with dynamic adjustments that cause unnormalized sums.
        // DQN: high perf => dynamic ~1.2, effective = 0.5 * 1.2 = 0.6
        // PPO: low perf  => dynamic ~0.6, effective = 0.3 * 0.6 = 0.18
        // TFT: default   => dynamic  1.0, effective = 0.2 * 1.0 = 0.2
        // Raw sum = 0.98, NOT 1.0
        let mut weights = HashMap::new();

        let mut dqn_w = ModelWeight::new("DQN".to_owned(), 0.5);
        dqn_w.performance_metrics.sharpe_ratio = 2.0;
        dqn_w.performance_metrics.accuracy = 0.7;
        dqn_w.update_dynamic_weight();
        weights.insert("DQN".to_owned(), dqn_w);

        let mut ppo_w = ModelWeight::new("PPO".to_owned(), 0.3);
        ppo_w.performance_metrics.sharpe_ratio = 0.3;
        ppo_w.performance_metrics.accuracy = 0.3;
        ppo_w.update_dynamic_weight();
        weights.insert("PPO".to_owned(), ppo_w);

        let mut tft_w = ModelWeight::new("TFT".to_owned(), 0.2);
        // Keep default performance (sharpe=1.0, accuracy=0.5) => dynamic ~0.95
        tft_w.update_dynamic_weight();
        weights.insert("TFT".to_owned(), tft_w);

        // Verify raw effective weights do NOT sum to 1.0
        let raw_sum: f64 = weights.values().map(|w| w.effective_weight()).sum();
        assert!(
            (raw_sum - 1.0).abs() > 0.01,
            "Raw effective weights should NOT sum to 1.0 (got {}), \
             otherwise the test does not exercise normalization",
            raw_sum
        );

        // After aggregation, model_votes weights must sum to 1.0
        let decision = aggregator
            .aggregate(predictions, &weights)
            .await
            .ok()
            .expect("aggregation should succeed");

        let vote_weight_sum: f64 = decision.model_votes.values().map(|v| v.weight).sum();
        assert!(
            (vote_weight_sum - 1.0).abs() < 1e-10,
            "Normalized vote weights must sum to 1.0, got {}",
            vote_weight_sum
        );

        // Each individual weight must be positive
        for (model_id, vote) in &decision.model_votes {
            assert!(
                vote.weight > 0.0,
                "Weight for {} must be positive, got {}",
                model_id,
                vote.weight
            );
        }
    }

    #[test]
    fn test_normalize_effective_weights_uniform_fallback() {
        // When weights map is empty, all models get 1/N
        let predictions = vec![
            ModelPrediction::new("A".to_owned(), 0.5, 0.8),
            ModelPrediction::new("B".to_owned(), 0.3, 0.7),
        ];
        let weights = HashMap::new();

        let normalized = SignalAggregator::normalize_effective_weights(&predictions, &weights);

        assert_eq!(normalized.len(), 2);
        let sum: f64 = normalized.values().sum();
        assert!(
            (sum - 1.0).abs() < 1e-10,
            "Uniform fallback weights must sum to 1.0, got {}",
            sum
        );
        for (_, w) in &normalized {
            assert!(
                (*w - 0.5).abs() < 1e-10,
                "Each weight should be 0.5 for 2 models, got {}",
                w
            );
        }
    }

    #[test]
    fn test_normalize_effective_weights_single_model() {
        let predictions = vec![ModelPrediction::new("SOLO".to_owned(), 0.9, 0.95)];
        let mut weights = HashMap::new();
        weights.insert(
            "SOLO".to_owned(),
            ModelWeight::new("SOLO".to_owned(), 0.7),
        );

        let normalized = SignalAggregator::normalize_effective_weights(&predictions, &weights);

        assert_eq!(normalized.len(), 1);
        let w = normalized.get("SOLO").copied().unwrap_or(0.0);
        assert!(
            (w - 1.0).abs() < 1e-10,
            "Single model must get weight 1.0, got {}",
            w
        );
    }

    /// Verify DST-aware trading session boundaries.
    ///
    /// June 15 2026 is during EDT (UTC-4), so 13:30 UTC = 09:30 ET (Regular).
    /// The old hardcoded UTC-5 would have mapped 13:30 UTC to 08:30 ET (PreMarket) -- wrong.
    #[test]
    fn test_trading_session_dst_correctness() {
        // --- EDT (summer): UTC-4 ---
        // June 15 2026 13:30 UTC  =>  09:30 ET  =>  Regular
        let june_1330 = Utc.with_ymd_and_hms(2026, 6, 15, 13, 30, 0).unwrap();
        assert_eq!(
            EnsembleCoordinator::trading_session_at(june_1330),
            TradingSession::Regular,
            "13:30 UTC in June (EDT) should be 09:30 ET = Regular, not PreMarket"
        );

        // June 15 2026 20:00 UTC  =>  16:00 ET  =>  AfterHours
        let june_2000 = Utc.with_ymd_and_hms(2026, 6, 15, 20, 0, 0).unwrap();
        assert_eq!(
            EnsembleCoordinator::trading_session_at(june_2000),
            TradingSession::AfterHours,
            "20:00 UTC in June (EDT) should be 16:00 ET = AfterHours"
        );

        // June 15 2026 08:00 UTC  =>  04:00 ET  =>  PreMarket
        let june_0800 = Utc.with_ymd_and_hms(2026, 6, 15, 8, 0, 0).unwrap();
        assert_eq!(
            EnsembleCoordinator::trading_session_at(june_0800),
            TradingSession::PreMarket,
            "08:00 UTC in June (EDT) should be 04:00 ET = PreMarket"
        );

        // --- EST (winter): UTC-5 ---
        // January 15 2026 14:30 UTC  =>  09:30 ET  =>  Regular
        let jan_1430 = Utc.with_ymd_and_hms(2026, 1, 15, 14, 30, 0).unwrap();
        assert_eq!(
            EnsembleCoordinator::trading_session_at(jan_1430),
            TradingSession::Regular,
            "14:30 UTC in January (EST) should be 09:30 ET = Regular"
        );

        // January 15 2026 21:00 UTC  =>  16:00 ET  =>  AfterHours
        let jan_2100 = Utc.with_ymd_and_hms(2026, 1, 15, 21, 0, 0).unwrap();
        assert_eq!(
            EnsembleCoordinator::trading_session_at(jan_2100),
            TradingSession::AfterHours,
            "21:00 UTC in January (EST) should be 16:00 ET = AfterHours"
        );

        // January 15 2026 03:00 UTC  =>  22:00 ET (prev day)  =>  AfterHours
        let jan_0300 = Utc.with_ymd_and_hms(2026, 1, 15, 3, 0, 0).unwrap();
        assert_eq!(
            EnsembleCoordinator::trading_session_at(jan_0300),
            TradingSession::AfterHours,
            "03:00 UTC in January (EST) should be 22:00 ET = AfterHours"
        );
    }
}