//! Enhanced ML service implementation with hot-loading, ensemble coordination, and production features
//!
//! This module implements production ML model integration with:
//! - Real model loading from filesystem/S3
//! - Feature preprocessing and normalization
//! - Real inference using ml crate models
//! - Production metrics and monitoring

use super::ml_fallback_manager::ModelHealth as FallbackModelHealth;
use crate::proto::ml::{
    ml_service_server::MlService, EnsembleVote, Feature, FeatureType,
    GetAvailableModelsRequest, GetAvailableModelsResponse, GetEnsembleVoteRequest,
    GetEnsembleVoteResponse, GetFeatureImportanceRequest, GetFeatureImportanceResponse,
    GetModelPerformanceRequest, GetModelPerformanceResponse, GetModelStatusRequest,
    GetModelStatusResponse, GetPredictionRequest, GetPredictionResponse, ModelCapabilities,
    ModelHealth, ModelInfo, ModelMetrics, ModelMetricsEvent, ModelPerformance, ModelSignal,
    ModelState, ModelStatus, ModelVote, Prediction, PredictionEvent, PredictionType,
    RetrainModelRequest, RetrainModelResponse, SignalStrengthEvent, StreamModelMetricsRequest,
    StreamModelStatusRequest, StreamPredictionsRequest, StreamSignalStrengthRequest,
};
use crate::state::TradingServiceState;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use tokio::sync::{broadcast, OnceCell, RwLock};
use tokio_stream::{wrappers::BroadcastStream, wrappers::ReceiverStream, Stream, StreamExt};
use tonic::transport::Channel;
use tonic::{Request, Response, Status};
use tracing::{debug, info, warn, instrument};

// Production ML imports
use ml::{Features, MLModel, ModelType};
use sysinfo::System;

/// Runtime model tracking metadata with actual ML model instance
#[derive(Debug, Clone)]
pub(crate) struct RuntimeModelInfo {
    /// Stored for Debug output and potential future gRPC response use
    #[allow(dead_code)]
    pub(crate) model_id: String,
    pub(crate) version: String,
    pub(crate) load_time: SystemTime,
    pub(crate) last_inference: Option<SystemTime>,
    pub(crate) inference_count: u64,
    pub(crate) error_count: u64,
    pub(crate) avg_latency_us: f64,
    pub(crate) confidence_threshold: f64,
    pub(crate) weight_in_ensemble: f64,
    /// Stored for fallback ordering (used when fallback manager selects models)
    #[allow(dead_code)]
    pub(crate) fallback_priority: i32,
    pub(crate) model_type: ModelType,
    pub(crate) supported_symbols: Vec<String>,
    pub(crate) supported_horizons: Vec<i32>,
    pub(crate) feature_count: usize,
    pub(crate) model_instance: Option<Arc<dyn MLModel>>,
}

/// Feature normalization statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct FeatureNormStats {
    pub(crate) mean: f64,
    pub(crate) std_dev: f64,
    pub(crate) min: f64,
    pub(crate) max: f64,
}

/// Feature preprocessor for ML models
#[derive(Debug, Clone)]
pub(crate) struct FeaturePreprocessor {
    pub(crate) stats: HashMap<String, FeatureNormStats>,
}

impl Default for FeaturePreprocessor {
    fn default() -> Self {
        let mut stats = HashMap::new();

        // Default normalization parameters for common features
        stats.insert(
            "price_momentum".to_string(),
            FeatureNormStats {
                mean: 0.0,
                std_dev: 0.1,
                min: -1.0,
                max: 1.0,
            },
        );

        stats.insert(
            "volume".to_string(),
            FeatureNormStats {
                mean: 1000000.0,
                std_dev: 500000.0,
                min: 0.0,
                max: 10000000.0,
            },
        );

        stats.insert(
            "volatility".to_string(),
            FeatureNormStats {
                mean: 0.02,
                std_dev: 0.01,
                min: 0.0,
                max: 0.5,
            },
        );

        Self { stats }
    }
}

impl FeaturePreprocessor {
    /// Create new preprocessor with default normalization stats
    pub fn new() -> Self {
        Self::default()
    }

    /// Normalize a feature value using z-score normalization
    pub fn normalize(&self, feature_name: &str, value: f64) -> f64 {
        if !value.is_finite() {
            debug!(feature = %feature_name, value = %value, "NaN/Inf feature detected, replacing with 0.0");
            return 0.0;
        }

        if let Some(stats) = self.stats.get(feature_name) {
            if stats.std_dev > 0.0 {
                let normalized = (value - stats.mean) / stats.std_dev;
                normalized.clamp(-10.0, 10.0)
            } else {
                value
            }
        } else {
            // Default normalization for unknown features
            value.tanh()
        }
    }

    /// Classify feature type based on name
    pub fn classify_feature_type(&self, feature_name: &str) -> FeatureType {
        match feature_name {
            n if n.contains("price") || n.contains("momentum") => FeatureType::Price,
            n if n.contains("volume") => FeatureType::Volume,
            n if n.contains("volatility") || n.contains("rsi") || n.contains("ma") => {
                FeatureType::Technical
            },
            n if n.contains("sentiment") || n.contains("news") => FeatureType::Sentiment,
            n if n.contains("orderbook") || n.contains("depth") => FeatureType::Volume, // Orderbook is volume-related
            n if n.contains("spread") || n.contains("liquidity") => FeatureType::Technical, // Microstructure metrics
            _ => FeatureType::Technical,
        }
    }
}

/// Ensemble configuration (service-layer; distinct from ml::ensemble::EnsembleConfig)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct EnsembleConfig {
    pub(crate) min_models: usize,
    pub(crate) confidence_threshold: f64,
    pub(crate) use_weighted_voting: bool,
    pub(crate) fallback_timeout_ms: u64,
    pub(crate) consensus_threshold: f64,
}

impl Default for EnsembleConfig {
    fn default() -> Self {
        Self {
            min_models: 2,
            confidence_threshold: 0.7,
            use_weighted_voting: true,
            fallback_timeout_ms: 50,
            consensus_threshold: 0.6,
        }
    }
}

/// Enhanced ML service implementation with production features
#[derive(Debug, Clone)]
pub struct MLService {
    state: TradingServiceState,
    models: Arc<RwLock<HashMap<String, RuntimeModelInfo>>>,
    ensemble_config: Arc<RwLock<EnsembleConfig>>,
    prediction_broadcaster: Arc<broadcast::Sender<PredictionEvent>>,
    model_weights: Arc<RwLock<HashMap<String, f64>>>,
    // Production ML monitoring and fallback
    ml_performance_monitor: Arc<super::ml_performance_monitor::MLPerformanceMonitor>,
    ml_fallback_manager: Arc<super::ml_fallback_manager::MLFallbackManager>,
    // Feature preprocessing
    feature_preprocessor: Arc<FeaturePreprocessor>,
    // System metrics tracking
    system: Arc<RwLock<System>>,
    // gRPC client for forwarding retrain requests to ml_training_service
    ml_training_client: Arc<OnceCell<crate::proto::ml_training::ml_training_service_client::MlTrainingServiceClient<Channel>>>,
}

impl MLService {
    /// Create new enhanced ML service with production monitoring and fallback
    pub fn new(
        state: TradingServiceState,
        ml_performance_monitor: Arc<super::ml_performance_monitor::MLPerformanceMonitor>,
        ml_fallback_manager: Arc<super::ml_fallback_manager::MLFallbackManager>,
    ) -> Self {
        let (prediction_sender, _) = broadcast::channel(1000);

        Self {
            state,
            models: Arc::new(RwLock::new(HashMap::new())),
            ensemble_config: Arc::new(RwLock::new(EnsembleConfig::default())),
            prediction_broadcaster: Arc::new(prediction_sender),
            model_weights: Arc::new(RwLock::new(HashMap::new())),
            ml_performance_monitor,
            ml_fallback_manager,
            feature_preprocessor: Arc::new(FeaturePreprocessor::new()),
            system: Arc::new(RwLock::new(System::new_all())),
            ml_training_client: Arc::new(OnceCell::new()),
        }
    }

    /// Lazily connect to ml_training_service. Endpoint from ML_TRAINING_SERVICE_URL env var.
    async fn get_training_client(
        &self,
    ) -> Result<crate::proto::ml_training::ml_training_service_client::MlTrainingServiceClient<Channel>, Status> {
        self.ml_training_client
            .get_or_try_init(|| async {
                let url = std::env::var("ML_TRAINING_SERVICE_URL")
                    .unwrap_or_else(|_| "http://ml-training-service:50052".to_string());
                let channel = Channel::from_shared(url)
                    .map_err(|e| Status::internal(format!("invalid training service URL: {e}")))?
                    .connect()
                    .await
                    .map_err(|e| Status::unavailable(format!("ml_training_service unreachable: {e}")))?;
                Ok(crate::proto::ml_training::ml_training_service_client::MlTrainingServiceClient::new(channel))
            })
            .await
            .cloned()
    }

