foxhunt/crates/common/src/ml_strategy.rs

//! Shared ML Strategy for Foxhunt Trading System
//!
//! This module provides a unified ML strategy implementation that is used by both
//! trading service and backtesting service to ensure consistent ML predictions
//! across all services. This eliminates code duplication and ensures ONE SINGLE SYSTEM.
//!
//! # Architecture
//!
//! ```text
//! SharedMLStrategy
//!     ├─ MLModelAdapter (abstraction over ml crate models)
//!     ├─ ProductionFeatureExtractor225 (consistent 225-feature engineering)
//!     ├─ EnsembleCoordinator (weighted voting)
//!     └─ ModelPerformanceTracker (metrics)
//! ```

use anyhow::Result;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// Trait for pluggable 225-feature extraction
///
/// This trait allows applications to inject the production-grade feature extractor
/// from the `ml` crate without creating circular dependencies.
pub trait ProductionFeatureExtractor225: Send + Sync {
    /// Update internal state with new market data
    fn update(&mut self, price: f64, volume: f64, timestamp: DateTime<Utc>) -> Result<()>;

    /// Extract 225-dimensional feature vector from current state
    fn extract_features(&mut self) -> Result<Vec<f64>>;
}

/// ML prediction result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MLPrediction {
    /// Model identifier
    pub model_id: String,
    /// Prediction value (0.0-1.0)
    pub prediction_value: f64,
    /// Confidence score (0.0-1.0)
    pub confidence: f64,
    /// Features used for prediction
    pub features: Vec<f64>,
    /// Prediction timestamp
    pub timestamp: DateTime<Utc>,
    /// Inference latency in microseconds
    pub inference_latency_us: u64,
}

/// ML model performance metrics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct MLModelPerformance {
    /// Model identifier
    pub model_id: String,
    /// Total predictions made
    pub total_predictions: u64,
    /// Correct predictions
    pub correct_predictions: u64,
    /// Average inference latency
    pub avg_latency_us: f64,
    /// Average confidence score
    pub avg_confidence: f64,
    /// Model accuracy percentage
    pub accuracy_percentage: f64,
    /// Returns generated
    pub returns: Vec<f64>,
    /// Sharpe ratio
    pub sharpe_ratio: f64,
    /// Maximum drawdown
    pub max_drawdown: f64,
}

/// Trait for ML model adapters
pub trait MLModelAdapter: Send + Sync {
    /// Get model prediction
    fn predict(&self, features: &[f64]) -> Result<MLPrediction>;

    /// Get model identifier
    fn model_id(&self) -> &str;

    /// Validate prediction against actual outcome
    fn validate_prediction(&mut self, prediction: &MLPrediction, actual_outcome: bool);
}

/// Shared ML strategy implementation (ONE SINGLE SYSTEM)
pub struct SharedMLStrategy {
    /// Available ML models
    models: Arc<RwLock<HashMap<String, Box<dyn MLModelAdapter>>>>,
    /// Feature extractor - pluggable 225-feature extractor (inject from ml crate)
    feature_extractor_225: Option<Arc<RwLock<Box<dyn ProductionFeatureExtractor225>>>>,
    /// Model performance tracking
    model_performance: Arc<RwLock<HashMap<String, MLModelPerformance>>>,
    /// Minimum confidence threshold
    min_confidence_threshold: f64,
}

impl std::fmt::Debug for SharedMLStrategy {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SharedMLStrategy")
            .field("min_confidence_threshold", &self.min_confidence_threshold)
            .field("model_count", &"<async_lock>")
            .finish()
    }
}

impl SharedMLStrategy {
    /// Create strategy with injected models and production feature extractor.
    ///
    /// # Arguments
    /// - `extractor`: Production-grade 225-feature extractor (inject from ml crate)
    /// - `models`: Pre-built model adapters for ensemble prediction
    /// - `min_confidence_threshold`: Minimum confidence for predictions
    ///
    /// If `models` is empty, predictions will return an empty vec (graceful degradation).
    pub fn new(
        extractor: Box<dyn ProductionFeatureExtractor225>,
        models: Vec<Box<dyn MLModelAdapter>>,
        min_confidence_threshold: f64,
    ) -> Self {
        let model_map: HashMap<String, Box<dyn MLModelAdapter>> = models
            .into_iter()
            .map(|m| (m.model_id().to_string(), m))
            .collect();

        Self {
            models: Arc::new(RwLock::new(model_map)),
            feature_extractor_225: Some(Arc::new(RwLock::new(extractor))),
            model_performance: Arc::new(RwLock::new(HashMap::new())),
            min_confidence_threshold,
        }
    }

