foxhunt/crates/ml-ensemble/src/model.rs

//! Ensemble Model Management for Trading Signal Aggregation
//!
//! Provides a unified interface for managing multiple ML models, aggregating their
//! signals, and maintaining performance tracking for ultra-low latency trading.

use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};

use crossbeam::atomic::AtomicCell;
use serde::{Deserialize, Serialize};

use crate::aggregator::ModelSignal;
use crate::{HealthStatus, MLError};
use common::trading::MarketRegime;

/// Configuration for ensemble models
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnsembleConfig {
    pub min_models: usize,
    pub max_models: usize,
    pub signal_timeout_ms: u64,
    pub aggregation_method: AggregationMethod,
    pub confidence_threshold: f32,
}

impl Default for EnsembleConfig {
    fn default() -> Self {
        Self {
            min_models: 2,
            max_models: 10,
            signal_timeout_ms: 1000,
            aggregation_method: AggregationMethod::WeightedAverage,
            confidence_threshold: 0.5,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AggregationMethod {
    WeightedAverage,
    MajorityVote,
    AdaptiveWeighted,
}

// Use ModelSignal and SignalMetadata from aggregator.rs

/// Ensemble signal result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnsembleSignal {
    pub value: f32,
    pub confidence: f32,
    pub contributing_models: usize,
    pub timestamp: u64,
    pub regime: MarketRegime,
}

/// Health status information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthInfo {
    pub status: HealthStatus,
    pub total_models: usize,
    pub active_models: usize,
    pub last_signal_time: Option<u64>,
    pub error_rate: f32,
}

/// Ensemble metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnsembleMetrics {
    pub total_models: usize,
    pub active_models: usize,
    pub current_regime: MarketRegime,
    pub total_signals: u64,
    pub successful_aggregations: u64,
    pub average_latency_us: f32,
}

/// Model information
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ModelInfo {
    model_id: String,
    model_type: String,
    version: String,
    features: Vec<String>,
    expected_latency_us: u64,
    last_signal_time: Option<u64>,
    signal_count: u64,
    error_count: u64,
}

/// Main ensemble model struct
#[derive(Debug)]
pub struct EnsembleModel {
    config: EnsembleConfig,
    models: Arc<RwLock<HashMap<String, ModelInfo>>>,
    signals: Arc<RwLock<Vec<ModelSignal>>>,
    current_regime: Arc<AtomicCell<MarketRegime>>,
    metrics: Arc<RwLock<EnsembleMetrics>>,
}

impl EnsembleModel {
    /// Create a new ensemble model
    ///
    /// # Errors
    ///
    /// Returns `MLError` if:
    /// - Configuration validation fails
    /// - Internal data structure initialization fails
    /// - Memory allocation fails
    pub fn new(config: EnsembleConfig) -> Result<Self, MLError> {
        Ok(Self {
            config,
            models: Arc::new(RwLock::new(HashMap::new())),
            signals: Arc::new(RwLock::new(Vec::new())),
            current_regime: Arc::new(AtomicCell::new(MarketRegime::Normal)),
            metrics: Arc::new(RwLock::new(EnsembleMetrics {
                total_models: 0,
                active_models: 0,
                current_regime: MarketRegime::Normal,
                total_signals: 0,
                successful_aggregations: 0,
                average_latency_us: 0.0,
            })),
        })
    }

    /// Register a new model
    pub fn register_model(
        &self,
        model_id: &str,
        model_type: &str,
        version: &str,
        features: Vec<String>,
        expected_latency_us: u64,
    ) -> Result<(), MLError> {
        let mut models = self
            .models
            .write()
            .map_err(|e| MLError::LockError(e.to_string()))?;

        let model_info = ModelInfo {
            model_id: model_id.to_string(),
            model_type: model_type.to_string(),
            version: version.to_string(),
            features,
            expected_latency_us,
            last_signal_time: None,
            signal_count: 0,
            error_count: 0,
        };

        models.insert(model_id.to_string(), model_info);

        // Update metrics
        let mut metrics = self
            .metrics
            .write()
            .map_err(|e| MLError::LockError(e.to_string()))?;
        metrics.total_models = models.len();
        metrics.active_models = models.len();

        Ok(())
    }

    /// Unregister a model
    ///
    /// # Errors
    ///
    /// Returns `MLError` if:
    /// - Model ID not found in registry
    /// - Lock acquisition fails
    /// - Metrics update fails
    /// - Model removal fails
    pub fn unregister_model(&self, model_id: &str) -> Result<(), MLError> {
        let mut models = self
            .models
            .write()
            .map_err(|e| MLError::LockError(e.to_string()))?;

        if models.remove(model_id).is_none() {
            return Err(MLError::ModelNotFound(model_id.to_string()));
        }

        // Update metrics
        let mut metrics = self
            .metrics
            .write()
            .map_err(|e| MLError::LockError(e.to_string()))?;
        metrics.total_models = models.len();
        metrics.active_models = models.len();

        Ok(())
    }