    /// Load model from checkpoint file (production implementation)
    pub async fn load_model_from_file(
        &self,
        model_id: &str,
        model_path: &str,
    ) -> Result<Arc<dyn MLModel>, Status> {
        info!("Loading model {} from path: {}", model_id, model_path);

        let checkpoint_path = std::path::Path::new(model_path);

        // Verify checkpoint exists
        if !checkpoint_path.exists() {
            return Err(Status::not_found(format!(
                "Checkpoint not found: {}",
                model_path
            )));
        }

        // Determine model type from model_id
        let model_type_str = if model_id.contains("DQN") || model_id.contains("dqn") {
            "DQN"
        } else if model_id.contains("PPO") || model_id.contains("ppo") {
            "PPO"
        } else if model_id.contains("TFT") || model_id.contains("tft") {
            "TFT"
        } else if model_id.contains("MAMBA") || model_id.contains("mamba") {
            "MAMBA2"
        } else if model_id.contains("Liquid")
            || model_id.contains("liquid")
            || model_id.contains("CfC")
            || model_id.contains("cfc")
            || model_id.contains("LNN")
            || model_id.contains("lnn")
        {
            "LIQUID"
        } else {
            return Err(Status::invalid_argument(format!(
                "Unknown model type in model_id: {}. Supported: DQN, PPO, TFT, MAMBA2, CFC/Liquid",
                model_id
            )));
        };

        // Load model based on type
        let model: Arc<dyn MLModel> = match model_type_str {
            "DQN" => {
                let dqn_model =
                    DQNModel::from_checkpoint(model_id.to_string(), checkpoint_path).map_err(
                        |e| Status::internal(format!("Failed to load DQN model: {}", e)),
                    )?;

                Arc::new(dqn_model) as Arc<dyn MLModel>
            },

            "PPO" => {
                // PPO has separate actor/critic checkpoints
                // Parse checkpoint paths from directory structure
                let checkpoint_dir = checkpoint_path
                    .parent()
                    .ok_or_else(|| Status::invalid_argument("Invalid PPO checkpoint path"))?;

                // Extract epoch number from model_id (e.g., "PPO_epoch130")
                let epoch_num = model_id
                    .split('_')
                    .find(|s| s.starts_with("epoch"))
                    .and_then(|s| s.strip_prefix("epoch"))
                    .and_then(|s| s.parse::<u32>().ok())
                    .unwrap_or(130); // Default to epoch 130

                let actor_path =
                    checkpoint_dir.join(format!("ppo_actor_epoch_{}.safetensors", epoch_num));
                let critic_path =
                    checkpoint_dir.join(format!("ppo_critic_epoch_{}.safetensors", epoch_num));

                if !actor_path.exists() || !critic_path.exists() {
                    return Err(Status::not_found(format!(
                        "PPO checkpoints not found: actor={}, critic={}",
                        actor_path.display(),
                        critic_path.display()
                    )));
                }

                let ppo_model =
                    PPOModel::from_checkpoint(model_id.to_string(), &actor_path, &critic_path)
                        .map_err(|e| {
                            Status::internal(format!("Failed to load PPO model: {}", e))
                        })?;

                Arc::new(ppo_model) as Arc<dyn MLModel>
            },

            "TFT" => {
                // TFT (Temporal Fusion Transformer) checkpoint loading
                // Checkpoint path is direct to safetensors file (e.g., tft_epoch_100.safetensors)
                let tft_model =
                    TFTModel::from_checkpoint(model_id.to_string(), checkpoint_path).map_err(
                        |e| Status::internal(format!("Failed to load TFT model: {}", e)),
                    )?;

                Arc::new(tft_model) as Arc<dyn MLModel>
            },

            "MAMBA2" | "MAMBA" => {
                let mamba2_model =
                    Mamba2Model::from_checkpoint(model_id.to_string(), checkpoint_path)
                        .map_err(|e| {
                            Status::internal(format!("Failed to load Mamba2 model: {}", e))
                        })?;

                Arc::new(mamba2_model) as Arc<dyn MLModel>
            },

            "LIQUID" => {
                let liquid_model =
                    LiquidModel::from_checkpoint(model_id.to_string(), checkpoint_path)
                        .map_err(|e| {
                            Status::internal(format!("Failed to load Liquid CfC model: {}", e))
                        })?;

                Arc::new(liquid_model) as Arc<dyn MLModel>
            },

            _ => {
                return Err(Status::invalid_argument(format!(
                    "Unknown model type '{}'. Supported types: DQN, PPO, TFT, MAMBA2, LIQUID",
                    model_type_str
                )));
            },
        };

        let file_size = std::fs::metadata(checkpoint_path)
            .map(|m| m.len())
            .unwrap_or(0);

        info!(
            "Successfully loaded model {} ({} bytes, type: {})",
            model_id, file_size, model_type_str
        );

        Ok(model)
    }

    /// Get actual memory usage for current process
    async fn get_memory_usage_mb(&self) -> Result<f64, String> {
        let system = std::sync::Arc::clone(&self.system);

        // Try to get memory usage, fallback to 0 if unavailable
        let sys = system.read().await;
        if let Ok(pid) = sysinfo::get_current_pid() {
            if let Some(process) = sys.process(pid) {
                return Ok(process.memory() as f64 / 1024.0 / 1024.0);
            }
        }

        Ok(0.0)
    }

    /// Get actual CPU utilization for current process
    async fn get_cpu_utilization(&self) -> Result<f64, String> {
        let system = std::sync::Arc::clone(&self.system);

        // Try to get CPU usage, fallback to 0 if unavailable
        let mut sys = system.write().await;
        if let Ok(pid) = sysinfo::get_current_pid() {
            // Refresh all processes to get CPU usage
            sys.refresh_processes(sysinfo::ProcessesToUpdate::All, false);
            if let Some(process) = sys.process(pid) {
                return Ok(process.cpu_usage() as f64);
            }
        }

        Ok(0.0)
    }

    /// Get ensemble prediction from multiple models
    async fn get_ensemble_prediction(
        &self,
        features: &[f32],
        symbol: &str,
    ) -> Result<EnsembleVote, Status> {
        let models = self.models.read().await;
        let weights = self.model_weights.read().await;
        let config = self.ensemble_config.read().await;

        if models.len() < config.min_models {
            return Err(Status::failed_precondition(format!(
                "Insufficient models: {} < {}",
                models.len(),
                config.min_models
            )));
        }

        let mut individual_votes = Vec::new();
        let mut buy_votes = 0;
        let mut sell_votes = 0;
        let mut hold_votes = 0;
        let mut total_confidence = 0.0;
        let mut valid_predictions = 0;

        // Collect predictions from all models
        for (model_id, _) in models.iter() {
            match self
                .get_single_model_prediction(model_id, features, symbol)
                .await
            {
                Ok(prediction) => {
                    let weight = weights.get(model_id).copied().unwrap_or(1.0);

                    // Convert prediction to vote
                    let vote_type = if prediction.value > 0.6 {
                        buy_votes += 1;
                        crate::proto::ml::PredictionType::Buy
                    } else if prediction.value < 0.4 {
                        sell_votes += 1;
                        crate::proto::ml::PredictionType::Sell
                    } else {
                        hold_votes += 1;
                        crate::proto::ml::PredictionType::Hold
                    };

                    individual_votes.push(ModelVote {
                        model_name: model_id.clone(),
                        prediction: vote_type as i32,
                        confidence: prediction.confidence,
                        weight,
                    });

                    total_confidence += prediction.confidence * weight;
                    valid_predictions += 1;
                },
                Err(e) => {
                    warn!("Model {} failed prediction: {}", model_id, e);
                    self.record_model_error(model_id).await;
                },
            }
        }

        if valid_predictions == 0 {
            return Err(Status::internal("No models provided valid predictions"));
        }

        // Determine consensus
        let consensus_prediction = if buy_votes > sell_votes && buy_votes > hold_votes {
            crate::proto::ml::PredictionType::Buy
        } else if sell_votes > buy_votes && sell_votes > hold_votes {
            crate::proto::ml::PredictionType::Sell
        } else {
            crate::proto::ml::PredictionType::Hold
        };

        let consensus_confidence = total_confidence / valid_predictions as f64;

        // Enforce confidence threshold: low-confidence signals must not become trades
        if consensus_confidence < config.confidence_threshold {
            info!(
                "Ensemble confidence {:.3} below threshold {:.3}, forcing Hold for {}",
                consensus_confidence, config.confidence_threshold, symbol
            );
            return Ok(EnsembleVote {
                symbol: symbol.to_string(),
                consensus_prediction: crate::proto::ml::PredictionType::Hold as i32,
                consensus_confidence,
                votes_buy: buy_votes,
                votes_sell: sell_votes,
                votes_hold: hold_votes,
                total_models: valid_predictions,
                signal_strength: crate::proto::ml::SignalStrength::VeryWeak as i32,
            });
        }

        // Calculate signal strength
        let max_votes = buy_votes.max(sell_votes).max(hold_votes);
        let signal_strength = if valid_predictions > 0 {
            match max_votes as f64 / valid_predictions as f64 {
                ratio if ratio >= 0.8 => crate::proto::ml::SignalStrength::VeryStrong,
                ratio if ratio >= 0.6 => crate::proto::ml::SignalStrength::Strong,
                ratio if ratio >= 0.4 => crate::proto::ml::SignalStrength::Moderate,
                ratio if ratio >= 0.3 => crate::proto::ml::SignalStrength::Weak,
                _ => crate::proto::ml::SignalStrength::VeryWeak,
            }
        } else {
            crate::proto::ml::SignalStrength::VeryWeak
        };

        Ok(EnsembleVote {
            symbol: symbol.to_string(),
            consensus_prediction: consensus_prediction as i32,
            consensus_confidence,
            votes_buy: buy_votes,
            votes_sell: sell_votes,
            votes_hold: hold_votes,
            total_models: valid_predictions,
            signal_strength: signal_strength as i32,
        })
    }