    /// Get ensemble prediction from all models
    pub async fn get_ensemble_prediction(
        &self,
        price: f64,
        volume: f64,
        timestamp: DateTime<Utc>,
    ) -> Result<Vec<MLPrediction>> {
        // Extract features using production 225-feature extractor
        let features: Vec<f64> = match self.feature_extractor_225 {
            Some(ref prod_extractor) => {
                let mut extractor = prod_extractor.write().await;
                extractor.update(price, volume, timestamp)?;
                extractor.extract_features()?
            }
            None => {
                anyhow::bail!("No feature extractor configured")
            }
        };

        // Validate feature count
        if features.len() != 225 {
            tracing::error!(
                "Feature extraction returned {} features, expected 225. \
                 Models will receive incorrect input!",
                features.len()
            );
        }

        let mut predictions = Vec::new();

        // Get predictions from all models
        let models = self.models.read().await;
        for (model_id, model) in models.iter() {
            match model.predict(&features) {
                Ok(prediction) => {
                    if prediction.confidence >= self.min_confidence_threshold {
                        predictions.push(prediction);
                    }
                }
                Err(e) => {
                    tracing::warn!("Model {} failed to predict: {}", model_id, e);
                }
            }
        }

        Ok(predictions)
    }

    /// Calculate weighted ensemble vote
    pub fn calculate_ensemble_vote(&self, predictions: &[MLPrediction]) -> Option<(f64, f64)> {
        if predictions.is_empty() {
            return None;
        }

        let total_confidence: f64 = predictions.iter().map(|p| p.confidence).sum();
        if total_confidence == 0.0 {
            return None;
        }

        // Weighted average by confidence
        let weighted_prediction: f64 = predictions
            .iter()
            .map(|p| p.prediction_value * p.confidence)
            .sum::<f64>()
            / total_confidence;

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

        Some((weighted_prediction, average_confidence))
    }

    /// Validate predictions against actual market outcomes
    pub async fn validate_predictions(&self, predictions: &[MLPrediction], actual_return: f64) {
        let actual_outcome = actual_return > 0.0; // Positive return = good outcome

        let mut models = self.models.write().await;
        let mut performance = self.model_performance.write().await;

        for prediction in predictions {
            if let Some(model) = models.get_mut(&prediction.model_id) {
                model.validate_prediction(prediction, actual_outcome);
            }

            // Update performance tracking
            let perf = performance
                .entry(prediction.model_id.clone())
                .or_insert_with(|| MLModelPerformance {
                    model_id: prediction.model_id.clone(),
                    ..Default::default()
                });

            perf.total_predictions += 1;

            let predicted_positive = prediction.prediction_value > 0.5;
            if predicted_positive == actual_outcome {
                perf.correct_predictions += 1;
            }

            perf.accuracy_percentage = if perf.total_predictions > 0 {
                (perf.correct_predictions as f64 / perf.total_predictions as f64) * 100.0
            } else {
                0.0
            };

            // Update average confidence
            let total_samples = perf.total_predictions as f64;
            perf.avg_confidence = (perf.avg_confidence * (total_samples - 1.0)
                + prediction.confidence)
                / total_samples;

            // Update average latency
            perf.avg_latency_us = (perf.avg_latency_us * (total_samples - 1.0)
                + prediction.inference_latency_us as f64)
                / total_samples;
        }
    }

    /// Get performance summary for all models
    pub async fn get_performance_summary(&self) -> HashMap<String, MLModelPerformance> {
        self.model_performance.read().await.clone()
    }

    /// Add a model to the strategy
    pub async fn add_model(&self, model_id: String, model: Box<dyn MLModelAdapter>) {
        let mut models = self.models.write().await;
        models.insert(model_id, model);
    }