    /// Submit a signal from a model
    ///
    /// # Errors
    ///
    /// Returns `MLError` if:
    /// - Lock acquisition fails
    /// - Model ID not found
    /// - Signal validation fails
    /// - Timestamp is invalid
    /// - Signal storage fails
    pub fn submit_signal(&self, signal: ModelSignal) -> Result<(), MLError> {
        let mut signals = self
            .signals
            .write()
            .map_err(|e| MLError::LockError(e.to_string()))?;

        // Update model info
        {
            let mut models = self
                .models
                .write()
                .map_err(|e| MLError::LockError(e.to_string()))?;
            if let Some(model_info) = models.get_mut(&signal.model_id) {
                model_info.last_signal_time = Some(signal.timestamp);
                model_info.signal_count += 1;
            }
        }

        signals.push(signal);

        // Clean old signals
        let current_time = current_timestamp();
        signals.retain(|s| current_time - s.timestamp < self.config.signal_timeout_ms);

        Ok(())
    }

    /// Aggregate signals from all models
    ///
    /// # Errors
    ///
    /// Returns `MLError` if:
    /// - Lock acquisition fails
    /// - Insufficient models have submitted signals
    /// - Signal aggregation calculation fails
    /// - Confidence calculation fails
    /// - No valid signals available
    /// - Consensus cannot be reached
    /// - Regime detection fails
    pub fn aggregate_signals(&self) -> Result<EnsembleSignal, MLError> {
        let signals = self
            .signals
            .read()
            .map_err(|e| MLError::LockError(e.to_string()))?;
        let _models = self
            .models
            .read()
            .map_err(|e| MLError::LockError(e.to_string()))?;

        let current_time = current_timestamp();
        let recent_signals: Vec<_> = signals
            .iter()
            .filter(|s| current_time - s.timestamp < self.config.signal_timeout_ms)
            .collect();

        if recent_signals.len() < self.config.min_models {
            return Err(MLError::InsufficientData(format!(
                "Need at least {} models, got {}",
                self.config.min_models,
                recent_signals.len()
            )));
        }

        let (aggregated_value, aggregated_confidence) = match self.config.aggregation_method {
            AggregationMethod::WeightedAverage => {
                self.weighted_average_aggregation(&recent_signals)
            },
            AggregationMethod::MajorityVote => self.majority_vote_aggregation(&recent_signals),
            AggregationMethod::AdaptiveWeighted => {
                self.adaptive_weighted_aggregation(&recent_signals)
            },
        };

        Ok(EnsembleSignal {
            value: aggregated_value,
            confidence: aggregated_confidence,
            contributing_models: recent_signals.len(),
            timestamp: current_time,
            regime: self.current_regime.load(),
        })
    }

    fn weighted_average_aggregation(&self, signals: &[&ModelSignal]) -> (f32, f32) {
        let total_weight: f32 = signals.iter().map(|s| s.confidence).sum();
        if total_weight == 0.0 {
            return (0.0, 0.0);
        }

        let weighted_value: f32 =
            signals.iter().map(|s| s.signal * s.confidence).sum::<f32>() / total_weight;

        let average_confidence: f32 = total_weight / signals.len() as f32;

        (weighted_value as f32, average_confidence as f32)
    }

    fn majority_vote_aggregation(&self, signals: &[&ModelSignal]) -> (f32, f32) {
        // Simple majority vote for binary signals
        let positive_count = signals.iter().filter(|s| s.signal > 0.0).count();
        let total_count = signals.len();

        let majority_value = if positive_count * 2 > total_count {
            1.0_f32
        } else {
            -1.0_f32
        };
        let confidence =
            (positive_count.max(total_count - positive_count) as f32) / (total_count as f32);

        (majority_value, confidence)
    }

    fn adaptive_weighted_aggregation(&self, signals: &[&ModelSignal]) -> (f32, f32) {
        // Use recent performance to adjust weights
        self.weighted_average_aggregation(signals) // Simplified - same as weighted for now
    }

    /// Update market regime
    pub fn update_market_regime(&self, regime: MarketRegime) {
        self.current_regime.store(regime);

        // Update metrics
        if let Ok(mut metrics) = self.metrics.write() {
            metrics.current_regime = regime;
        }
    }

    /// Get list of registered models
    pub fn list_models(&self) -> Vec<String> {
        if let Ok(models) = self.models.read() {
            models.keys().cloned().collect()
        } else {
            Vec::new()
        }
    }

    /// Health check
    pub fn health_check(&self) -> HealthInfo {
        let models = self.models.read().map_err(|e| {
            tracing::error!("Models RwLock poisoned in health_check: {}", e);
            e
        }).ok();
        let signals = self.signals.read().map_err(|e| {
            tracing::error!("Signals RwLock poisoned in health_check: {}", e);
            e
        }).ok();

        let total_models = models.as_ref().map(|m| m.len()).unwrap_or(0);
        let active_models = total_models; // Simplified

        let last_signal_time = signals
            .as_ref()
            .and_then(|s| s.last().map(|sig| sig.timestamp));

        let status = if total_models >= self.config.min_models {
            HealthStatus::Healthy
        } else if total_models > 0 {
            HealthStatus::Degraded
        } else {
            HealthStatus::Unhealthy
        };

        HealthInfo {
            status,
            total_models,
            active_models,
            last_signal_time,
            error_rate: 0.0, // Simplified
        }
    }

    /// Get current metrics
    pub fn get_metrics(&self) -> Result<EnsembleMetrics, MLError> {
        let metrics = self
            .metrics
            .read()
            .map_err(|e| MLError::LockError(e.to_string()))?;
        Ok(metrics.clone())
    }
}

/// Helper function to get current timestamp
fn current_timestamp() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis() as u64
}