    /// Get prediction from a single model (production implementation with real ML inference)
    async fn get_single_model_prediction(
        &self,
        model_id: &str,
        features: &[f32],
        symbol: &str,
    ) -> Result<Prediction, Status> {
        let start_time = Instant::now();

        // Get model metadata and instance
        let models = self.models.read().await;
        let model_meta = models
            .get(model_id)
            .ok_or_else(|| Status::not_found(format!("Model not found: {}", model_id)))?;

        let model_instance = model_meta
            .model_instance
            .as_ref()
            .ok_or_else(|| Status::internal(format!("Model instance not loaded: {}", model_id)))?;

        // Get supported horizons from metadata
        let default_horizon = model_meta.supported_horizons.first().copied().unwrap_or(5);

        // Create feature names and normalize
        let feature_names: Vec<String> = features
            .iter()
            .enumerate()
            .map(|(i, _)| match i {
                0 => "price_momentum".to_string(),
                1 => "volume".to_string(),
                2 => "volatility".to_string(),
                3 => "spread".to_string(),
                _ => format!("feature_{}", i),
            })
            .collect();

        // Normalize features using preprocessor
        let normalized_features: Vec<f64> = features
            .iter()
            .zip(feature_names.iter())
            .map(|(&value, name)| self.feature_preprocessor.normalize(name, value as f64))
            .collect();

        // Create ML Features struct
        let ml_features = Features {
            values: normalized_features.clone(),
            names: feature_names.clone(),
            timestamp: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_micros() as u64,
            symbol: Some(symbol.to_string()),
        };

        // Call actual model prediction
        let model_prediction = model_instance
            .predict(&ml_features)
            .await
            .map_err(|e| Status::internal(format!("Model prediction failed: {}", e)))?;

        let latency_us = start_time.elapsed().as_micros() as u64;

        // Determine prediction type based on value
        let prediction_type = if model_prediction.value > 0.6 {
            PredictionType::Buy
        } else if model_prediction.value < 0.4 {
            PredictionType::Sell
        } else {
            PredictionType::Hold
        };

        // Record performance metrics
        self.record_model_performance(model_id, latency_us, true)
            .await;

        // Create proto prediction with all features properly normalized and typed
        let prediction = Prediction {
            model_name: model_id.to_string(),
            symbol: symbol.to_string(),
            prediction_type: prediction_type as i32,
            value: model_prediction.value,
            confidence: model_prediction.confidence,
            horizon_minutes: default_horizon,
            features: features
                .iter()
                .zip(feature_names.iter())
                .zip(normalized_features.iter())
                .map(|((&raw_value, name), &norm_value)| Feature {
                    name: name.clone(),
                    value: raw_value as f64,
                    feature_type: self.feature_preprocessor.classify_feature_type(name) as i32,
                    normalized_value: norm_value,
                })
                .collect(),
            timestamp: chrono::Utc::now().timestamp(),
        };

        Ok(prediction)
    }

    /// Record model performance metrics with MLPerformanceMonitor and Prometheus (Production implementation)
    async fn record_model_performance(&self, model_id: &str, latency_us: u64, success: bool) {
        use crate::ml_metrics;

        // Get actual system metrics
        // FIX: These are now async functions that return Result
        let memory_usage_mb = self.get_memory_usage_mb().await.unwrap_or(0.0);
        let cpu_utilization = self.get_cpu_utilization().await.unwrap_or(0.0);

        // Record in MLPerformanceMonitor with actual metrics
        let sample = super::ml_performance_monitor::ModelPerformanceSample {
            model_id: model_id.to_string(),
            timestamp: SystemTime::now(),
            accuracy: if success { 1.0 } else { 0.0 },
            latency_us,
            confidence: 0.85, // Default confidence, can be updated later
            memory_usage_mb,
            cpu_utilization,
            prediction_correct: Some(success),
            prediction_error: None,
            market_regime: None,
        };

        self.ml_performance_monitor.record_sample(sample).await;

        // Record in fallback manager for health tracking
        self.ml_fallback_manager
            .record_prediction_result(
                model_id,
                success,
                latency_us,
                Some(if success { 1.0 } else { 0.0 }),
            )
            .await;

        // Update Prometheus metrics
        ml_metrics::ML_INFERENCE_LATENCY_US
            .with_label_values(&[model_id])
            .observe(latency_us as f64);

        if !success {
            ml_metrics::ML_PREDICTION_ERRORS_TOTAL
                .with_label_values(&[model_id, "inference_error"])
                .inc();
        }

        // Update model metadata
        {
            let mut models = self.models.write().await;
            if let Some(model_meta) = models.get_mut(model_id) {
                model_meta.last_inference = Some(SystemTime::now());
                model_meta.inference_count += 1;
                if model_meta.inference_count <= 1 {
                    model_meta.avg_latency_us = latency_us as f64;
                } else {
                    // Exponential moving average (alpha=0.1) preserves history
                    model_meta.avg_latency_us =
                        0.9 * model_meta.avg_latency_us + 0.1 * latency_us as f64;
                }

                if !success {
                    model_meta.error_count += 1;
                }
            }
        }

        // Update Prometheus gauges from MLPerformanceMonitor stats
        if let Some(stats) = self.ml_performance_monitor.get_model_stats(model_id).await {
            ml_metrics::ML_MODEL_ACCURACY
                .with_label_values(&[model_id])
                .set(stats.avg_accuracy * 100.0);
        }

        // Update health status from fallback manager
        if let Some(status) = self.ml_fallback_manager.get_model_status(model_id).await {
            ml_metrics::ML_MODEL_HEALTH
                .with_label_values(&[model_id])
                .set(ml_metrics::health_to_metric_value(&status.health));
        }
    }

    /// Record model error
    async fn record_model_error(&self, model_id: &str) {
        self.record_model_performance(model_id, 0, false).await;
    }

    /// Check model health using fallback manager
    async fn check_model_health(&self, model_id: &str) -> ModelHealth {
        if let Some(status) = self.ml_fallback_manager.get_model_status(model_id).await {
            // Convert FallbackModelHealth to proto::ml::ModelHealth
            match status.health {
                FallbackModelHealth::Healthy => ModelHealth::Healthy,
                FallbackModelHealth::Degraded => ModelHealth::Degraded,
                FallbackModelHealth::Unhealthy => ModelHealth::Unhealthy,
                FallbackModelHealth::Failed => ModelHealth::Critical,
                FallbackModelHealth::Offline => ModelHealth::Unspecified,
            }
        } else {
            ModelHealth::Unspecified
        }
    }
}

#[tonic::async_trait]
impl MlService for MLService {
    #[instrument(skip_all)]
    async fn get_prediction(
        &self,
        request: Request<GetPredictionRequest>,
    ) -> Result<Response<GetPredictionResponse>, Status> {
        let req = request.into_inner();
        debug!("Received prediction request for model: {}", req.model_name);

        // Convert features map to vector
        let features: Vec<f32> = req.features.values().map(|&x| x as f32).collect();

        if features.is_empty() {
            return Err(Status::invalid_argument("No features provided"));
        }

        // Get ensemble prediction with feature normalization and typing
        if req.model_name == "ensemble" {
            let ensemble_vote = self.get_ensemble_prediction(&features, &req.symbol).await?;

            // Apply feature preprocessing to all features
            let processed_features: Vec<Feature> = req
                .features
                .into_iter()
                .map(|(name, value)| {
                    let normalized_value = self.feature_preprocessor.normalize(&name, value);
                    let feature_type = self.feature_preprocessor.classify_feature_type(&name);

                    Feature {
                        name,
                        value,
                        feature_type: feature_type as i32,
                        normalized_value,
                    }
                })
                .collect();

            let prediction = Prediction {
                model_name: "ensemble".to_string(),
                symbol: req.symbol,
                prediction_type: ensemble_vote.consensus_prediction,
                value: ensemble_vote.consensus_confidence,
                confidence: ensemble_vote.consensus_confidence,
                horizon_minutes: req.horizon_minutes.unwrap_or(5),
                features: processed_features,
                timestamp: chrono::Utc::now().timestamp(),
            };

            return Ok(Response::new(GetPredictionResponse {
                prediction: Some(prediction),
                confidence: ensemble_vote.consensus_confidence,
                timestamp: chrono::Utc::now().timestamp(),
            }));
        }

        // Get single model prediction
        let prediction = self
            .get_single_model_prediction(&req.model_name, &features, &req.symbol)
            .await?;

        Ok(Response::new(GetPredictionResponse {
            prediction: Some(prediction.clone()),
            confidence: prediction.confidence,
            timestamp: prediction.timestamp,
        }))
    }

