foxhunt/ml/src/traits.rs

//! Core traits for ML models in the Foxhunt HFT system
//!
//! These traits provide a unified interface for all ML models, enabling
//! consistent integration with the trading engine and performance monitoring.

use crate::{InferenceResult, ModelMetadata, TrainingMetrics, ValidationMetrics};
use async_trait::async_trait;
use ndarray::Array2;
use serde::{Deserialize, Serialize};

// DO NOT RE-EXPORT - Use explicit imports at usage sites
// Note: MLModel trait is defined separately in tgnn::traits

/// Core ML model trait for all models in the system
#[async_trait]
pub trait MLModelCore {
    type Config;

    /// Get model metadata
    fn metadata(&self) -> &ModelMetadata;

    /// Check if model is ready for inference
    fn is_ready(&self) -> bool;

    /// Train the model
    async fn train(
        &mut self,
        features: &Array2<f64>,
        targets: &Array2<f64>,
    ) -> Result<TrainingMetrics, crate::MLError>;

    /// Predict using the model
    async fn predict(&self, features: &[f64]) -> Result<InferenceResult, crate::MLError>;

    /// Validate model performance
    async fn validate(
        &self,
        features: &Array2<f64>,
        targets: &Array2<f64>,
    ) -> Result<ValidationMetrics, crate::MLError>;

    /// Update model with new data (online learning)
    async fn update(
        &mut self,
        features: &Array2<f64>,
        targets: &Array2<f64>,
    ) -> Result<(), crate::MLError>;

    /// Save model to file
    async fn save(&self, path: &str) -> Result<(), crate::MLError>;

    /// Load model from file
    async fn load(&mut self, path: &str) -> Result<(), crate::MLError>;

    /// Get model configuration
    fn config(&self) -> Self::Config;

    /// Set model configuration
    fn set_config(&mut self, config: Self::Config) -> Result<(), crate::MLError>;
}

/// Performance metrics tracking for ML models
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PerformanceMetrics {
    pub average_inference_latency_us: u64,
    pub max_inference_latency_us: u64,
    pub throughput_pps: u64,
    pub memory_usage_bytes: u64,
    pub total_predictions: u64,
    pub error_rate: f64,
}

impl Default for PerformanceMetrics {
    fn default() -> Self {
        Self {
            average_inference_latency_us: 0,
            max_inference_latency_us: 0,
            throughput_pps: 0,
            memory_usage_bytes: 0,
            total_predictions: 0,
            error_rate: 0.0,
        }
    }
}

/// Streaming statistics for real-time monitoring
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamingStats {
    pub total_processed: u64,
    pub min_latency_us: u64,
    pub max_latency_us: u64,
    pub avg_latency_us: u64,
    pub throughput_pps: f64,
}

impl Default for StreamingStats {
    fn default() -> Self {
        Self {
            total_processed: 0,
            min_latency_us: u64::MAX,
            max_latency_us: 0,
            avg_latency_us: 0,
            throughput_pps: 0.0,
        }
    }
}

/// Trait for models that track performance metrics
pub trait ModelPerformance {
    /// Get current performance metrics
    fn performance_metrics(&self) -> PerformanceMetrics;

    /// Reset performance counters
    fn reset_performance_counters(&mut self);

    /// Check if model meets performance targets
    fn meets_performance_targets(&self) -> bool {
        let metrics = self.performance_metrics();
        metrics.average_inference_latency_us < 100 && // Sub-100μs target
        metrics.throughput_pps > 100_000 // Over 100K predictions per second
    }
}

/// Trait for models that can predict from Array2 features (batch prediction)
#[async_trait]
pub trait BatchPredict {
    /// Predict from batch of features
    async fn predict_batch(
        &self,
        features: &Array2<f64>,
    ) -> Result<Vec<InferenceResult>, crate::MLError>;
}

/// Trait for graph-based models
pub trait GraphModel {
    /// Add node to the model's internal graph
    fn add_node(&mut self, node_id: String, features: Vec<f64>) -> Result<(), crate::MLError>;

    /// Remove node from the model's internal graph
    fn remove_node(&mut self, node_id: &str) -> Result<(), crate::MLError>;

    /// Add edge between nodes
    fn add_edge(&mut self, from: &str, to: &str, weight: f64) -> Result<(), crate::MLError>;

    /// Get neighbors of a node
    fn get_neighbors(&self, node_id: &str) -> Option<Vec<String>>;

    /// Update graph structure
    fn update_graph(&mut self) -> Result<(), crate::MLError>;
}

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

    #[test]
    fn test_streaming_stats_default() {
        let stats = StreamingStats::default();
        assert_eq!(stats.total_processed, 0);
        assert_eq!(stats.min_latency_us, u64::MAX);
    }

    #[test]
    fn test_performance_metrics_targets() {
        let mut metrics = PerformanceMetrics::default();

        // Should not meet targets initially
        metrics.average_inference_latency_us = 0;
        metrics.throughput_pps = 0;

        // Mock a struct that implements ModelPerformance for testing
        struct MockModel {
            metrics: PerformanceMetrics,
        }

        impl ModelPerformance for MockModel {
            fn performance_metrics(&self) -> PerformanceMetrics {
                self.metrics.clone()
            }

            fn reset_performance_counters(&mut self) {
                self.metrics = PerformanceMetrics::default();
            }
        }

        let mock = MockModel { metrics };
        assert!(!mock.meets_performance_targets());

        // Update to meet targets
        let mut good_metrics = PerformanceMetrics::default();
        good_metrics.average_inference_latency_us = 50; // Under 100μs
        good_metrics.throughput_pps = 150_000; // Over 100K

        let good_mock = MockModel {
            metrics: good_metrics,
        };
        assert!(good_mock.meets_performance_targets());
    }
}