foxhunt/ml/src/observability/metrics.rs

//! Production observability and monitoring for ML inference pipeline
//!
//! This module provides comprehensive metrics collection, monitoring, and alerting
//! for ML models in production HFT environment. Critical for maintaining sub-50μs
//! latency targets and ensuring model reliability.

use anyhow::{Context, Result};
use prometheus::{
    CounterVec, Gauge, GaugeVec, HistogramOpts, HistogramVec, IntCounter, IntCounterVec,
    IntGaugeVec, Opts, Registry,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::RwLock;

use crate::{MLError, MLResult, ModelPrediction, ModelType};

/// Comprehensive metrics collection for ML operations
#[derive(Clone)]
pub struct MLMetricsCollector {
    registry: Arc<Registry>,

    // Latency metrics
    inference_latency: HistogramVec,
    prediction_latency: HistogramVec,
    model_load_latency: HistogramVec,

    // Throughput metrics
    predictions_total: CounterVec,
    inference_requests_total: IntCounterVec,
    successful_predictions: IntCounterVec,
    failed_predictions: IntCounterVec,

    // Model performance metrics
    model_confidence: GaugeVec,
    prediction_accuracy: GaugeVec,
    drift_detection_score: GaugeVec,

    // Resource utilization
    gpu_utilization: Gauge,
    cpu_utilization: Gauge,
    memory_usage_mb: Gauge,

    // Model health
    model_status: IntGaugeVec,
    last_prediction_time: GaugeVec,
    error_rate: GaugeVec,

    // Feature quality
    feature_quality_score: GaugeVec,
    missing_features_total: IntCounterVec,
    invalid_features_total: IntCounterVec,

    // Business metrics
    trading_pnl: Gauge,
    position_sizing_errors: IntCounter,
    risk_violations: IntCounterVec,
}

impl MLMetricsCollector {
    /// Create new metrics collector with Prometheus registry
    pub fn new() -> Result<Self> {
        let registry = Arc::new(Registry::new());

        // Latency histograms with HFT-appropriate buckets (microseconds)
        let latency_buckets = vec![
            1.0, 5.0, 10.0, 25.0, 50.0, 100.0, 250.0, 500.0, 1000.0, 2500.0, 5000.0,
        ];

        let inference_latency = HistogramVec::new(
            HistogramOpts::new(
                "ml_inference_latency_microseconds",
                "ML inference latency in microseconds",
            )
            .buckets(latency_buckets.clone()),
            &["model_type", "model_name", "symbol"],
        )?;

        let prediction_latency = HistogramVec::new(
            HistogramOpts::new(
                "ml_prediction_latency_microseconds",
                "ML prediction processing latency in microseconds",
            )
            .buckets(latency_buckets.clone()),
            &["model_type", "operation"],
        )?;

        let model_load_latency = HistogramVec::new(
            HistogramOpts::new(
                "ml_model_load_latency_seconds",
                "Model loading latency in seconds",
            )
            .buckets(vec![0.1, 0.5, 1.0, 2.0, 5.0, 10.0, 30.0]),
            &["model_type", "model_name"],
        )?;

        // Counters
        let predictions_total = CounterVec::new(
            Opts::new("ml_predictions_total", "Total ML predictions made"),
            &["model_type", "model_name", "result"],
        )?;

        let inference_requests_total = IntCounterVec::new(
            Opts::new("ml_inference_requests_total", "Total inference requests"),
            &["model_type", "model_name", "symbol"],
        )?;

        let successful_predictions = IntCounterVec::new(
            Opts::new("ml_successful_predictions_total", "Successful predictions"),
            &["model_type", "model_name"],
        )?;

        let failed_predictions = IntCounterVec::new(
            Opts::new("ml_failed_predictions_total", "Failed predictions"),
            &["model_type", "model_name", "error_type"],
        )?;

        // Gauges
        let model_confidence = GaugeVec::new(
            Opts::new(
                "ml_model_confidence",
                "Current model confidence score (0-1)",
            ),
            &["model_type", "model_name"],
        )?;

        let prediction_accuracy = GaugeVec::new(
            Opts::new(
                "ml_prediction_accuracy",
                "Model prediction accuracy over time window",
            ),
            &["model_type", "model_name", "time_window"],
        )?;

        let drift_detection_score = GaugeVec::new(
            Opts::new("ml_drift_detection_score", "Model drift detection score"),
            &["model_type", "model_name", "feature_group"],
        )?;

        let gpu_utilization =
            Gauge::new("ml_gpu_utilization_percent", "GPU utilization percentage")?;

        let cpu_utilization =
            Gauge::new("ml_cpu_utilization_percent", "CPU utilization percentage")?;

        let memory_usage_mb = Gauge::new("ml_memory_usage_megabytes", "Memory usage in megabytes")?;

        let model_status = IntGaugeVec::new(
            Opts::new("ml_model_status", "Model status (1=healthy, 0=unhealthy)"),
            &["model_type", "model_name"],
        )?;

        let last_prediction_time = GaugeVec::new(
            Opts::new(
                "ml_last_prediction_timestamp",
                "Timestamp of last prediction",
            ),
            &["model_type", "model_name"],
        )?;

        let error_rate = GaugeVec::new(
            Opts::new("ml_error_rate", "Error rate over time window"),
            &["model_type", "model_name", "time_window"],
        )?;

        let feature_quality_score = GaugeVec::new(
            Opts::new("ml_feature_quality_score", "Feature quality score (0-1)"),
            &["feature_group", "symbol"],
        )?;

        let missing_features_total = IntCounterVec::new(
            Opts::new(
                "ml_missing_features_total",
                "Total missing features detected",
            ),
            &["feature_name", "symbol"],
        )?;

        let invalid_features_total = IntCounterVec::new(
            Opts::new(
                "ml_invalid_features_total",
                "Total invalid features detected",
            ),
            &["feature_name", "validation_rule", "symbol"],
        )?;

        // Business metrics
        let trading_pnl = Gauge::new(
            "ml_trading_pnl_dollars",
            "Current trading P&L from ML predictions",
        )?;

        let position_sizing_errors = IntCounter::new(
            "ml_position_sizing_errors_total",
            "Position sizing errors from ML models",
        )?;

        let risk_violations = IntCounterVec::new(
            Opts::new("ml_risk_violations_total", "Risk management violations"),
            &["violation_type", "model_name"],
        )?;

        // Register all metrics
        registry.register(Box::new(inference_latency.clone()))?;
        registry.register(Box::new(prediction_latency.clone()))?;
        registry.register(Box::new(model_load_latency.clone()))?;
        registry.register(Box::new(predictions_total.clone()))?;
        registry.register(Box::new(inference_requests_total.clone()))?;
        registry.register(Box::new(successful_predictions.clone()))?;
        registry.register(Box::new(failed_predictions.clone()))?;
        registry.register(Box::new(model_confidence.clone()))?;
        registry.register(Box::new(prediction_accuracy.clone()))?;
        registry.register(Box::new(drift_detection_score.clone()))?;
        registry.register(Box::new(gpu_utilization.clone()))?;
        registry.register(Box::new(cpu_utilization.clone()))?;
        registry.register(Box::new(memory_usage_mb.clone()))?;
        registry.register(Box::new(model_status.clone()))?;
        registry.register(Box::new(last_prediction_time.clone()))?;
        registry.register(Box::new(error_rate.clone()))?;
        registry.register(Box::new(feature_quality_score.clone()))?;
        registry.register(Box::new(missing_features_total.clone()))?;
        registry.register(Box::new(invalid_features_total.clone()))?;
        registry.register(Box::new(trading_pnl.clone()))?;
        registry.register(Box::new(position_sizing_errors.clone()))?;
        registry.register(Box::new(risk_violations.clone()))?;

        Ok(Self {
            registry,
            inference_latency,
            prediction_latency,
            model_load_latency,
            predictions_total,
            inference_requests_total,
            successful_predictions,
            failed_predictions,
            model_confidence,
            prediction_accuracy,
            drift_detection_score,
            gpu_utilization,
            cpu_utilization,
            memory_usage_mb,
            model_status,
            last_prediction_time,
            error_rate,
            feature_quality_score,
            missing_features_total,
            invalid_features_total,
            trading_pnl,
            position_sizing_errors,
            risk_violations,
        })
    }

    /// Record inference latency
    pub fn record_inference_latency(
        &self,
        model_type: ModelType,
        model_name: &str,
        symbol: Option<&str>,
        latency_us: f64,
    ) {
        let labels = [
            model_type.to_string(),
            model_name.to_string(),
            symbol.unwrap_or("unknown").to_string(),
        ];
        let label_refs: Vec<&str> = labels.iter().map(|s| s.as_str()).collect();
        self.inference_latency
            .with_label_values(&label_refs)
            .observe(latency_us);
    }

    /// Record successful prediction
    pub fn record_successful_prediction(
        &self,
        model_type: ModelType,
        model_name: &str,
        prediction: &ModelPrediction,
        latency_us: f64,
    ) {
        // Update counters
        let labels1 = [
            model_type.to_string(),
            model_name.to_string(),
            "success".to_string(),
        ];
        let label_refs1: Vec<&str> = labels1.iter().map(|s| s.as_str()).collect();
        self.predictions_total.with_label_values(&label_refs1).inc();

        let labels2 = [model_type.to_string(), model_name.to_string()];
        let label_refs2: Vec<&str> = labels2.iter().map(|s| s.as_str()).collect();
        self.successful_predictions
            .with_label_values(&label_refs2)
            .inc();
        // Update latency
        self.record_inference_latency(model_type, model_name, None, latency_us);

        // Update confidence
        let labels3 = [model_type.to_string(), model_name.to_string()];
        let label_refs3: Vec<&str> = labels3.iter().map(|s| s.as_str()).collect();
        self.model_confidence
            .with_label_values(&label_refs3)
            .set(prediction.confidence);

        // Update last prediction time
        let timestamp = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap_or_default()
            .as_secs() as f64;

        self.last_prediction_time
            .with_label_values(&label_refs3)
            .set(timestamp);
    }

    /// Record failed prediction
    pub fn record_failed_prediction(
        &self,
        model_type: ModelType,
        model_name: &str,
        error: &MLError,
    ) {
        let error_type = match error {
            MLError::ValidationError { .. } => "validation",
            MLError::InferenceError(..) => "inference",
            MLError::ModelError(..) => "model",
            MLError::ConfigError { .. } => "config",
            MLError::TensorCreationError { .. } => "tensor",
            _ => "other",
        };

        let labels = [
            model_type.to_string(),
            model_name.to_string(),
            "failure".to_string(),
        ];
        let label_refs: Vec<&str> = labels.iter().map(|s| s.as_str()).collect();
        self.predictions_total.with_label_values(&label_refs).inc();

        let labels2 = [
            model_type.to_string(),
            model_name.to_string(),
            error_type.to_string(),
        ];
        let label_refs2: Vec<&str> = labels2.iter().map(|s| s.as_str()).collect();
        self.failed_predictions
            .with_label_values(&label_refs2)
            .inc();
    }

    /// Update model health status
    pub fn update_model_status(&self, model_type: ModelType, model_name: &str, is_healthy: bool) {
        let labels = [model_type.to_string(), model_name.to_string()];
        let label_refs: Vec<&str> = labels.iter().map(|s| s.as_str()).collect();
        self.model_status
            .with_label_values(&label_refs)
            .set(if is_healthy { 1 } else { 0 });
    }

    /// Record resource utilization
    pub fn record_resource_utilization(&self, gpu_percent: f64, cpu_percent: f64, memory_mb: f64) {
        self.gpu_utilization.set(gpu_percent);
        self.cpu_utilization.set(cpu_percent);
        self.memory_usage_mb.set(memory_mb);
    }

    /// Record drift detection score
    pub fn record_drift_score(
        &self,
        model_type: ModelType,
        model_name: &str,
        feature_group: &str,
        score: f64,
    ) {
        let labels = [
            model_type.to_string(),
            model_name.to_string(),
            feature_group.to_string(),
        ];
        let label_refs: Vec<&str> = labels.iter().map(|s| s.as_str()).collect();
        self.drift_detection_score
            .with_label_values(&label_refs)
            .set(score);
    }

    /// Record feature quality metrics
    pub fn record_feature_quality(&self, feature_group: &str, symbol: &str, quality_score: f64) {
        self.feature_quality_score
            .with_label_values(&[feature_group, symbol])
            .set(quality_score);
    }

    /// Record missing feature
    pub fn record_missing_feature(&self, feature_name: &str, symbol: &str) {
        self.missing_features_total
            .with_label_values(&[feature_name, symbol])
            .inc();
    }

    /// Record invalid feature
    pub fn record_invalid_feature(&self, feature_name: &str, validation_rule: &str, symbol: &str) {
        self.invalid_features_total
            .with_label_values(&[feature_name, validation_rule, symbol])
            .inc();
    }

    /// Update trading P&L
    pub fn update_trading_pnl(&self, pnl_dollars: f64) {
        self.trading_pnl.set(pnl_dollars);
    }

    /// Record position sizing error
    pub fn record_position_sizing_error(&self) {
        self.position_sizing_errors.inc();
    }

    /// Record risk violation
    pub fn record_risk_violation(&self, violation_type: &str, model_name: &str) {
        self.risk_violations
            .with_label_values(&[violation_type, model_name])
            .inc();
    }

    /// Get Prometheus metrics registry for HTTP exposure
    pub fn get_registry(&self) -> Arc<Registry> {
        self.registry.clone()
    }

    /// Generate metrics report
    pub async fn generate_report(&self) -> MLMetricsReport {
        // Simplified metrics report to avoid protobuf complexity
        let summary = HashMap::new();

        MLMetricsReport {
            timestamp: SystemTime::now(),
            summary,
            total_predictions: self.calculate_total_predictions(),
            average_latency_us: self.calculate_average_latency(),
            error_rate: self.calculate_error_rate(),
            health_score: self.calculate_health_score(),
        }
    }

    fn calculate_total_predictions(&self) -> u64 {
        // This would sum all prediction counters - simplified for now
        0
    }

    fn calculate_average_latency(&self) -> f64 {
        // This would calculate weighted average from histograms - simplified for now
        0.0
    }

    fn calculate_error_rate(&self) -> f64 {
        // This would calculate error rate from counters - simplified for now
        0.0
    }

    fn calculate_health_score(&self) -> f64 {
        // This would calculate overall system health - simplified for now
        1.0
    }
}

impl Default for MLMetricsCollector {
    fn default() -> Self {
        Self::new().expect("Failed to create metrics collector")
    }
}

/// Comprehensive metrics report
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MLMetricsReport {
    pub timestamp: SystemTime,
    pub summary: HashMap<String, Vec<f64>>,
    pub total_predictions: u64,
    pub average_latency_us: f64,
    pub error_rate: f64,
    pub health_score: f64,
}

/// Model type to string conversion for metrics
impl ToString for ModelType {
    fn to_string(&self) -> String {
        match self {
            ModelType::DQN => "dqn".to_string(),
            ModelType::MAMBA | ModelType::Mamba => "mamba".to_string(),
            ModelType::TFT => "tft".to_string(),
            ModelType::TGGN | ModelType::TGNN => "tgnn".to_string(),
            ModelType::LNN | ModelType::LiquidNet => "liquid".to_string(),
            ModelType::CompactDQN => "compact_dqn".to_string(),
            ModelType::DistilledMicroNet => "distilled".to_string(),
            ModelType::RainbowDQN => "rainbow_dqn".to_string(),
            ModelType::TLOB => "tlob".to_string(),
            ModelType::PPO => "ppo".to_string(),
            ModelType::Transformer => "transformer".to_string(),
            ModelType::Ensemble => "ensemble".to_string(),
        }
    }
}

/// Global metrics collector instance
static GLOBAL_METRICS: once_cell::sync::Lazy<Arc<RwLock<Option<MLMetricsCollector>>>> =
    once_cell::sync::Lazy::new(|| Arc::new(RwLock::new(None)));

/// Initialize global metrics collector
pub async fn initialize_metrics() -> Result<()> {
    let collector = MLMetricsCollector::new().context("Failed to create metrics collector")?;

    let mut global = GLOBAL_METRICS.write().await;
    *global = Some(collector);

    tracing::info!("ML metrics collector initialized");
    Ok(())
}

/// Get global metrics collector
pub async fn get_metrics_collector() -> Option<MLMetricsCollector> {
    let global = GLOBAL_METRICS.read().await;
    global.clone()
}

/// Record inference timing with automatic metrics collection
pub async fn record_inference_timing<F, T>(
    model_type: ModelType,
    model_name: &str,
    symbol: Option<&str>,
    operation: F,
) -> MLResult<T>
where
    F: std::future::Future<Output = MLResult<T>>,
{
    let start = Instant::now();
    let result = operation.await;
    let latency_us = start.elapsed().as_micros() as f64;

    if let Some(collector) = get_metrics_collector().await {
        match &result {
            Ok(_) => {
                collector.record_inference_latency(model_type, model_name, symbol, latency_us);
            }
            Err(error) => {
                collector.record_failed_prediction(model_type, model_name, error);
            }
        }
    }

    result
}

/// Performance monitoring wrapper for ML operations
pub struct MLPerformanceMonitor {
    collector: MLMetricsCollector,
    alert_thresholds: AlertThresholds,
}

#[derive(Debug, Clone)]
pub struct AlertThresholds {
    pub max_latency_us: f64,
    pub min_confidence: f64,
    pub max_error_rate: f64,
    pub min_health_score: f64,
}

impl Default for AlertThresholds {
    fn default() -> Self {
        Self {
            max_latency_us: 100.0, // 100μs max latency
            min_confidence: 0.7,   // 70% minimum confidence
            max_error_rate: 0.05,  // 5% maximum error rate
            min_health_score: 0.8, // 80% minimum health score
        }
    }
}

impl MLPerformanceMonitor {
    pub fn new(collector: MLMetricsCollector) -> Self {
        Self {
            collector,
            alert_thresholds: AlertThresholds::default(),
        }
    }

    pub fn with_thresholds(mut self, thresholds: AlertThresholds) -> Self {
        self.alert_thresholds = thresholds;
        self
    }

    /// Check if system meets performance requirements
    pub async fn check_performance_health(&self) -> PerformanceHealthCheck {
        let report = self.collector.generate_report().await;

        PerformanceHealthCheck {
            latency_ok: report.average_latency_us <= self.alert_thresholds.max_latency_us,
            error_rate_ok: report.error_rate <= self.alert_thresholds.max_error_rate,
            health_score_ok: report.health_score >= self.alert_thresholds.min_health_score,
            overall_healthy: report.average_latency_us <= self.alert_thresholds.max_latency_us
                && report.error_rate <= self.alert_thresholds.max_error_rate
                && report.health_score >= self.alert_thresholds.min_health_score,
            current_metrics: report,
        }
    }
}

#[derive(Debug, Clone)]
pub struct PerformanceHealthCheck {
    pub latency_ok: bool,
    pub error_rate_ok: bool,
    pub health_score_ok: bool,
    pub overall_healthy: bool,
    pub current_metrics: MLMetricsReport,
}

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

    #[tokio::test]
    async fn test_metrics_collector_creation() {
        let collector = MLMetricsCollector::new();
        assert!(collector.is_ok());
    }

    #[tokio::test]
    async fn test_global_metrics_initialization() {
        let result = initialize_metrics().await;
        assert!(result.is_ok());

        let collector = get_metrics_collector().await;
        assert!(collector.is_some());
    }

    #[test]
    fn test_model_type_string_conversion() {
        assert_eq!(ModelType::DQN.to_string(), "dqn");
        assert_eq!(ModelType::MAMBA.to_string(), "mamba");
        assert_eq!(ModelType::TLOB.to_string(), "tlob");
    }

    #[tokio::test]
    async fn test_performance_monitor() {
        let collector = MLMetricsCollector::new().unwrap();
        let monitor = MLPerformanceMonitor::new(collector);

        let health_check = monitor.check_performance_health().await;
        assert!(health_check.overall_healthy); // Should be healthy with no data
    }
}