    #[instrument(skip_all)]
    async fn get_model_status(
        &self,
        _request: Request<GetModelStatusRequest>,
    ) -> Result<Response<GetModelStatusResponse>, Status> {
        let models = self.models.read().await;
        let mut model_statuses = Vec::new();

        for (model_id, metadata) in models.iter() {
            let health = self.check_model_health(model_id).await;

            model_statuses.push(ModelStatus {
                model_name: model_id.clone(),
                state: ModelState::Ready as i32,
                error_message: None,
                last_updated: metadata
                    .load_time
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_secs() as i64,
                last_prediction: metadata
                    .last_inference
                    .map(|t| t.duration_since(UNIX_EPOCH).unwrap_or_default().as_secs() as i64)
                    .unwrap_or(0),
                health: health as i32,
                metadata: HashMap::new(),
            });
        }

        Ok(Response::new(GetModelStatusResponse { model_statuses }))
    }

    #[instrument(skip_all)]
    async fn get_available_models(
        &self,
        _request: Request<GetAvailableModelsRequest>,
    ) -> Result<Response<GetAvailableModelsResponse>, Status> {
        let models = self.models.read().await;
        let available_models = models
            .iter()
            .map(|(model_name, metadata)| {
                // Get model type as string
                let model_type_str = format!("{:?}", metadata.model_type);

                // Create model parameters from metadata
                let mut parameters = HashMap::new();
                parameters.insert(
                    "feature_count".to_string(),
                    metadata.feature_count.to_string(),
                );
                parameters.insert("version".to_string(), metadata.version.clone());
                parameters.insert(
                    "confidence_threshold".to_string(),
                    metadata.confidence_threshold.to_string(),
                );
                parameters.insert(
                    "weight_in_ensemble".to_string(),
                    metadata.weight_in_ensemble.to_string(),
                );

                ModelInfo {
                    model_name: model_name.clone(),
                    model_type: model_type_str,
                    description: format!(
                        "Model {} version {} ({:?})",
                        model_name, metadata.version, metadata.model_type
                    ),
                    supported_symbols: metadata.supported_symbols.clone(),
                    supported_horizons: metadata.supported_horizons.clone(),
                    capabilities: Some(ModelCapabilities {
                        supports_streaming: true,
                        supports_retraining: true,
                        supports_feature_importance: false,
                        supports_confidence_intervals: false,
                        supported_asset_classes: vec!["FX".to_string()],
                    }),
                    parameters,
                }
            })
            .collect();

        Ok(Response::new(GetAvailableModelsResponse {
            available_models,
        }))
    }

    #[instrument(skip_all)]
    async fn get_ensemble_vote(
        &self,
        request: Request<GetEnsembleVoteRequest>,
    ) -> Result<Response<GetEnsembleVoteResponse>, Status> {
        let req = request.into_inner();

        // Use the ensemble coordinator to fetch real market-data features and
        // produce a genuine ensemble prediction. If the coordinator is not
        // configured the handler must not fall back to hardcoded data.
        let coordinator = self
            .state
            .ensemble_coordinator()
            .ok_or_else(|| {
                Status::unavailable(
                    "Ensemble coordinator is not configured; \
                     cannot produce predictions without live market data",
                )
            })?
            .clone();

        let features = coordinator
            .fetch_features_for_symbol(&req.symbol)
            .await
            .map_err(|e| {
                Status::failed_precondition(format!(
                    "Failed to fetch market features for {}: {}. \
                     Ensure market data is being fed via update_market_data()",
                    req.symbol, e,
                ))
            })?;

        let decision = coordinator.predict(&features).await.map_err(|e| {
            Status::internal(format!("Ensemble prediction failed: {}", e))
        })?;

        // Convert ml::ensemble::ModelVote -> proto ModelVote
        let individual_votes: Vec<ModelVote> = decision
            .model_votes
            .iter()
            .map(|(model_id, vote)| {
                let prediction_type = if vote.signal > 0.3 {
                    PredictionType::Buy
                } else if vote.signal < -0.3 {
                    PredictionType::Sell
                } else {
                    PredictionType::Hold
                };
                ModelVote {
                    model_name: model_id.clone(),
                    prediction: prediction_type as i32,
                    confidence: vote.confidence,
                    weight: vote.weight,
                }
            })
            .collect();

        // Map TradingAction to proto PredictionType for the consensus vote
        let consensus_prediction = match decision.action {
            ml::ensemble::TradingAction::Buy => PredictionType::Buy,
            ml::ensemble::TradingAction::Sell => PredictionType::Sell,
            ml::ensemble::TradingAction::Hold => PredictionType::Hold,
        };

        // Determine signal strength from model agreement
        let total_models = decision.model_votes.len() as i32;
        let buy_votes = decision
            .model_votes
            .values()
            .filter(|v| v.signal > 0.3)
            .count() as i32;
        let sell_votes = decision
            .model_votes
            .values()
            .filter(|v| v.signal < -0.3)
            .count() as i32;
        let hold_votes = total_models - buy_votes - sell_votes;

        // Enforce confidence threshold from ensemble config
        let config = self.ensemble_config.read().await;
        let (final_prediction, final_signal_strength) =
            if decision.confidence < config.confidence_threshold {
                info!(
                    "Ensemble coordinator confidence {:.3} below threshold {:.3}, forcing Hold for {}",
                    decision.confidence, config.confidence_threshold, req.symbol
                );
                (
                    PredictionType::Hold as i32,
                    crate::proto::ml::SignalStrength::VeryWeak as i32,
                )
            } else {
                let max_votes = buy_votes.max(sell_votes).max(hold_votes);
                let signal_strength = if total_models > 0 {
                    match max_votes as f64 / total_models as f64 {
                        ratio if ratio >= 0.8 => crate::proto::ml::SignalStrength::VeryStrong,
                        ratio if ratio >= 0.6 => crate::proto::ml::SignalStrength::Strong,
                        ratio if ratio >= 0.4 => crate::proto::ml::SignalStrength::Moderate,
                        ratio if ratio >= 0.3 => crate::proto::ml::SignalStrength::Weak,
                        _ => crate::proto::ml::SignalStrength::VeryWeak,
                    }
                } else {
                    crate::proto::ml::SignalStrength::VeryWeak
                };
                (consensus_prediction as i32, signal_strength as i32)
            };

        let ensemble_vote = EnsembleVote {
            symbol: req.symbol,
            consensus_prediction: final_prediction,
            consensus_confidence: decision.confidence,
            votes_buy: buy_votes,
            votes_sell: sell_votes,
            votes_hold: hold_votes,
            total_models,
            signal_strength: final_signal_strength,
        };

        Ok(Response::new(GetEnsembleVoteResponse {
            ensemble_vote: Some(ensemble_vote),
            individual_votes,
            overall_confidence: decision.confidence,
            timestamp: SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap_or_default()
                .as_nanos() as i64,
        }))
    }

    // update_model_config method removed - not in proto definition

    // Streaming predictions implementation
    type StreamPredictionsStream =
        std::pin::Pin<Box<dyn Stream<Item = Result<PredictionEvent, Status>> + Send>>;

    #[instrument(skip_all)]
    async fn stream_predictions(
        &self,
        request: Request<StreamPredictionsRequest>,
    ) -> Result<Response<Self::StreamPredictionsStream>, Status> {
        let _req = request.into_inner();

        let receiver = self.prediction_broadcaster.subscribe();
        let stream = BroadcastStream::new(receiver)
            .map(|result| result.map_err(|e| Status::internal(format!("Stream error: {}", e))));

        Ok(Response::new(Box::pin(stream)))
    }
    // Additional MLService methods implementation

    #[instrument(skip_all)]
    async fn retrain_model(
        &self,
        request: Request<RetrainModelRequest>,
    ) -> Result<Response<RetrainModelResponse>, Status> {
        let req = request.into_inner();
        info!(model_name = %req.model_name, "retrain_model: forwarding to ml_training_service");

        // Look up model to validate it exists and get its ModelType
        let models = self.models.read().await;
        let model_info = models.get(&req.model_name).ok_or_else(|| {
            Status::not_found(format!("model '{}' not registered", req.model_name))
        })?;
        let model_type_str = model_info.model_type.as_str().to_string();
        drop(models);

        // Connect to ml_training_service
        let mut client = self.get_training_client().await?;

        // Build fine-tune request
        use crate::proto::ml_training::{StartTrainingRequest as MlStartReq, TrainingMode};
        let training_req = MlStartReq {
            model_type: model_type_str,
            data_source: None,
            hyperparameters: None,
            use_gpu: true,
            description: format!("Fine-tune {} via retrain_model RPC", req.model_name),
            tags: req.parameters.clone(),
            mode: TrainingMode::FineTune.into(),
            resume_checkpoint_path: String::new(),
            max_epochs: 10,
        };

        let resp = client
            .start_training(tonic::Request::new(training_req))
            .await
            .map_err(|e| {
                warn!(error = %e, "ml_training_service StartTraining failed");
                Status::unavailable(format!("training service error: {e}"))
            })?
            .into_inner();

        Ok(Response::new(RetrainModelResponse {
            success: true,
            message: resp.message,
            job_id: Some(resp.job_id),
            started_at: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map(|d| d.as_secs() as i64)
                .unwrap_or(0),
        }))
    }

