foxhunt/crates/ml/src/integration/model_registry.rs

//! # Model Registry
//!
//! Centralized registry for managing ML model deployments, versions,
//! and metadata in the Foxhunt HFT system.

use std::collections::HashMap;
use std::time::SystemTime;

use super::{ModelDeployment, ModelSearchCriteria, ModelState, ModelStatus};
#[cfg(test)]
use super::{ModelType, ServingMode};
use crate::MLError;

/// Model Registry for managing ML model deployments
#[derive(Debug)]
pub struct ModelRegistry {
    /// Active models with their status
    pub active_models: HashMap<String, ModelStatus>,
    /// Model deployments
    deployments: HashMap<String, ModelDeployment>,
}

impl ModelRegistry {
    /// Create a new model registry
    pub fn new() -> Self {
        Self {
            active_models: HashMap::new(),
            deployments: HashMap::new(),
        }
    }

    /// Register a model deployment
    pub async fn register_model(&mut self, deployment: ModelDeployment) -> Result<(), MLError> {
        let model_id = deployment.model_id.clone();

        // Create initial status
        let status = ModelStatus {
            model_id: model_id.clone(),
            status: ModelState::Loading,
            last_health_check: SystemTime::now(),
            deployment_time: SystemTime::now(),
            inference_count: 0,
            error_count: 0,
            avg_latency_us: 0.0,
            memory_usage_mb: 0.0,
            cpu_utilization: 0.0,
        };

        self.deployments.insert(model_id.clone(), deployment);
        self.active_models.insert(model_id, status);

        Ok(())
    }

    /// Get model status
    pub fn get_model_status(&self, model_id: &str) -> Option<&ModelStatus> {
        self.active_models.get(model_id)
    }

    /// List all active models
    pub fn list_active_models(&self) -> Vec<String> {
        self.active_models.keys().cloned().collect()
    }

    /// Search models by criteria
    pub fn search_models(&self, criteria: &ModelSearchCriteria) -> Vec<String> {
        self.deployments
            .iter()
            .filter(|(model_id, deployment)| {
                // Filter by model type
                if let Some(ref model_type) = criteria.model_type {
                    if &deployment.model_type != model_type {
                        return false;
                    }
                }

                // Filter by serving mode
                if let Some(ref serving_mode) = criteria.serving_mode {
                    if !deployment.serving_modes.contains(serving_mode) {
                        return false;
                    }
                }

                // Filter by max latency
                if let Some(max_latency) = criteria.max_latency_us {
                    if deployment.target_latency_us > max_latency {
                        return false;
                    }
                }

                // Filter by status
                if let Some(ref status) = criteria.status {
                    if let Some(model_status) = self.active_models.get(*model_id) {
                        if &model_status.status != status {
                            return false;
                        }
                    }
                }

                true
            })
            .map(|(model_id, _)| model_id.clone())
            .collect()
    }

    /// Calculate model score based on criteria
    pub fn calculate_model_score(&self, model_id: &str, criteria: &ModelSearchCriteria) -> f64 {
        if let Some(status) = self.active_models.get(model_id) {
            let mut score = 1.0;

            // Penalize high error rate
            if status.inference_count > 0 {
                let error_rate = status.error_count as f64 / status.inference_count as f64;
                score *= (1.0 - error_rate).max(0.0);
            }

            // Favor lower latency if criteria specifies max latency
            if let Some(max_latency) = criteria.max_latency_us {
                if status.avg_latency_us > 0.0 {
                    let latency_score =
                        (max_latency as f64 - status.avg_latency_us) / max_latency as f64;
                    score *= latency_score.max(0.0);
                }
            }

            // Favor lower resource usage
            score *= (1.0 - (status.cpu_utilization / 100.0)).max(0.0);

            score.min(1.0).max(0.0)
        } else {
            0.0
        }
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::File;
    use tempfile::tempdir;

    #[tokio::test]
    async fn test_model_registry_creation() {
        let registry = ModelRegistry::new();
        assert!(registry.list_active_models().is_empty());
    }

    #[tokio::test]
    async fn test_model_registration() -> Result<(), Box<dyn std::error::Error>> {
        let mut registry = ModelRegistry::new();

        // Create a temporary model file
        let temp_dir = tempdir()?;
        let model_path = temp_dir.path().join("test_model.onnx");
        File::create(&model_path)?;

        let deployment = ModelDeployment {
            model_id: "test_model".to_owned(),
            model_type: ModelType::CompactDQN,
            version: "1.0".to_owned(),
            serving_modes: vec![ServingMode::LowLatency],
            file_path: model_path.to_string_lossy().to_string(),
            target_latency_us: 1000,
            memory_requirement_mb: 100,
            compute_unit: "CPU".to_owned(),
            warm_up_samples: 10,
        };

        let result = registry.register_model(deployment).await;
        assert!(result.is_ok());

        let status = registry.get_model_status("test_model");
        assert!(status.is_some());
        if let Some(status) = status {
            assert_eq!(status.status, ModelState::Loading);
        }
        Ok(())
    }

    #[tokio::test]
    async fn test_model_search() -> Result<(), Box<dyn std::error::Error>> {
        let mut registry = ModelRegistry::new();

        // Create temporary model files
        let temp_dir = tempdir()?;
        let model1_path = temp_dir.path().join("model1.onnx");
        let model2_path = temp_dir.path().join("model2.onnx");
        File::create(&model1_path)?;
        File::create(&model2_path)?;

        // Register two models
        let deployment1 = ModelDeployment {
            model_id: "fast_model".to_owned(),
            model_type: ModelType::DistilledMicroNet,
            version: "1.0".to_owned(),
            serving_modes: vec![ServingMode::UltraLowLatency],
            file_path: model1_path.to_string_lossy().to_string(),
            target_latency_us: 50,
            memory_requirement_mb: 10,
            compute_unit: "CPU".to_owned(),
            warm_up_samples: 5,
        };

        let deployment2 = ModelDeployment {
            model_id: "accurate_model".to_owned(),
            model_type: ModelType::CompactDQN,
            version: "1.0".to_owned(),
            serving_modes: vec![ServingMode::LowLatency],
            file_path: model2_path.to_string_lossy().to_string(),
            target_latency_us: 1000,
            memory_requirement_mb: 100,
            compute_unit: "GPU".to_owned(),
            warm_up_samples: 20,
        };

        registry.register_model(deployment1).await?;
        registry.register_model(deployment2).await?;

        // Search for ultra-low latency models
        let criteria = ModelSearchCriteria {
            model_type: None,
            serving_mode: Some(ServingMode::UltraLowLatency),
            max_latency_us: Some(100),
            min_accuracy: None,
            tags: vec![],
            status: None,
        };

        let results = registry.search_models(&criteria);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0], "fast_model");
        Ok(())
    }

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

        // Add a model status
        let status = ModelStatus {
            model_id: "test_model".to_owned(),
            status: ModelState::Active,
            last_health_check: SystemTime::now(),
            deployment_time: SystemTime::now(),
            inference_count: 1000,
            error_count: 10, // 1% error rate
            avg_latency_us: 500.0,
            memory_usage_mb: 50.0,
            cpu_utilization: 30.0,
        };

        registry
            .active_models
            .insert("test_model".to_owned(), status);

        let criteria = ModelSearchCriteria {
            model_type: None,
            serving_mode: None,
            max_latency_us: Some(1000),
            min_accuracy: None,
            tags: vec![],
            status: None,
        };

        let score = registry.calculate_model_score("test_model", &criteria);
        assert!(score > 0.0);
        assert!(score <= 1.0);
    }
}