foxhunt/crates/ml-ensemble/src/inference_adapter.rs

//! Model inference adapter trait for ensemble prediction
//!
//! Defines the contract that each model adapter must implement
//! to participate in ensemble prediction.

use crate::MLResult;

/// Canonical feature vector for ensemble inference.
/// 51-dim layout matching DQN state_dim:
///   0-9: price features (OHLCV, VWAP, spread, tick)
///   10-19: technical indicators (RSI, MACD, Bollinger, ATR)
///   20-29: order book (depth levels, imbalance)
///   30-39: microstructure (trade flow, toxicity)
///   40-50: position/risk (PnL, exposure, drawdown)
#[derive(Debug, Clone)]
pub struct FeatureVector {
    pub values: Vec<f64>,
    pub timestamp: i64,
}

/// Prediction output from a single model adapter, normalized for ensemble aggregation.
#[derive(Debug, Clone)]
pub struct EnsemblePrediction {
    pub model_name: String,
    /// Directional signal: -1.0 (bearish) to 1.0 (bullish)
    pub direction: f64,
    /// Model confidence: 0.0 to 1.0
    pub confidence: f64,
    /// Model-specific metadata
    pub metadata: PredictionMeta,
}

/// Raw prediction with optional GPU logits for batched aggregation.
///
/// When `logits` is `Some`, the ensemble can aggregate on GPU before extraction.
/// When `logits` is `None`, falls back to `direction_scalar` (already CPU).
#[derive(Debug, Clone)]
pub struct RawPrediction {
    /// Pre-computed direction (used when logits are not available)
    pub direction_scalar: f64,
    /// Model confidence: 0.0 to 1.0
    pub confidence: f64,
    /// Raw GPU logits before sigmoid/softmax (if available).
    /// Vec<f32> for CPU-side aggregation after GPU extraction.
    pub logits: Option<Vec<f32>>,
}

/// Optional model-specific metadata attached to predictions.
#[derive(Debug, Clone, Default)]
pub struct PredictionMeta {
    /// Inference latency in microseconds
    pub latency_us: u64,
    /// Quantile forecasts (TFT only)
    pub quantiles: Option<Vec<f64>>,
    /// Attention weights (TFT only)
    pub attention_weights: Option<Vec<f64>>,
    /// Raw Q-values (DQN only)
    pub q_values: Option<Vec<f64>>,
}

/// Adapter trait for running inference on a loaded model.
///
/// Each model type (DQN, PPO, TFT, Mamba2) implements this trait
/// to handle its own feature-to-tensor mapping and output normalization.
pub trait ModelInferenceAdapter: Send + Sync {
    /// Human-readable model name for logging
    fn model_name(&self) -> &str;

    /// Run inference on a canonical feature vector.
    /// Returns a normalized ensemble prediction.
    fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction>;

    /// Whether this adapter has a loaded model ready for inference
    fn is_ready(&self) -> bool;

    /// Return raw logits + confidence for GPU-side aggregation.
    ///
    /// Default implementation falls back to [`predict()`](Self::predict) and wraps
    /// the result with `logits: None`.  Adapters that hold a GPU model can
    /// override this to return the pre-sigmoid/softmax logits as a `Vec<f32>`,
    /// enabling the ensemble to aggregate entirely on device.
    fn predict_raw(&self, features: &FeatureVector) -> MLResult<RawPrediction> {
        let pred = self.predict(features)?;
        Ok(RawPrediction {
            direction_scalar: pred.direction,
            confidence: pred.confidence,
            logits: None,
        })
    }

    /// Batched inference: process multiple feature vectors in a single GPU
    /// upload -> forward -> download cycle instead of per-sample roundtrips.
    ///
    /// Default implementation falls back to calling [`predict()`](Self::predict)
    /// in a loop. Adapters with GPU models should override this to create
    /// a single `[N, feature_dim]` tensor, run one forward pass, and extract
    /// all results at once.
    fn predict_batch(&self, batch: &[FeatureVector]) -> MLResult<Vec<EnsemblePrediction>> {
        batch.iter().map(|fv| self.predict(fv)).collect()
    }
}

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

    #[test]
    fn test_feature_vector_creation() {
        let fv = FeatureVector {
            values: vec![0.1; 51],
            timestamp: 1700000000_000_000,
        };
        assert_eq!(fv.values.len(), 51);
        assert_eq!(fv.timestamp, 1700000000_000_000);
    }

    #[test]
    fn test_ensemble_prediction_direction_bounds() {
        let pred = EnsemblePrediction {
            model_name: "test".to_owned(),
            direction: 0.75,
            confidence: 0.9,
            metadata: PredictionMeta::default(),
        };
        assert!(pred.direction >= -1.0 && pred.direction <= 1.0);
        assert!(pred.confidence >= 0.0 && pred.confidence <= 1.0);
    }

    #[test]
    fn test_prediction_meta_default() {
        let meta = PredictionMeta::default();
        assert_eq!(meta.latency_us, 0);
        assert!(meta.quantiles.is_none());
        assert!(meta.attention_weights.is_none());
        assert!(meta.q_values.is_none());
    }
}