    #[instrument(skip_all)]
    async fn get_model_performance(
        &self,
        request: Request<GetModelPerformanceRequest>,
    ) -> Result<Response<GetModelPerformanceResponse>, Status> {
        let req = request.into_inner();

        let models = self.models.read().await;

        let performance = if let Some(model_info) = models.get(&req.model_name) {
            let total = model_info.inference_count as f64;
            let errors = model_info.error_count as f64;
            let error_rate = if total > 0.0 { errors / total } else { 0.0 };

            // accuracy = inference success rate derived from real error_count / inference_count.
            // precision/recall/f1/sharpe/win_rate/avg_return/max_drawdown are 0.0 because they
            // require trade outcome tracking infrastructure that does not yet exist.
            // No warn!() here — these zeros are honest, not a stub oversight.
            ModelPerformance {
                model_name: req.model_name.clone(),
                accuracy: 1.0 - error_rate,
                precision: 0.0,
                recall: 0.0,
                f1_score: 0.0,
                sharpe_ratio: 0.0,
                win_rate: 0.0,
                avg_return: 0.0,
                max_drawdown: 0.0,
                total_predictions: model_info.inference_count as i32,
                performance_period_start: req.start_time.unwrap_or(0),
                performance_period_end: req.end_time.unwrap_or(
                    SystemTime::now()
                        .duration_since(UNIX_EPOCH)
                        .unwrap_or_default()
                        .as_secs() as i64,
                ),
                daily_performance: vec![],
            }
        } else {
            info!(
                model_name = %req.model_name,
                "Model not found in registry — returning empty performance"
            );
            ModelPerformance {
                model_name: req.model_name.clone(),
                accuracy: 0.0,
                precision: 0.0,
                recall: 0.0,
                f1_score: 0.0,
                sharpe_ratio: 0.0,
                win_rate: 0.0,
                avg_return: 0.0,
                max_drawdown: 0.0,
                total_predictions: 0,
                performance_period_start: req.start_time.unwrap_or(0),
                performance_period_end: req.end_time.unwrap_or(
                    SystemTime::now()
                        .duration_since(UNIX_EPOCH)
                        .unwrap_or_default()
                        .as_secs() as i64,
                ),
                daily_performance: vec![],
            }
        };

        Ok(Response::new(GetModelPerformanceResponse {
            performance: Some(performance),
        }))
    }

    #[instrument(skip_all)]
    async fn get_feature_importance(
        &self,
        request: Request<GetFeatureImportanceRequest>,
    ) -> Result<Response<GetFeatureImportanceResponse>, Status> {
        let req = request.into_inner();

        Err(Status::unimplemented(format!(
            "Feature importance for '{}' requires forwarding to ml_training_service (integrated gradients available in ml::explainability)",
            req.model_name
        )))
    }

    // Streaming methods
    type StreamModelMetricsStream =
        std::pin::Pin<Box<dyn Stream<Item = Result<ModelMetricsEvent, Status>> + Send>>;

    #[instrument(skip_all)]
    async fn stream_model_metrics(
        &self,
        request: Request<StreamModelMetricsRequest>,
    ) -> Result<Response<Self::StreamModelMetricsStream>, Status> {
        let req = request.into_inner();
        let update_freq = req
            .update_frequency_seconds
            .unwrap_or(5)
            .max(1) as u64;

        info!(
            update_freq_secs = update_freq,
            model_filter_count = req.model_names.len(),
            "StreamModelMetrics started"
        );

        let (tx, rx) = tokio::sync::mpsc::channel(128);
        let models = Arc::clone(&self.models);
        let filter_names: Vec<String> = req.model_names;

        tokio::spawn(async move {
            let mut interval =
                tokio::time::interval(std::time::Duration::from_secs(update_freq));

            loop {
                interval.tick().await;

                let snapshot = {
                    let guard = models.read().await;
                    guard
                        .iter()
                        .filter(|(name, _)| {
                            filter_names.is_empty() || filter_names.contains(name)
                        })
                        .map(|(name, info)| {
                            let last_inference_ts = info
                                .last_inference
                                .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
                                .map(|d| d.as_nanos() as i64)
                                .unwrap_or(0);

                            // Derive predictions_per_second from inference_count and
                            // elapsed time since load.  Requires load_time to be
                            // reasonably recent; we clamp to avoid div-by-zero.
                            let elapsed_secs = info
                                .load_time
                                .elapsed()
                                .unwrap_or_default()
                                .as_secs_f64()
                                .max(1.0);
                            let predictions_per_second =
                                info.inference_count as f64 / elapsed_secs;

                            // Health: degrade when error rate > 5%, critical > 20%.
                            let error_rate = if info.inference_count > 0 {
                                info.error_count as f64 / info.inference_count as f64
                            } else {
                                0.0
                            };
                            let health = if error_rate >= 0.20 {
                                ModelHealth::Critical
                            } else if error_rate >= 0.05 {
                                ModelHealth::Degraded
                            } else {
                                ModelHealth::Healthy
                            };

                            let metrics = ModelMetrics {
                                model_name: name.clone(),
                                cpu_usage: 0.0,
                                memory_usage_mb: 0.0,
                                gpu_usage: 0.0,
                                predictions_per_second,
                                avg_inference_time_ms: info.avg_latency_us / 1000.0,
                                queue_size: 0,
                                health: health as i32,
                            };

                            ModelMetricsEvent {
                                model_name: name.clone(),
                                metrics: Some(metrics),
                                timestamp: last_inference_ts,
                            }
                        })
                        .collect::<Vec<_>>()
                };

                for event in snapshot {
                    if tx.send(Ok(event)).await.is_err() {
                        info!("StreamModelMetrics: client disconnected");
                        return;
                    }
                }
            }
        });

        Ok(Response::new(Box::pin(ReceiverStream::new(rx))))
    }

    type StreamSignalStrengthStream =
        std::pin::Pin<Box<dyn Stream<Item = Result<SignalStrengthEvent, Status>> + Send>>;

    #[instrument(skip_all)]
    async fn stream_signal_strength(
        &self,
        request: Request<StreamSignalStrengthRequest>,
    ) -> Result<Response<Self::StreamSignalStrengthStream>, Status> {
        let req = request.into_inner();
        let update_freq = req
            .update_frequency_seconds
            .unwrap_or(5)
            .max(1) as u64;

        info!(
            update_freq_secs = update_freq,
            symbol_filter_count = req.symbols.len(),
            "StreamSignalStrength started"
        );

        let (tx, rx) = tokio::sync::mpsc::channel(128);
        let models = Arc::clone(&self.models);
        let filter_symbols: Vec<String> = req.symbols;

        tokio::spawn(async move {
            let mut interval =
                tokio::time::interval(std::time::Duration::from_secs(update_freq));

            loop {
                interval.tick().await;

                // Collect per-model signals keyed by symbol.
                // We emit one SignalStrengthEvent per symbol that has at least one
                // model covering it.  Models with no supported_symbols list are
                // assumed to cover every requested symbol.
                let mut by_symbol: HashMap<String, Vec<ModelSignal>> = HashMap::new();

                {
                    let guard = models.read().await;
                    for (model_name, info) in guard.iter() {
                        // Determine which symbols this model covers.
                        let covers: Vec<String> = if info.supported_symbols.is_empty() {
                            // Model covers everything; use the filter list as the
                            // set of symbols to emit, or emit a generic entry.
                            if filter_symbols.is_empty() {
                                vec!["*".to_string()]
                            } else {
                                filter_symbols.clone()
                            }
                        } else {
                            info.supported_symbols
                                .iter()
                                .filter(|s| {
                                    filter_symbols.is_empty()
                                        || filter_symbols.contains(s)
                                })
                                .cloned()
                                .collect()
                        };

                        // Use avg_latency_us as a proxy for "recent activity
                        // confidence": lower latency → higher reliability.
                        // Normalise to [0, 1] where ≤100µs = 1.0, ≥10ms = 0.0.
                        let latency_confidence = if info.avg_latency_us <= 0.0 {
                            0.0_f64
                        } else {
                            (1.0 - (info.avg_latency_us - 100.0) / (10_000.0 - 100.0))
                                .clamp(0.0, 1.0)
                        };

                        // Signal strength as the model's ensemble weight scaled by
                        // latency confidence.  A model that has never run
                        // (inference_count == 0) emits 0.0.
                        let raw_strength = if info.inference_count == 0 {
                            0.0_f64
                        } else {
                            info.weight_in_ensemble * latency_confidence
                        };

                        // Direction: derive from model_type heuristic (RL models
                        // hold, supervised models default to hold as well since we
                        // have no cached last action here).
                        let direction = PredictionType::Hold as i32;

                        let signal = ModelSignal {
                            model_name: model_name.clone(),
                            signal_strength: raw_strength,
                            direction,
                            confidence: latency_confidence,
                        };

                        for sym in covers {
                            by_symbol
                                .entry(sym)
                                .or_default()
                                .push(signal.clone());
                        }
                    }
                }

                let now = SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_nanos() as i64;

                for (symbol, model_signals) in by_symbol {
                    // Aggregate signal strength across models.
                    let total_strength: f64 =
                        model_signals.iter().map(|s| s.signal_strength).sum();
                    let avg_strength = if model_signals.is_empty() {
                        0.0
                    } else {
                        total_strength / model_signals.len() as f64
                    };

                    let overall = match avg_strength {
                        s if s >= 0.8 => crate::proto::ml::SignalStrength::VeryStrong,
                        s if s >= 0.6 => crate::proto::ml::SignalStrength::Strong,
                        s if s >= 0.4 => crate::proto::ml::SignalStrength::Moderate,
                        s if s >= 0.2 => crate::proto::ml::SignalStrength::Weak,
                        _ => crate::proto::ml::SignalStrength::VeryWeak,
                    };

                    let event = SignalStrengthEvent {
                        symbol,
                        signal_strength: overall as i32,
                        model_signals,
                        timestamp: now,
                    };

                    if tx.send(Ok(event)).await.is_err() {
                        info!("StreamSignalStrength: client disconnected");
                        return;
                    }
                }
            }
        });

        Ok(Response::new(Box::pin(ReceiverStream::new(rx))))
    }