    /// Get minimum confidence threshold
    pub fn min_confidence_threshold(&self) -> f64 {
        self.min_confidence_threshold
    }
}

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

    /// Minimal mock adapter for unit tests
    struct MockAdapter {
        id: String,
    }

    impl MockAdapter {
        fn new(id: &str) -> Self {
            Self {
                id: id.to_string(),
            }
        }
    }

    impl MLModelAdapter for MockAdapter {
        fn predict(&self, features: &[f64]) -> Result<MLPrediction> {
            let sum: f64 = features.iter().sum::<f64>() / features.len().max(1) as f64;
            let prediction_value = 1.0 / (1.0 + (-sum).exp());
            let confidence = 0.5 + (prediction_value - 0.5).abs() * 0.8;
            Ok(MLPrediction {
                model_id: self.id.clone(),
                prediction_value,
                confidence,
                features: features.to_vec(),
                timestamp: Utc::now(),
                inference_latency_us: 10,
            })
        }

        fn model_id(&self) -> &str {
            &self.id
        }

        fn validate_prediction(&mut self, _prediction: &MLPrediction, _actual_outcome: bool) {}
    }

    /// Minimal mock extractor for unit tests (returns 225 values)
    struct MockExtractor;

    impl ProductionFeatureExtractor225 for MockExtractor {
        fn update(&mut self, _price: f64, _volume: f64, _timestamp: DateTime<Utc>) -> Result<()> {
            Ok(())
        }
        fn extract_features(&mut self) -> Result<Vec<f64>> {
            Ok(vec![0.1; 225])
        }
    }

    #[tokio::test]
    async fn test_shared_ml_strategy_creation() {
        let strategy = SharedMLStrategy::new(
            Box::new(MockExtractor),
            vec![Box::new(MockAdapter::new("m1"))],
            0.6,
        );
        assert_eq!(strategy.min_confidence_threshold(), 0.6);
    }

    #[tokio::test]
    async fn test_ensemble_prediction() {
        let strategy = SharedMLStrategy::new(
            Box::new(MockExtractor),
            vec![Box::new(MockAdapter::new("m1"))],
            0.0,
        );
        let predictions = strategy
            .get_ensemble_prediction(100.0, 1000.0, Utc::now())
            .await
            .unwrap_or_default();
        assert!(!predictions.is_empty());
    }

    #[tokio::test]
    async fn test_ensemble_vote() {
        let strategy = SharedMLStrategy::new(
            Box::new(MockExtractor),
            vec![Box::new(MockAdapter::new("m1"))],
            0.0,
        );
        let predictions = vec![
            MLPrediction {
                model_id: "model1".to_string(),
                prediction_value: 0.8,
                confidence: 0.9,
                features: vec![],
                timestamp: Utc::now(),
                inference_latency_us: 50,
            },
            MLPrediction {
                model_id: "model2".to_string(),
                prediction_value: 0.6,
                confidence: 0.7,
                features: vec![],
                timestamp: Utc::now(),
                inference_latency_us: 60,
            },
        ];
        let (vote, confidence) = strategy
            .calculate_ensemble_vote(&predictions)
            .unwrap_or_default();
        assert!((0.6..=0.8).contains(&vote));
        assert!((0.7..=0.9).contains(&confidence));
    }

    #[tokio::test]
    async fn test_performance_tracking() {
        let strategy = SharedMLStrategy::new(
            Box::new(MockExtractor),
            vec![Box::new(MockAdapter::new("test_model"))],
            0.0,
        );
        let prediction = MLPrediction {
            model_id: "test_model".to_string(),
            prediction_value: 0.7,
            confidence: 0.8,
            features: vec![],
            timestamp: Utc::now(),
            inference_latency_us: 50,
        };
        strategy
            .validate_predictions(std::slice::from_ref(&prediction), 0.05)
            .await;
        let performance = strategy.get_performance_summary().await;
        let perf = performance.get("test_model").cloned().unwrap_or_default();
        assert_eq!(perf.total_predictions, 1);
        assert_eq!(perf.correct_predictions, 1);
        assert_eq!(perf.accuracy_percentage, 100.0);
    }

    #[tokio::test]
    async fn test_empty_models_graceful() {
        let strategy =
            SharedMLStrategy::new(Box::new(MockExtractor), vec![], 0.0);
        let predictions = strategy
            .get_ensemble_prediction(100.0, 1000.0, Utc::now())
            .await
            .unwrap_or_default();
        assert!(
            predictions.is_empty(),
            "No models should produce no predictions"
        );
    }
}