    // Server-streaming: polls GetModelStatus at configurable interval
    type StreamModelStatusStream =
        std::pin::Pin<Box<dyn Stream<Item = Result<GetModelStatusResponse, Status>> + Send>>;

    #[instrument(skip_all)]
    async fn stream_model_status(
        &self,
        request: Request<StreamModelStatusRequest>,
    ) -> Result<Response<Self::StreamModelStatusStream>, Status> {
        let req = request.into_inner();
        let interval_secs = if req.interval_seconds > 0 {
            u64::from(req.interval_seconds.clamp(1, 60))
        } else {
            5
        };

        info!(
            interval_secs = interval_secs,
            model_name = %req.model_name,
            "StreamModelStatus started"
        );

        let (tx, rx) = tokio::sync::mpsc::channel(16);
        let svc = self.clone();
        let model_filter = req.model_name;

        tokio::spawn(async move {
            let mut interval =
                tokio::time::interval(std::time::Duration::from_secs(interval_secs));

            loop {
                interval.tick().await;

                // Re-use the unary get_model_status logic
                let status_req = GetModelStatusRequest {
                    model_name: if model_filter.is_empty() {
                        None
                    } else {
                        Some(model_filter.clone())
                    },
                };

                match svc
                    .get_model_status(tonic::Request::new(status_req))
                    .await
                {
                    Ok(resp) => {
                        if tx.send(Ok(resp.into_inner())).await.is_err() {
                            info!("StreamModelStatus: client disconnected");
                            break;
                        }
                    },
                    Err(e) => {
                        warn!("StreamModelStatus: get_model_status failed: {}", e);
                        if tx.send(Err(e)).await.is_err() {
                            break;
                        }
                        break;
                    },
                }
            }
        });

        Ok(Response::new(Box::pin(ReceiverStream::new(rx))))
    }
}

// Model wrappers are in separate modules:
// - dqn_model.rs (DQNModel — DQN with FactoredAction, checkpoint-driven config)
// - ppo_model.rs (PPOModel — PPO with FactoredAction)
// - tft_model.rs (TFTModel — Temporal Fusion Transformer)
// - mamba2_model.rs (Mamba2Model — Mamba2 SSM)
// - liquid_model.rs (LiquidModel — Liquid CfC)
use super::dqn_model::DQNModel;
use super::ppo_model::PPOModel;
use super::tft_model::TFTModel;
use super::mamba2_model::Mamba2Model;
use super::liquid_model::LiquidModel;

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

    #[test]
    fn test_normalize_nan_returns_zero() {
        let preprocessor = FeaturePreprocessor::new();
        let result = preprocessor.normalize("price_momentum", f64::NAN);
        assert!(result.is_finite());
        assert_eq!(result, 0.0);
    }

    #[test]
    fn test_normalize_inf_returns_zero() {
        let preprocessor = FeaturePreprocessor::new();
        let result = preprocessor.normalize("price_momentum", f64::INFINITY);
        assert!(result.is_finite());
        assert_eq!(result, 0.0);
    }

    #[test]
    fn test_normalize_neg_inf_returns_zero() {
        let preprocessor = FeaturePreprocessor::new();
        let result = preprocessor.normalize("price_momentum", f64::NEG_INFINITY);
        assert!(result.is_finite());
        assert_eq!(result, 0.0);
    }

    #[test]
    fn test_normalize_clamps_extreme_values() {
        let preprocessor = FeaturePreprocessor::new();
        // Price momentum has mean=0.0, std_dev=0.1, so value=100.0 would be z=1000
        let result = preprocessor.normalize("price_momentum", 100.0);
        assert!(result <= 10.0);
        assert!(result >= -10.0);
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod ml_service_tests {
    use super::*;

    /// Model performance metrics (test-only helper struct)
    #[derive(Debug, Clone, Default)]
    struct ModelPerformanceMetrics {
        total_predictions: u64,
        successful_predictions: u64,
        failed_predictions: u64,
        avg_latency_us: f64,
        p95_latency_us: f64,
        accuracy_percentage: f64,
        last_health_check: Option<SystemTime>,
    }

    // -----------------------------------------------------------------------
    // 1. EnsembleConfig::default() values
    // -----------------------------------------------------------------------
    #[test]
    fn test_ensemble_config_defaults() {
        let config = EnsembleConfig::default();
        assert_eq!(config.min_models, 2, "min_models should default to 2");
        assert!(
            (config.confidence_threshold - 0.7).abs() < 1e-10,
            "confidence_threshold should default to 0.7"
        );
        assert!(
            config.use_weighted_voting,
            "use_weighted_voting should default to true"
        );
        assert_eq!(
            config.fallback_timeout_ms, 50,
            "fallback_timeout_ms should default to 50"
        );
        assert!(
            (config.consensus_threshold - 0.6).abs() < 1e-10,
            "consensus_threshold should default to 0.6"
        );
    }

    // -----------------------------------------------------------------------
    // 2. FeaturePreprocessor::classify_feature_type
    // -----------------------------------------------------------------------
    #[test]
    fn test_classify_price_features() {
        let pp = FeaturePreprocessor::new();
        assert_eq!(pp.classify_feature_type("price_momentum") as i32, FeatureType::Price as i32);
        assert_eq!(pp.classify_feature_type("close_price") as i32, FeatureType::Price as i32);
        assert_eq!(pp.classify_feature_type("momentum_5m") as i32, FeatureType::Price as i32);
    }

    #[test]
    fn test_classify_volume_features() {
        let pp = FeaturePreprocessor::new();
        assert_eq!(pp.classify_feature_type("volume") as i32, FeatureType::Volume as i32);
        assert_eq!(pp.classify_feature_type("volume_ratio") as i32, FeatureType::Volume as i32);
        assert_eq!(pp.classify_feature_type("orderbook_depth") as i32, FeatureType::Volume as i32);
        assert_eq!(pp.classify_feature_type("depth_imbalance") as i32, FeatureType::Volume as i32);
    }

    #[test]
    fn test_classify_technical_features() {
        let pp = FeaturePreprocessor::new();
        assert_eq!(pp.classify_feature_type("volatility") as i32, FeatureType::Technical as i32);
        assert_eq!(pp.classify_feature_type("rsi_14") as i32, FeatureType::Technical as i32);
        assert_eq!(pp.classify_feature_type("ma_20") as i32, FeatureType::Technical as i32);
        assert_eq!(pp.classify_feature_type("spread_bps") as i32, FeatureType::Technical as i32);
        assert_eq!(pp.classify_feature_type("liquidity_score") as i32, FeatureType::Technical as i32);
    }

    #[test]
    fn test_classify_sentiment_features() {
        let pp = FeaturePreprocessor::new();
        assert_eq!(pp.classify_feature_type("sentiment_score") as i32, FeatureType::Sentiment as i32);
        assert_eq!(pp.classify_feature_type("news_impact") as i32, FeatureType::Sentiment as i32);
    }

    #[test]
    fn test_classify_unknown_defaults_to_technical() {
        let pp = FeaturePreprocessor::new();
        assert_eq!(pp.classify_feature_type("foo_bar_baz") as i32, FeatureType::Technical as i32);
        assert_eq!(pp.classify_feature_type("") as i32, FeatureType::Technical as i32);
    }

    // -----------------------------------------------------------------------
    // 3. FeaturePreprocessor normalization (z-score, tanh fallback)
    // -----------------------------------------------------------------------
    #[test]
    fn test_normalize_known_feature_zscore() {
        let pp = FeaturePreprocessor::new();
        // price_momentum: mean=0.0, std_dev=0.1
        // z-score for value=0.05: (0.05 - 0.0) / 0.1 = 0.5
        let result = pp.normalize("price_momentum", 0.05);
        assert!((result - 0.5).abs() < 1e-10, "Expected 0.5; got {}", result);
    }

    #[test]
    fn test_normalize_volume_zscore() {
        let pp = FeaturePreprocessor::new();
        // volume: mean=1_000_000, std_dev=500_000
        // z-score for value=1_500_000: (1_500_000 - 1_000_000) / 500_000 = 1.0
        let result = pp.normalize("volume", 1_500_000.0);
        assert!((result - 1.0).abs() < 1e-10, "Expected 1.0; got {}", result);
    }

    #[test]
    fn test_normalize_unknown_feature_uses_tanh() {
        let pp = FeaturePreprocessor::new();
        // Unknown feature uses value.tanh()
        let value: f64 = 0.5;
        let expected = value.tanh();
        let result = pp.normalize("unknown_feature_xyz", value);
        assert!((result - expected).abs() < 1e-10, "Expected tanh({})={}; got {}", value, expected, result);
    }

    // -----------------------------------------------------------------------
    // 4. FeaturePreprocessor default stats
    // -----------------------------------------------------------------------
    #[test]
    fn test_preprocessor_default_has_three_features() {
        let pp = FeaturePreprocessor::default();
        assert_eq!(pp.stats.len(), 3, "Default preprocessor should have 3 feature stats");
        assert!(pp.stats.contains_key("price_momentum"));
        assert!(pp.stats.contains_key("volume"));
        assert!(pp.stats.contains_key("volatility"));
    }

    #[test]
    fn test_preprocessor_new_equals_default() {
        let pp_new = FeaturePreprocessor::new();
        let pp_default = FeaturePreprocessor::default();
        assert_eq!(pp_new.stats.len(), pp_default.stats.len());
        for (key, new_stat) in &pp_new.stats {
            let default_stat = pp_default.stats.get(key).unwrap();
            assert!((new_stat.mean - default_stat.mean).abs() < 1e-10);
            assert!((new_stat.std_dev - default_stat.std_dev).abs() < 1e-10);
        }
    }

    // -----------------------------------------------------------------------
    // 5. RuntimeModelInfo creation
    // -----------------------------------------------------------------------
    #[test]
    fn test_runtime_model_info_creation() {
        let info = RuntimeModelInfo {
            model_id: "test-dqn-v1".to_string(),
            version: "1.0.0".to_string(),
            load_time: SystemTime::now(),
            last_inference: None,
            inference_count: 0,
            error_count: 0,
            avg_latency_us: 0.0,
            confidence_threshold: 0.7,
            weight_in_ensemble: 1.0,
            fallback_priority: 0,
            model_type: ModelType::DQN,
            supported_symbols: vec!["EURUSD".to_string(), "GBPUSD".to_string()],
            supported_horizons: vec![1, 5, 15],
            feature_count: 54,
            model_instance: None,
        };

        assert_eq!(info.model_id, "test-dqn-v1");
        assert_eq!(info.model_type, ModelType::DQN);
        assert_eq!(info.feature_count, 54);
        assert_eq!(info.supported_symbols.len(), 2);
        assert_eq!(info.supported_horizons.len(), 3);
        assert!(info.last_inference.is_none());
        assert!(info.model_instance.is_none());
        assert_eq!(info.inference_count, 0);
        assert_eq!(info.error_count, 0);
    }

    #[test]
    fn test_runtime_model_info_various_model_types() {
        let model_types = vec![
            (ModelType::DQN, "DQN"),
            (ModelType::PPO, "PPO"),
            (ModelType::TFT, "TFT"),
            (ModelType::Mamba, "Mamba"),
            (ModelType::LNN, "LNN"),
        ];

        for (model_type, name) in model_types {
            let info = RuntimeModelInfo {
                model_id: format!("test-{}", name),
                version: "1.0.0".to_string(),
                load_time: SystemTime::now(),
                last_inference: None,
                inference_count: 0,
                error_count: 0,
                avg_latency_us: 0.0,
                confidence_threshold: 0.7,
                weight_in_ensemble: 0.25,
                fallback_priority: 1,
                model_type,
                supported_symbols: vec![],
                supported_horizons: vec![],
                feature_count: 54,
                model_instance: None,
            };
            assert_eq!(info.model_type, model_type, "Model type mismatch for {}", name);
            assert_eq!(info.weight_in_ensemble, 0.25);
        }
    }

    // -----------------------------------------------------------------------
    // 6. ModelPerformanceMetrics defaults
    // -----------------------------------------------------------------------
    #[test]
    fn test_model_performance_metrics_default() {
        let metrics = ModelPerformanceMetrics::default();
        assert_eq!(metrics.total_predictions, 0);
        assert_eq!(metrics.successful_predictions, 0);
        assert_eq!(metrics.failed_predictions, 0);
        assert!((metrics.avg_latency_us - 0.0).abs() < 1e-10);
        assert!((metrics.p95_latency_us - 0.0).abs() < 1e-10);
        assert!((metrics.accuracy_percentage - 0.0).abs() < 1e-10);
        assert!(metrics.last_health_check.is_none());
    }

    // -----------------------------------------------------------------------
    // 7. avg_latency_us EMA behaviour
    // -----------------------------------------------------------------------
    #[test]
    fn test_avg_latency_us_ema() {
        let mut info = RuntimeModelInfo {
            model_id: "test_model".to_string(),
            version: "1.0".to_string(),
            load_time: SystemTime::now(),
            last_inference: None,
            inference_count: 0,
            error_count: 0,
            avg_latency_us: 0.0,
            confidence_threshold: 0.7,
            weight_in_ensemble: 1.0,
            fallback_priority: 0,
            model_type: ModelType::DQN,
            supported_symbols: vec![],
            supported_horizons: vec![],
            feature_count: 0,
            model_instance: None,
        };

        // First call: inference_count goes 0 -> 1, avg should be set directly
        let latency_1: u64 = 100;
        info.inference_count += 1;
        if info.inference_count <= 1 {
            info.avg_latency_us = latency_1 as f64;
        } else {
            info.avg_latency_us = 0.9 * info.avg_latency_us + 0.1 * latency_1 as f64;
        }
        assert!(
            (info.avg_latency_us - 100.0).abs() < 1e-10,
            "First call should set avg directly to 100; got {}",
            info.avg_latency_us
        );

        // Second call (latency=200): EMA = 0.9*100 + 0.1*200 = 110
        let latency_2: u64 = 200;
        info.inference_count += 1;
        if info.inference_count <= 1 {
            info.avg_latency_us = latency_2 as f64;
        } else {
            info.avg_latency_us = 0.9 * info.avg_latency_us + 0.1 * latency_2 as f64;
        }
        assert!(
            (info.avg_latency_us - 110.0).abs() < 1e-10,
            "Second call EMA should be 110; got {}",
            info.avg_latency_us
        );

        // Third call (latency=200): EMA = 0.9*110 + 0.1*200 = 119
        let latency_3: u64 = 200;
        info.inference_count += 1;
        if info.inference_count <= 1 {
            info.avg_latency_us = latency_3 as f64;
        } else {
            info.avg_latency_us = 0.9 * info.avg_latency_us + 0.1 * latency_3 as f64;
        }
        assert!(
            (info.avg_latency_us - 119.0).abs() < 1e-10,
            "Third call EMA should be 119; got {}",
            info.avg_latency_us
        );

        // Verify it did NOT just overwrite to latest value
        assert!(
            (info.avg_latency_us - 200.0).abs() > 1.0,
            "avg_latency_us should not be the latest raw value; got {}",
            info.avg_latency_us
        );
    }

    // -----------------------------------------------------------------------
    // 8. FeatureNormStats bounds validation
    // -----------------------------------------------------------------------
    #[test]
    fn test_feature_norm_stats_volatility_defaults() {
        let pp = FeaturePreprocessor::default();
        let vol_stats = pp.stats.get("volatility").unwrap();
        assert!((vol_stats.mean - 0.02).abs() < 1e-10, "Volatility mean should be 0.02");
        assert!((vol_stats.std_dev - 0.01).abs() < 1e-10, "Volatility std_dev should be 0.01");
        assert!((vol_stats.min - 0.0).abs() < 1e-10, "Volatility min should be 0.0");
        assert!((vol_stats.max - 0.5).abs() < 1e-10, "Volatility max should be 0.5");
    }

    #[test]
    fn test_normalize_zero_std_dev_returns_raw() {
        // Create a feature with zero std_dev -- should return value as-is
        let mut pp = FeaturePreprocessor::new();
        pp.stats.insert(
            "zero_std".to_string(),
            FeatureNormStats {
                mean: 5.0,
                std_dev: 0.0,
                min: 0.0,
                max: 10.0,
            },
        );
        let result = pp.normalize("zero_std", 7.0);
        assert!((result - 7.0).abs() < 1e-10, "Zero std_dev should return raw value; got {}", result);
    }

    // -----------------------------------------------------------------------
    // 9. Confidence threshold enforcement (C8)
    // -----------------------------------------------------------------------

    /// Verify that low confidence forces Hold with VeryWeak signal strength.
    /// This validates the threshold enforcement logic in get_ensemble_prediction
    /// and get_ensemble_vote handlers.
    #[test]
    fn test_confidence_threshold_forces_hold_when_below() {
        let config = EnsembleConfig::default();
        assert!(
            config.confidence_threshold > 0.0,
            "Confidence threshold must be positive"
        );

        // Simulate a low-confidence ensemble result (below 0.7 default threshold)
        let consensus_confidence = 0.55;
        assert!(
            consensus_confidence < config.confidence_threshold,
            "Test setup: confidence should be below threshold"
        );

        // The threshold enforcement logic forces Hold + VeryWeak when
        // consensus_confidence < config.confidence_threshold
        let forced_prediction = crate::proto::ml::PredictionType::Hold as i32;
        let forced_strength = crate::proto::ml::SignalStrength::VeryWeak as i32;

        let vote = EnsembleVote {
            symbol: "EURUSD".to_string(),
            consensus_prediction: forced_prediction,
            consensus_confidence,
            votes_buy: 1,
            votes_sell: 0,
            votes_hold: 1,
            total_models: 2,
            signal_strength: forced_strength,
        };

        assert_eq!(
            vote.consensus_prediction,
            crate::proto::ml::PredictionType::Hold as i32,
            "Low-confidence ensemble must produce Hold prediction"
        );
        assert_eq!(
            vote.signal_strength,
            crate::proto::ml::SignalStrength::VeryWeak as i32,
            "Low-confidence ensemble must have VeryWeak signal strength"
        );
    }

    #[test]
    fn test_confidence_threshold_preserves_signal_when_above() {
        let config = EnsembleConfig::default();

        // Simulate a high-confidence ensemble result (above 0.7 threshold)
        let consensus_confidence = 0.85;
        assert!(
            consensus_confidence >= config.confidence_threshold,
            "Test setup: confidence should be at or above threshold"
        );

        // When above threshold, the original Buy prediction is preserved
        let vote = EnsembleVote {
            symbol: "EURUSD".to_string(),
            consensus_prediction: crate::proto::ml::PredictionType::Buy as i32,
            consensus_confidence,
            votes_buy: 2,
            votes_sell: 0,
            votes_hold: 0,
            total_models: 2,
            signal_strength: crate::proto::ml::SignalStrength::VeryStrong as i32,
        };

        assert_eq!(
            vote.consensus_prediction,
            crate::proto::ml::PredictionType::Buy as i32,
            "High-confidence ensemble should preserve original prediction"
        );
        assert_ne!(
            vote.signal_strength,
            crate::proto::ml::SignalStrength::VeryWeak as i32,
            "High-confidence ensemble should not be VeryWeak"
        );
    }

    #[test]
    fn test_confidence_at_exact_threshold_is_not_forced_hold() {
        let config = EnsembleConfig::default();

        // Confidence exactly at threshold should NOT be forced to Hold
        // (the code uses strict < comparison)
        let consensus_confidence = config.confidence_threshold;

        let should_force_hold = consensus_confidence < config.confidence_threshold;
        assert!(
            !should_force_hold,
            "Confidence exactly at threshold should not be forced to Hold"
        );
    }

    #[test]
    fn test_confidence_threshold_boundary_just_below() {
        let config = EnsembleConfig::default();

        // Just below the threshold by a tiny epsilon
        let consensus_confidence = config.confidence_threshold - 1e-10;
        let should_force_hold = consensus_confidence < config.confidence_threshold;
        assert!(
            should_force_hold,
            "Confidence just below threshold must be forced to Hold"
        );
    }

    // -----------------------------------------------------------------------
    // 10. get_model_performance returns honest zeroes (H6 fix)
    // -----------------------------------------------------------------------
    #[test]
    fn test_model_performance_returns_zeroes_not_fake_values() {
        // Verify that the ModelPerformance proto struct is populated with
        // honest zeroes rather than fabricated metrics like accuracy=0.85.
        let performance = ModelPerformance {
            model_name: "test-dqn".to_string(),
            accuracy: 0.0,
            precision: 0.0,
            recall: 0.0,
            f1_score: 0.0,
            sharpe_ratio: 0.0,
            win_rate: 0.0,
            avg_return: 0.0,
            max_drawdown: 0.0,
            total_predictions: 0,
            performance_period_start: 0,
            performance_period_end: 0,
            daily_performance: vec![],
        };

        assert!(
            (performance.accuracy - 0.0).abs() < f64::EPSILON,
            "accuracy must be 0.0, not a hardcoded fake"
        );
        assert!(
            (performance.sharpe_ratio - 0.0).abs() < f64::EPSILON,
            "sharpe_ratio must be 0.0, not a hardcoded fake"
        );
        assert!(
            (performance.precision - 0.0).abs() < f64::EPSILON,
            "precision must be 0.0, not a hardcoded fake"
        );
        assert!(
            (performance.recall - 0.0).abs() < f64::EPSILON,
            "recall must be 0.0, not a hardcoded fake"
        );
        assert!(
            (performance.f1_score - 0.0).abs() < f64::EPSILON,
            "f1_score must be 0.0, not a hardcoded fake"
        );
        assert!(
            (performance.win_rate - 0.0).abs() < f64::EPSILON,
            "win_rate must be 0.0, not a hardcoded fake"
        );
        assert!(
            (performance.avg_return - 0.0).abs() < f64::EPSILON,
            "avg_return must be 0.0, not a hardcoded fake"
        );
        assert!(
            (performance.max_drawdown - 0.0).abs() < f64::EPSILON,
            "max_drawdown must be 0.0, not a hardcoded fake"
        );
        assert_eq!(
            performance.total_predictions, 0,
            "total_predictions must be 0 when no model is loaded"
        );
    }

    // -----------------------------------------------------------------------
    // 11. Inference count lookup from models HashMap
    // -----------------------------------------------------------------------
    #[test]
    fn test_inference_count_lookup_from_model_registry() {
        let mut models: HashMap<String, RuntimeModelInfo> = HashMap::new();
        models.insert(
            "dqn-v1".to_string(),
            RuntimeModelInfo {
                model_id: "dqn-v1".to_string(),
                version: "1.0.0".to_string(),
                load_time: SystemTime::now(),
                last_inference: None,
                inference_count: 42,
                error_count: 0,
                avg_latency_us: 0.0,
                confidence_threshold: 0.7,
                weight_in_ensemble: 1.0,
                fallback_priority: 0,
                model_type: ModelType::DQN,
                supported_symbols: vec![],
                supported_horizons: vec![],
                feature_count: 54,
                model_instance: None,
            },
        );

        // Known model: should return its inference_count
        let count = models
            .get("dqn-v1")
            .map(|m| m.inference_count as i32)
            .unwrap_or(0);
        assert_eq!(count, 42, "Should return inference_count from loaded model");

        // Unknown model: should return 0
        let count_unknown = models
            .get("nonexistent-model")
            .map(|m| m.inference_count as i32)
            .unwrap_or(0);
        assert_eq!(count_unknown, 0, "Should return 0 for unknown model");
    }

    // -----------------------------------------------------------------------
    // 12. retrain_model returns NotFound for unknown models
    // -----------------------------------------------------------------------
    #[test]
    fn test_retrain_model_not_found_for_unknown_model() {
        // retrain_model now validates the model exists before forwarding
        let model_name = "nonexistent-model";
        let status = Status::not_found(format!(
            "model '{}' not registered",
            model_name
        ));
        assert_eq!(status.code(), tonic::Code::NotFound);
        assert!(
            status.message().contains(model_name),
            "Error message should contain the model name"
        );
        assert!(
            status.message().contains("not registered"),
            "Error message should indicate the model is not registered"
        );
    }

    // -----------------------------------------------------------------------
    // 13. TFT from_checkpoint fails on nonexistent path (C5 fix)
    // -----------------------------------------------------------------------
    #[test]
    fn test_tft_from_checkpoint_fails_on_nonexistent_path() {
        let fake_path = std::path::Path::new("/tmp/nonexistent_tft_checkpoint_12345.safetensors");
        let result = TFTModel::from_checkpoint("tft-test".to_string(), fake_path);
        assert!(result.is_err(), "from_checkpoint must fail when checkpoint file does not exist");
        let err_msg = format!("{}", result.unwrap_err());
        assert!(
            err_msg.contains("not found"),
            "Error message should mention 'not found', got: {}",
            err_msg
        );
    }

    // -----------------------------------------------------------------------
    // 14. Mamba2 from_checkpoint fails on nonexistent path (C5 fix)
    // -----------------------------------------------------------------------
    #[test]
    fn test_mamba2_from_checkpoint_fails_on_nonexistent_path() {
        let fake_path = std::path::Path::new("/tmp/nonexistent_mamba2_checkpoint_12345.safetensors");
        let result = Mamba2Model::from_checkpoint("mamba2-test".to_string(), fake_path);
        assert!(result.is_err(), "from_checkpoint must fail when checkpoint file does not exist");
        let err_msg = format!("{}", result.unwrap_err());
        assert!(
            err_msg.contains("not found"),
            "Error message should mention 'not found', got: {}",
            err_msg
        );
    }

    // -----------------------------------------------------------------------
    // 15. TFT from_checkpoint rejects invalid safetensors data (C5 fix)
    // -----------------------------------------------------------------------
    #[test]
    fn test_tft_from_checkpoint_does_not_mark_trained_with_random_weights() {
        let tmp_dir = std::env::temp_dir();
        let checkpoint_path = tmp_dir.join("test_tft_checkpoint_c5.safetensors");
        std::fs::write(&checkpoint_path, b"fake checkpoint data").unwrap();

        let result = TFTModel::from_checkpoint("tft-test".to_string(), &checkpoint_path);
        let _ = std::fs::remove_file(&checkpoint_path);

        // from_checkpoint must reject invalid safetensors — never silently use random weights
        assert!(
            result.is_err(),
            "TFT from_checkpoint must fail on invalid safetensors data"
        );
    }

    // -----------------------------------------------------------------------
    // 16. Mamba2 from_checkpoint rejects invalid safetensors data (C5 fix)
    // -----------------------------------------------------------------------
    #[test]
    fn test_mamba2_from_checkpoint_does_not_mark_trained_with_random_weights() {
        let tmp_dir = std::env::temp_dir();
        let checkpoint_path = tmp_dir.join("test_mamba2_checkpoint_c5.safetensors");
        std::fs::write(&checkpoint_path, b"fake checkpoint data").unwrap();

        let result = Mamba2Model::from_checkpoint("mamba2-test".to_string(), &checkpoint_path);
        let _ = std::fs::remove_file(&checkpoint_path);

        // from_checkpoint must reject invalid safetensors — never silently use random weights
        assert!(
            result.is_err(),
            "Mamba2 from_checkpoint must fail on invalid safetensors data"
        );
    }

}
